added core functional

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -1,4 +1,8 @@
+import torch
+import pickle
+import nmslib
 import gradio as gr
+from sentence_transformers import SentenceTransformer
 
 
 K = 5
@@ -20,16 +24,32 @@ def create_demo(callback):
 
 
 class Callback:
-    def __init__(self):
-        pass
+    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):
-        res = ["These are completely different companies."] * K
-        return res
+        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():
-    callback = Callback()
+    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()
 

requirements.txt

@@ -0,0 +1,4 @@
+torch
+gradio
+nmslib
+sentence_transformers