Spaces:
Runtime error
Runtime error
File size: 1,641 Bytes
fa438f6 986a798 fa438f6 986a798 fa438f6 986a798 fa438f6 986a798 fa438f6 986a798 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import torch
import pickle
import nmslib
import gradio as gr
from sentence_transformers import SentenceTransformer
K = 5
def create_demo(callback):
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
fn = gr.Textbox(label="Company name", placeholder="Enter company name here...")
with gr.Row():
with gr.Column():
outs = [gr.Text(show_label=False) for _ in range(K)]
outs[0].label = "Similar company names"
outs[0].show_label = True
btn = gr.Button("Find similar companies", variant="primary")
btn.click(callback, inputs=fn, outputs=outs)
return demo
class Callback:
def __init__(self, model, data):
self.index = nmslib.init(method='hnsw', space='cosinesimil')
self.index.addDataPointBatch(data["emb"])
self.index.createIndex({'post': 2}, print_progress=True)
self.model = model
self.data = data
def __call__(self, input_name):
emb = self.model.encode(input_name)
ids, _ = self.index.knnQuery(emb, k=K)
names = [self.data["names"][id] for id in ids]
return names
def load_data(filename):
with open(filename, "rb") as file:
data = pickle.load(file)
return data
def main():
data = load_data("data.pickle")
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("Vsevolod/company-names-similarity-sentence-transformer").to(device)
callback = Callback(model, data)
demo = create_demo(callback)
demo.launch()
if __name__ == "__main__":
main()
|