File size: 813 Bytes
eb68b9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
from sentence_transformers import SentenceTransformer
import numpy as np
class SemanticMemory:
def __init__(self):
self.memory_db = []
self.model = SentenceTransformer('all-MiniLM-L6-v2')
def store_memory(self, conversation, metadata=None):
embedding = self.model.encode(conversation)
self.memory_db.append({'embedding': embedding, 'conversation': conversation, 'metadata': metadata})
def retrieve_relevant(self, query, k=3):
query_emb = self.model.encode(query)
def cosine_similarity(a, b):
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
scores = [(cosine_similarity(query_emb, m['embedding']), m['conversation']) for m in self.memory_db]
scores.sort(reverse=True)
return [s[1] for s in scores[:k]] |