|
import os |
|
from langchain.embeddings.openai import OpenAIEmbeddings |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain.vectorstores import Chroma |
|
from langchain.document_loaders import PyPDFium2Loader |
|
from langchain.chains.question_answering import load_qa_chain |
|
|
|
from langchain.chat_models import ChatOpenAI |
|
|
|
|
|
class PDFQuery: |
|
def __init__(self): |
|
os.environ["OPENAI_API_KEY"] = "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" |
|
self.embeddings = OpenAIEmbeddings() |
|
self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=200) |
|
|
|
self.llm = ChatOpenAI(temperature=0) |
|
self.chain = None |
|
self.db = None |
|
|
|
def ask(self, question: str) -> str: |
|
if self.chain is None: |
|
response = "Please, add a document." |
|
else: |
|
docs = self.db.get_relevant_documents(question) |
|
response = self.chain.run(input_documents=docs, question=question) |
|
return response |
|
|
|
def ingest(self, file_path: os.PathLike) -> None: |
|
loader = PyPDFium2Loader(file_path) |
|
documents = loader.load() |
|
splitted_documents = self.text_splitter.split_documents(documents) |
|
self.db = Chroma.from_documents(splitted_documents, self.embeddings).as_retriever() |
|
|
|
self.chain = load_qa_chain(ChatOpenAI(temperature=0), chain_type="stuff") |
|
|
