doc-rag / app.py
SRINI123's picture
Update app.py
da91893 verified
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_groq import ChatGroq
from langchain.prompts import PromptTemplate
import os
GROQ_API = os.getenv("GROQ_API")
# Define embeddings and text splitter
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
# Initialize the LLM (ChatGroq)
llm = ChatGroq(
model="mixtral-8x7b-32768",
temperature=1.25,
max_tokens=512,
timeout=None,
api_key=GROQ_API
)
# Define the custom prompt template
custom_prompt = PromptTemplate.from_template("""
You are a helpful AI assistant.
Answer the question using only the following context:
{context}
Question: {input}
Provide a detailed and accurate response.
""")
# Create the document chain
document_chain = create_stuff_documents_chain(
llm=llm,
prompt=custom_prompt,
)
# Function to process the uploaded PDF and create the retriever
def process_pdf(file):
# Load the PDF file
loader = PyPDFLoader(file.name)
pages = loader.load_and_split(text_splitter)
# Create the vector store
vector_store = Chroma.from_documents(pages, embeddings)
# Create the retriever
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
# Create the QA chain
qa_chain = create_retrieval_chain(
retriever=retriever,
combine_docs_chain=document_chain
)
return qa_chain
# Function to get the answer from the QA chain
def get_answer(file, user_prompt):
# Process the PDF and create the QA chain
qa_chain = process_pdf(file)
# Get the answer
response = qa_chain.invoke({"input": user_prompt})
return response["answer"]
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# PDF Q&A with ChatGroq and LangChain")
with gr.Row():
pdf_input = gr.File(label="Upload PDF", type="filepath")
question_input = gr.Textbox(label="Ask a question", placeholder="Type your question here...")
output = gr.Textbox(label="Answer", interactive=False)
submit_button = gr.Button("Submit")
submit_button.click(fn=get_answer, inputs=[pdf_input, question_input], outputs=output)
# Launch the Gradio app
demo.launch()