Spaces:
Runtime error
Runtime error
File size: 2,612 Bytes
28a47b6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import os
import pickle
from langchain import FAISS
from langchain.retrievers import TimeWeightedVectorStoreRetriever
from langchain.embeddings.base import Embeddings
import faiss
from context import Context
from model import load_embedding_from_config
import math
from langchain.docstore import InMemoryDocstore
# reference:
# https://python.langchain.com/en/latest/use_cases/agent_simulations/characters.html#create-a-generative-character
def relevance_score_fn(score: float) -> float:
"""Return a similarity score on a scale [0, 1]."""
# This will differ depending on a few things:
# - the distance / similarity metric used by the VectorStore
# - the scale of your embeddings (OpenAI's are unit norm. Many others are not!)
# This function converts the euclidean norm of normalized embeddings
# (0 is most similar, sqrt(2) most dissimilar)
# to a similarity function (0 to 1)
return 1.0 - score / math.sqrt(2)
# reference:
# https://python.langchain.com/en/latest/use_cases/agent_simulations/characters.html#create-a-generative-character
def create_new_memory_retriever(ctx: Context):
"""Create a new vector store retriever unique to the agent."""
# Define your embedding model
embeddings_model = load_embedding_from_config(ctx.settings.model.embedding)
# Initialize the vectorstore as empty
embedding_size = 1536
index = faiss.IndexFlatL2(embedding_size)
vectorstore = FAISS(
embeddings_model.embed_query,
index,
InMemoryDocstore({}),
{},
relevance_score_fn=relevance_score_fn,
)
return TimeWeightedVectorStoreRetriever(
vectorstore=vectorstore, other_score_keys=["importance"], k=15
)
class Retriever(TimeWeightedVectorStoreRetriever):
embedding_model: Embeddings
def faiss_path(self, path) -> str:
return path + "/faiss"
def mem_path(self, path) -> str:
return path + "/memory.pickle"
def try_load_memory(self, path: str) -> bool:
if not os.path.isdir(path):
return False
faiss_path = self.faiss_path(path)
faiss: FAISS = self.vectorstore
faiss.load_local(faiss_path, self.embedding_model)
mem_path = self.mem_path(path)
with open(mem_path, "rb") as mem_file:
self.memory_stream = pickle.load(mem_file)
return True
def dump_memory(self, path: str) -> bool:
faiss: FAISS = self.vectorstore
faiss.save_local(self.faiss_path(path))
with open(self.mem_path(path), "wb") as mem_file:
pickle.dump(self.memory_stream, mem_file)
|