| 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() | |