Update Similarity.py

Similarity.py

@@ -1,15 +1,19 @@
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
 
-class Similarity:
+class SimpleEmbedder:
     def __init__(self):
-        self.model = None
-
-    def load_model(self):
-        if self.model is None:
-            print("Loading SentenceTransformer model...")
-            self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-            print("Model loaded.")
+        print("Loading tokenizer and model...")
+        self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+        self.model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+        print("Loaded.")
 
     def embed_text(self, text):
-        self.load_model()
-        return self.model.encode(text, convert_to_tensor=True)
+        inputs = self.tokenizer(text, return_tensors='pt')
+        outputs = self.model(**inputs)
+        # Mean pooling
+        embeddings = outputs.last_hidden_state.mean(dim=1)
+        return embeddings
+
+embedder = SimpleEmbedder()
+print(embedder.embed_text("Hello world"))