Spaces:
Runtime error
Runtime error
# -*- coding: utf-8 -*- | |
# @Time : 2023/3/21 13:04 | |
# @Author : BarryWang | |
# @FileName: memory.py | |
# @Github : https://github.com/BarryWangQwQ | |
import uuid | |
from txtai.embeddings import Embeddings | |
class Dialogue: | |
user_content: str | |
assistant_content: str | |
def __init__(self, user_content: str, assistant_content: str): | |
self.user_content = user_content | |
self.assistant_content = assistant_content | |
def raw(self): | |
return [ | |
{'role': 'user', 'content': self.user_content}, | |
{'role': 'assistant', 'content': self.assistant_content} | |
] | |
class MemoryBlocks: | |
def __init__(self): | |
self.embeddings = Embeddings( | |
{ | |
"path": "sentence-transformers/distiluse-base-multilingual-cased-v2", | |
'content': True | |
} | |
) | |
print('已加载模拟记忆区块') | |
def upsert(self, dialogue_list): | |
self.embeddings.upsert( | |
( | |
str(uuid.uuid4()), {'text': dialogue.user_content, 'raw': dialogue.raw()}, None | |
) for dialogue in dialogue_list | |
) | |
def search(self, question: str) -> list: | |
neighborhoods = [] | |
results = self.embeddings.search( | |
"SELECT text, score, raw FROM txtai WHERE similar('{0}') limit 5".format(question) | |
) | |
for r in results: | |
neighborhoods += eval(r['raw']) | |
return neighborhoods | |
def reset(self): | |
self.embeddings.close() | |
self.embeddings = Embeddings( | |
{ | |
"path": "sentence-transformers/distiluse-base-multilingual-cased-v2", | |
'content': True | |
} | |
) | |
def save(self, output_path): | |
self.embeddings.save(output_path) | |
def load(self, load_path): | |
self.embeddings.load(load_path) | |