File size: 1,339 Bytes
3115bc6
 
 
 
 
 
 
 
d3db4aa
3115bc6
 
 
 
d02c531
3115bc6
 
 
d3db4aa
3115bc6
 
 
 
 
 
 
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
from langchain.document_loaders import PyPDFDirectoryLoader
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain import HuggingFaceHub
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.question_answering import load_qa_chain
import streamlit as st
from langchain.chains import RetrievalQA
import os
def main():
    pdf_folder_path = "./PDFfiles"
    loader = PyPDFDirectoryLoader(pdf_folder_path)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter (chunk_size=1000, chunk_overlap=200)
    texts = text_splitter.split_documents(docs)
    embeddings = HuggingFaceEmbeddings()
    vectordb = Chroma.from_documents(documents=texts,embedding=embeddings)
    llm = HuggingFaceHub(repo_id="google/flan-t5-small", model_kwargs={"temperature":0,"max_length":200}, huggingfacehub_api_token=os.getenv("API_TOKEN")) # type: ignore
    qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff",retriever=vectordb.as_retriever(search_type="mmr", search_kwargs={'fetch_k': 30}), return_source_documents=True)
    query = st.text_input("Ask a question: ")
    qa(inputs=query,return_only_outputs=True)
    result = qa({"query": query})
    st.write(result["result"])
    st.write(result["source_documents"][0])
main()