-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
86 lines (74 loc) · 3.12 KB
/
main.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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import tempfile
import streamlit as st
from langchain_community.document_loaders import PyPDFLoader
import os
from langchain_community.vectorstores.chroma import Chroma
from langchain_google_genai import GoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from dotenv import load_dotenv
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain.prompts import ChatPromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
load_dotenv()
# Passing uploaded pdf document to PyPDFLoader
def pdf_loader(pdf):
if pdf:
temp = None
with tempfile.NamedTemporaryFile(delete=False) as temp_pdf:
temp_pdf.write(pdf.read())
loader = PyPDFLoader(file_path=temp_pdf.name)
pdf_content = loader.load()
temp = pdf_content
os.remove(temp_pdf.name)
return temp
# RAG chain for sending the loaded document into the llm for genetrating the response
def RAG_chain(document, query):
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(documents=document)
vectorstore = Chroma.from_documents(
documents = splits,
embedding = GoogleGenerativeAIEmbeddings(model='models/embedding-001'),
persist_directory='PDF-Chatbot/db'
)
retriever = vectorstore.as_retriever()
template = """
{context}
from the above context answer the user query [{question}] in a best possible way
"""
prompt = ChatPromptTemplate.from_template(
template=template
)
llm = GoogleGenerativeAI(
model='models/text-bison-001'
)
rag_chain = (
{'context': retriever, 'question': RunnablePassthrough()}
| prompt
| llm
| StrOutputParser()
)
return rag_chain.invoke(query)
def PDFChatbot(pdf, query):
document = pdf_loader(pdf)
return RAG_chain(document = document, query = query)
# main function of the program comprises of basic UI for user interaction, where the end user ask the prompt query related to the pdf uploaded .....
if __name__=='__main__':
try:
st.title("PDF Chatbot 📚")
with st.sidebar:
st.title('PDF Chatbot 📃')
google_api = st.text_input('Enter Google AI API-KEY:', type='password')
if not (google_api.startswith('AI') and len(google_api)==39):
st.warning('Please enter your credentials!', icon='⚠️')
else:
st.success('Proceed to entering your prompt message!', icon='👉')
st.markdown("You can make your API Token key from here → [Link](https://makersuite.google.com/app/apikey)")
os.environ['GOOGLE_API_KEY'] = google_api
pdf = st.file_uploader('Upload PDF file', type='pdf')
prompt = st.text_input('Enter the question you want to ask to the LLM', placeholder='Enter prompt your here.....')
result = PDFChatbot(pdf = pdf, query = prompt)
if st.button("Submit"):
st.balloons()
st.write(result)
except:
pass