mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
import csv
|
|
|
|
from langchain.chains.question_answering import load_qa_chain
|
|
from langchain_chroma import Chroma
|
|
from langchain_community.embeddings import CohereEmbeddings
|
|
from langchain_core.prompts import PromptTemplate
|
|
|
|
from .chat import chat
|
|
|
|
csv_file = open("data/books_with_blurbs.csv", "r")
|
|
csv_reader = csv.reader(csv_file)
|
|
csv_data = list(csv_reader)
|
|
parsed_data = [
|
|
{
|
|
"id": x[0],
|
|
"title": x[1],
|
|
"author": x[2],
|
|
"year": x[3],
|
|
"publisher": x[4],
|
|
"blurb": x[5],
|
|
}
|
|
for x in csv_data
|
|
]
|
|
parsed_data[1]
|
|
|
|
embeddings = CohereEmbeddings(model="embed-english-v3.0")
|
|
|
|
docsearch = Chroma.from_texts(
|
|
[x["title"] for x in parsed_data], embeddings, metadatas=parsed_data
|
|
).as_retriever()
|
|
|
|
|
|
prompt_template = """
|
|
{context}
|
|
|
|
Use the book reccommendations to suggest books for the user to read.
|
|
Only use the titles of the books, do not make up titles. Format the response as
|
|
a bulleted list prefixed by a relevant message.
|
|
|
|
User: {message}"""
|
|
|
|
PROMPT = PromptTemplate(
|
|
template=prompt_template, input_variables=["context", "message"]
|
|
)
|
|
|
|
book_rec_chain = {
|
|
"input_documents": lambda x: docsearch.invoke(x["message"]),
|
|
"message": lambda x: x["message"],
|
|
} | load_qa_chain(chat, chain_type="stuff", prompt=PROMPT)
|