Update memory.py

memory.py

@@ -1,22 +1,21 @@
-import os
 from transformers import AutoTokenizer, AutoModel
+import torch
+import os
 
 class MemoryVectorStore:
-    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
-        self.cache_dir = "/tmp/hf_cache"
-        os.makedirs(self.cache_dir, exist_ok=False)
-
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_name, cache_dir=self.cache_dir
-        )
-        self.model = AutoModel.from_pretrained(
-            model_name, cache_dir=self.cache_dir
-        )
-
-        self.memory = []
+    def __init__(self):
+        # Download model to project dir
+        model_name = "sentence-transformers/all-MiniLM-L6-v2"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="./hf-cache")
+        self.model = AutoModel.from_pretrained(model_name, cache_dir="./hf-cache")
 
-    def add(self, text):
-        self.memory.append(text)
+    def encode(self, texts):
+        # Ensure list of texts
+        if isinstance(texts, str):
+            texts = [texts]
 
-    def retrieve(self, query):
-        return self.memory
+        with torch.no_grad():
+            inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+            outputs = self.model(**inputs)
+            embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
+            return embeddings.numpy()