Upload m3.py

m3.py

@@ -1,9 +1,10 @@
-import os, xml.etree.ElementTree as ET, torch, torch.nn as nn, torch.nn.functional as F, faiss, numpy as np
+import os, xml.etree.ElementTree as ET, torch, torch.nn as nn, torch.nn.functional as F, numpy as np
 from typing import List, Dict, Any, Optional
 from collections import defaultdict
 from accelerate import Accelerator
 from transformers import AutoTokenizer, AutoModel
 from termcolor import colored
+from sklearn.metrics.pairwise import cosine_similarity
 
 class DM(nn.Module):
     def __init__(self, s: Dict[str, List[Dict[str, Any]]]):
@@ -108,7 +109,7 @@ def cmf(folder_path: str) -> DM:
 def ceas(folder_path: str, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
     t = AutoTokenizer.from_pretrained(model_name)
     m = AutoModel.from_pretrained(model_name)
-    vs = faiss.IndexFlatL2(384)
+    embeddings = []
     ds = []
     for r, d, f in os.walk(folder_path):
         for file in f:
@@ -123,20 +124,22 @@ def ceas(folder_path: str, model_name: str = "sentence-transformers/all-MiniLM-L
                             i = t(text, return_tensors="pt", truncation=True, padding=True)
                             with torch.no_grad():
                                 emb = m(**i).last_hidden_state.mean(dim=1).numpy()
-                            vs.add(emb)
+                            embeddings.append(emb)
                             ds.append(text)
                 except Exception as e:
                     print(colored(f"Error processing {fp}: {str(e)}", 'red'))
-    return vs, ds
+    embeddings = np.vstack(embeddings)
+    return embeddings, ds
 
-def qvs(query: str, vs, ds, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
+def qvs(query: str, embeddings, ds, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
     t = AutoTokenizer.from_pretrained(model_name)
     m = AutoModel.from_pretrained(model_name)
     i = t(query, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         qe = m(**i).last_hidden_state.mean(dim=1).numpy()
-    D, I = vs.search(qe, k=5)
-    return [ds[i] for i in I[0]]
+    similarities = cosine_similarity(qe, embeddings)
+    top_k_indices = similarities[0].argsort()[-5:][::-1]
+    return [ds[i] for i in top_k_indices]
 
 def main():
     fp = 'data'
@@ -148,7 +151,7 @@ def main():
     si = torch.randn(1, ife)
     o = m(si)
     print(colored(f"Sample output shape: {o.shape}", 'green'))
-    vs, ds = ceas(fp)
+    embeddings, ds = ceas(fp)
     a = Accelerator()
     o = torch.optim.Adam(m.parameters(), lr=0.001)
     c = nn.CrossEntropyLoss()
@@ -169,7 +172,7 @@ def main():
         al = tl / len(td)
         print(colored(f"Epoch {e+1}/{ne}, Average Loss: {al:.4f}", 'blue'))
     uq = "example query text"
-    r = qvs(uq, vs, ds)
+    r = qvs(uq, embeddings, ds)
     print(colored(f"Query results: {r}", 'magenta'))
 
 if __name__ == "__main__":