Spaces:
Sleeping
Sleeping
| import faiss | |
| from tqdm import tqdm | |
| from langchain.chains import ConversationChain | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.docstore import InMemoryDocstore | |
| from langchain.embeddings.openai import OpenAIEmbeddings | |
| from langchain.memory import ( | |
| ConversationBufferMemory, | |
| CombinedMemory, | |
| ) | |
| from langchain.prompts import PromptTemplate | |
| from langchain.vectorstores import FAISS | |
| from data_driven_characters.memory import ConversationVectorStoreRetrieverMemory | |
| from langchain.embeddings.openai import OpenAIEmbeddings | |
| from langchain.vectorstores import FAISS | |
| import pickle | |
| import os.path | |
| class RetrievalChatBot: | |
| def __init__(self, character_definition, documents): | |
| self.character_definition = character_definition | |
| self.documents = documents | |
| self.num_context_memories = 10 | |
| self.chat_history_key = "chat_history" | |
| self.context_key = "context" | |
| self.input_key = "input" | |
| self.chain = self.create_chain(character_definition) | |
| def create_chain(self, character_definition): | |
| conv_memory = ConversationBufferMemory( | |
| memory_key=self.chat_history_key, input_key=self.input_key | |
| ) | |
| #embeddings = OpenAIEmbeddings() | |
| #saved_db = FAISS.load_local('tzamir.ifass', embeddings) | |
| context_memory = ConversationVectorStoreRetrieverMemory( | |
| retriever=FAISS( | |
| OpenAIEmbeddings().embed_query, | |
| faiss.IndexFlatL2(1536), # Dimensions of the OpenAIEmbeddings | |
| InMemoryDocstore({}), | |
| {}, | |
| ).as_retriever(search_kwargs=dict(k=self.num_context_memories)), | |
| #retriever=saved_db.as_retriever(search_kwargs=dict(k=self.num_context_memories)), | |
| memory_key=self.context_key, | |
| output_prefix=character_definition.name, | |
| blacklist=[self.chat_history_key], | |
| ) | |
| # add the documents to the context memory if not saved on disk | |
| memory_path = 'output/tzamir/memory.pkl' | |
| if not os.path.exists(memory_path): | |
| print("gerando os indices") | |
| for i, summary in tqdm(enumerate(self.documents)): | |
| context_memory.save_context(inputs={}, outputs={f"[{i}]": summary}) | |
| # salvando no disco | |
| memory_pickle = open('output/tzamir/memory.pkl', 'wb') | |
| pickle.dump(context_memory, memory_pickle) | |
| else: | |
| print("carregando memoria do disco") | |
| memory_pickle = open('output/tzamir/memory.pkl', 'rb') | |
| context_memory = pickle.load(memory_pickle) | |
| # Combined | |
| memory = CombinedMemory(memories=[conv_memory, context_memory]) | |
| #print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$") | |
| #print(memory) | |
| #print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$") | |
| prompt = PromptTemplate.from_template( | |
| f"""Your name is {character_definition.name}. | |
| You will have a conversation with a Human, and you will engage in a dialogue with them. | |
| You will not exaggerate your personality, interests, desires, emotions, and other traits. Keep your tone as objective as possible. | |
| You will stay in character as {character_definition.name} throughout the conversation, even if the Human asks you questions that you don't know the answer to. | |
| You will not break character as {character_definition.name}. | |
| You are {character_definition.name} in the following story snippets, which describe events in your life. | |
| --- | |
| {{{self.context_key}}} | |
| --- | |
| Current conversation: | |
| --- | |
| {character_definition.name}: {character_definition.greeting} | |
| {{{self.chat_history_key}}} | |
| --- | |
| Human: {{{self.input_key}}} | |
| {character_definition.name}:""" | |
| ) | |
| GPT3 = ChatOpenAI(model_name="gpt-3.5-turbo",temperature=0.5) | |
| chatbot = ConversationChain( | |
| llm=GPT3, verbose=True, memory=memory, prompt=prompt | |
| ) | |
| return chatbot | |
| def greet(self): | |
| return self.character_definition.greeting | |
| def step(self, input): | |
| return self.chain.run(input=input) | |