forked from evannaderi/Epic-Law
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathquery_data.py
33 lines (26 loc) · 1.21 KB
/
query_data.py
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
# PDF Loaders. If unstructured gives you a hard time, try PyPDFLoader
from langchain.document_loaders import UnstructuredPDFLoader, OnlinePDFLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
from langchain.vectorstores import Chroma, Pinecone
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.chains.question_answering import load_qa_chain
from dotenv import load_dotenv
load_dotenv()
loader = UnstructuredPDFLoader('AFA10.pdf')
api_key = os.getenv("OPENAI_API_KEY")
data = loader.load()
print(f'You have {len(data)} documents in your data')
text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=50)
texts = text_splitter.split_documents(data)
embeddings_model = OpenAIEmbeddings(openai_api_key=api_key)
doc_search = Chroma.from_documents(texts, embeddings_model)
print(f'You now have {len(texts)} texts in your data')
llm = OpenAI(temperature=0, openai_api_key=api_key)
chain = load_qa_chain(llm, chain_type="stuff")
# replace with your own query
query = "Which is not covered under bodily injury?"
docs = doc_search.similarity_search(query)
res = chain.run(input_documents=docs, question=query)
print(res)