import gradio as gr
import numpy as np
from usearch.index import Index
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
from sentencex import segment

model = SentenceTransformer("Corran/SciGenAllMiniLM")

train = load_dataset("Corran/SciGenColbertTriplets")['train']

rf = sorted(list(set(train['query'])))
rf_emb = model.encode(rf)
index = Index(ndim=rf_emb[0].size)
index.add(range(len(classes)), classes_emb)

def get_matches(input):
  global index, model, rf
  emb = model.encode(input,batch_size=128)
  matches = index.search(emb,4)
  if type(input)==list and len(input)>1:
    matches = [m[0] for m in matches]
  else:
    matches = [m for m in matches]
  return [(rf[m.key],m.distance) for m in matches]



def return_rf_scores(paragraph):
    
    sentences = list(segment("en", paragraph))
    matches = get_matches(sentences)
    
    output = {}
    for s,m in zip(sentences,matches):
        output[s] = m
        
    return output

demo = gr.Interface(fn=return_rf_scores, inputs="text", outputs="json")
demo.launch()