mirror of https://github.com/hwchase17/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
304 lines
9.5 KiB
Plaintext
304 lines
9.5 KiB
Plaintext
2 years ago
|
{
|
||
|
"cells": [
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "94e33ebe",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
2 years ago
|
"# How to create a custom Memory class\n",
|
||
2 years ago
|
"Although there are a few predefined types of memory in LangChain, it is highly possible you will want to add your own type of memory that is optimal for your application. This notebook covers how to do that."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "bdfd0305",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"For this notebook, we will add a custom memory type to `ConversationChain`. In order to add a custom memory class, we need to import the base memory class and subclass it."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
|
"execution_count": 1,
|
||
|
"id": "6d787ef2",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
|
"from langchain import OpenAI, ConversationChain\n",
|
||
2 years ago
|
"from langchain.schema import BaseMemory\n",
|
||
2 years ago
|
"from pydantic import BaseModel\n",
|
||
|
"from typing import List, Dict, Any"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "9489e5e1",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"In this example, we will write a custom memory class that uses spacy to extract entities and save information about them in a simple hash table. Then, during the conversation, we will look at the input text, extract any entities, and put any information about them into the context.\n",
|
||
|
"\n",
|
||
|
"* Please note that this implementation is pretty simple and brittle and probably not useful in a production setting. Its purpose is to showcase that you can add custom memory implementations.\n",
|
||
|
"\n",
|
||
|
"For this, we will need spacy."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": null,
|
||
|
"id": "48a5dd13",
|
||
2 years ago
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
|
"# !pip install spacy\n",
|
||
|
"# !python -m spacy download en_core_web_lg"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 5,
|
||
2 years ago
|
"id": "ff065f58",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
|
"import spacy\n",
|
||
1 year ago
|
"\n",
|
||
|
"nlp = spacy.load(\"en_core_web_lg\")"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 9,
|
||
2 years ago
|
"id": "1d45d429",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
2 years ago
|
"class SpacyEntityMemory(BaseMemory, BaseModel):\n",
|
||
2 years ago
|
" \"\"\"Memory class for storing information about entities.\"\"\"\n",
|
||
|
"\n",
|
||
|
" # Define dictionary to store information about entities.\n",
|
||
|
" entities: dict = {}\n",
|
||
|
" # Define key to pass information about entities into prompt.\n",
|
||
2 years ago
|
" memory_key: str = \"entities\"\n",
|
||
1 year ago
|
"\n",
|
||
2 years ago
|
" def clear(self):\n",
|
||
|
" self.entities = {}\n",
|
||
2 years ago
|
"\n",
|
||
|
" @property\n",
|
||
2 years ago
|
" def memory_variables(self) -> List[str]:\n",
|
||
|
" \"\"\"Define the variables we are providing to the prompt.\"\"\"\n",
|
||
|
" return [self.memory_key]\n",
|
||
2 years ago
|
"\n",
|
||
2 years ago
|
" def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]:\n",
|
||
|
" \"\"\"Load the memory variables, in this case the entity key.\"\"\"\n",
|
||
2 years ago
|
" # Get the input text and run through spacy\n",
|
||
|
" doc = nlp(inputs[list(inputs.keys())[0]])\n",
|
||
|
" # Extract known information about entities, if they exist.\n",
|
||
1 year ago
|
" entities = [\n",
|
||
|
" self.entities[str(ent)] for ent in doc.ents if str(ent) in self.entities\n",
|
||
|
" ]\n",
|
||
2 years ago
|
" # Return combined information about entities to put into context.\n",
|
||
2 years ago
|
" return {self.memory_key: \"\\n\".join(entities)}\n",
|
||
2 years ago
|
"\n",
|
||
|
" def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:\n",
|
||
|
" \"\"\"Save context from this conversation to buffer.\"\"\"\n",
|
||
|
" # Get the input text and run through spacy\n",
|
||
|
" text = inputs[list(inputs.keys())[0]]\n",
|
||
|
" doc = nlp(text)\n",
|
||
|
" # For each entity that was mentioned, save this information to the dictionary.\n",
|
||
|
" for ent in doc.ents:\n",
|
||
|
" ent_str = str(ent)\n",
|
||
|
" if ent_str in self.entities:\n",
|
||
|
" self.entities[ent_str] += f\"\\n{text}\"\n",
|
||
|
" else:\n",
|
||
|
" self.entities[ent_str] = text"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "429ba264",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"We now define a prompt that takes in information about entities as well as user input"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 10,
|
||
2 years ago
|
"id": "c05159b6",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
|
"from langchain.prompts.prompt import PromptTemplate\n",
|
||
|
"\n",
|
||
|
"template = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
|
||
|
"\n",
|
||
|
"Relevant entity information:\n",
|
||
|
"{entities}\n",
|
||
|
"\n",
|
||
|
"Conversation:\n",
|
||
|
"Human: {input}\n",
|
||
|
"AI:\"\"\"\n",
|
||
1 year ago
|
"prompt = PromptTemplate(input_variables=[\"entities\", \"input\"], template=template)"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "db611041",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"And now we put it all together!"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 11,
|
||
2 years ago
|
"id": "f08dc8ed",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": [
|
||
|
"llm = OpenAI(temperature=0)\n",
|
||
1 year ago
|
"conversation = ConversationChain(\n",
|
||
|
" llm=llm, prompt=prompt, verbose=True, memory=SpacyEntityMemory()\n",
|
||
|
")"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "92a5f685",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"In the first example, with no prior knowledge about Harrison, the \"Relevant entity information\" section is empty."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 12,
|
||
2 years ago
|
"id": "5b96e836",
|
||
|
"metadata": {},
|
||
|
"outputs": [
|
||
|
{
|
||
|
"name": "stdout",
|
||
|
"output_type": "stream",
|
||
|
"text": [
|
||
|
"\n",
|
||
|
"\n",
|
||
2 years ago
|
"\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
|
||
2 years ago
|
"Prompt after formatting:\n",
|
||
|
"\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
|
||
|
"\n",
|
||
|
"Relevant entity information:\n",
|
||
|
"\n",
|
||
|
"\n",
|
||
|
"Conversation:\n",
|
||
|
"Human: Harrison likes machine learning\n",
|
||
|
"AI:\u001b[0m\n",
|
||
|
"\n",
|
||
2 years ago
|
"\u001b[1m> Finished ConversationChain chain.\u001b[0m\n"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"data": {
|
||
|
"text/plain": [
|
||
2 years ago
|
"\" That's great to hear! Machine learning is a fascinating field of study. It involves using algorithms to analyze data and make predictions. Have you ever studied machine learning, Harrison?\""
|
||
2 years ago
|
]
|
||
|
},
|
||
2 years ago
|
"execution_count": 12,
|
||
2 years ago
|
"metadata": {},
|
||
|
"output_type": "execute_result"
|
||
|
}
|
||
|
],
|
||
|
"source": [
|
||
|
"conversation.predict(input=\"Harrison likes machine learning\")"
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "b1faa743",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"Now in the second example, we can see that it pulls in information about Harrison."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
2 years ago
|
"execution_count": 13,
|
||
2 years ago
|
"id": "4bca7070",
|
||
|
"metadata": {},
|
||
|
"outputs": [
|
||
|
{
|
||
|
"name": "stdout",
|
||
|
"output_type": "stream",
|
||
|
"text": [
|
||
|
"\n",
|
||
|
"\n",
|
||
2 years ago
|
"\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
|
||
2 years ago
|
"Prompt after formatting:\n",
|
||
|
"\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. You are provided with information about entities the Human mentions, if relevant.\n",
|
||
|
"\n",
|
||
|
"Relevant entity information:\n",
|
||
|
"Harrison likes machine learning\n",
|
||
|
"\n",
|
||
|
"Conversation:\n",
|
||
|
"Human: What do you think Harrison's favorite subject in college was?\n",
|
||
|
"AI:\u001b[0m\n",
|
||
|
"\n",
|
||
2 years ago
|
"\u001b[1m> Finished ConversationChain chain.\u001b[0m\n"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"data": {
|
||
|
"text/plain": [
|
||
2 years ago
|
"' From what I know about Harrison, I believe his favorite subject in college was machine learning. He has expressed a strong interest in the subject and has mentioned it often.'"
|
||
2 years ago
|
]
|
||
|
},
|
||
2 years ago
|
"execution_count": 13,
|
||
2 years ago
|
"metadata": {},
|
||
|
"output_type": "execute_result"
|
||
|
}
|
||
|
],
|
||
|
"source": [
|
||
1 year ago
|
"conversation.predict(\n",
|
||
|
" input=\"What do you think Harrison's favorite subject in college was?\"\n",
|
||
|
")"
|
||
2 years ago
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "markdown",
|
||
|
"id": "58b856e3",
|
||
|
"metadata": {},
|
||
|
"source": [
|
||
|
"Again, please note that this implementation is pretty simple and brittle and probably not useful in a production setting. Its purpose is to showcase that you can add custom memory implementations."
|
||
|
]
|
||
|
},
|
||
|
{
|
||
|
"cell_type": "code",
|
||
|
"execution_count": null,
|
||
|
"id": "a1994600",
|
||
|
"metadata": {},
|
||
|
"outputs": [],
|
||
|
"source": []
|
||
|
}
|
||
|
],
|
||
|
"metadata": {
|
||
|
"kernelspec": {
|
||
|
"display_name": "Python 3 (ipykernel)",
|
||
|
"language": "python",
|
||
|
"name": "python3"
|
||
|
},
|
||
|
"language_info": {
|
||
|
"codemirror_mode": {
|
||
|
"name": "ipython",
|
||
|
"version": 3
|
||
|
},
|
||
|
"file_extension": ".py",
|
||
|
"mimetype": "text/x-python",
|
||
|
"name": "python",
|
||
|
"nbconvert_exporter": "python",
|
||
|
"pygments_lexer": "ipython3",
|
||
2 years ago
|
"version": "3.9.1"
|
||
2 years ago
|
}
|
||
|
},
|
||
|
"nbformat": 4,
|
||
|
"nbformat_minor": 5
|
||
|
}
|