Aventra / semantic_memory.py
Tpayne101's picture
Create semantic_memory.py
eb68b9c verified
raw
history blame contribute delete
813 Bytes
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]]