Upload indexer.py

indexer.py

@@ -0,0 +1,64 @@
+import pickle
+import faiss
+import numpy as np
+# from grammar import remove_verbs, clean_text
+from utils import *
+from sentence_transformers import SentenceTransformer
+
+
+class FAISS:
+    def __init__(self, dimensions: int):
+        self.dimensions = dimensions
+        self.index = faiss.IndexFlatL2(dimensions)
+        self.vectors = {}
+        self.counter = 0
+        self.model_name = 'paraphrase-multilingual-MiniLM-L12-v2'
+        self.sentence_encoder = SentenceTransformer(self.model_name)
+
+    def init_vectors(self, path):
+        with open(path, 'rb') as pkl_file:
+            self.vectors = pickle.load(pkl_file)
+
+    def init_index(self, path):
+        self.index = faiss.read_index(path)
+
+    def add(self, text, idx, pop, emb=None):
+        if emb is None:
+            text_vec = self.sentence_encoder.encode([text])
+        else:
+            text_vec = emb
+        self.index.add(text_vec)
+        self.vectors[self.counter] = (idx, text, pop, text_vec)
+        self.counter += 1
+
+    def search(self, v: list, k: int = 10):
+        result = []
+        distance, item_index = self.index.search(v, k)
+        for dist, i in zip(distance[0], item_index[0]):
+            if i == -1:
+                break
+            else:
+                result.append((self.vectors[i][0], self.vectors[i][1], self.vectors[i][2], dist))
+
+        return result
+
+    def suggest_tags(self, query, top_n=10, k=30) -> list:
+
+        emb = self.sentence_encoder.encode([query.lower()])
+        r = self.search(emb, k)
+
+        result = []
+        for i in r:
+            if check(query, i[1]):
+                result.append(i)
+        # надо добавить вес относительно длины
+        result = sorted(result, key=lambda x: x[0] * 0.3 - x[-1], reverse=True)
+        total_result = []
+        for i in range(len(result)):
+            flag = True
+            for j in result[i + 1:]:
+                flag &= sweet_check(result[i][1], j[1])
+            if flag:
+                total_result.append(result[i][1])
+
+        return total_result[:top_n]