Spaces:
Running
Running
import gradio as gr | |
from langchain_chroma import Chroma | |
from langchain_openai import OpenAIEmbeddings | |
from langchain_core.documents import Document | |
from langchain_openai import ChatOpenAI | |
from langchain_core.prompts import ChatPromptTemplate | |
from langchain_core.runnables import RunnablePassthrough | |
from PyPDF2 import PdfReader | |
import os | |
# Function to process the uploaded PDF and convert it to documents | |
def pdf_to_documents(pdf_file): | |
reader = PdfReader(pdf_file.name) | |
pages = [page.extract_text() for page in reader.pages] | |
documents = [Document(page_content=page, metadata={"page_number": idx + 1}) for idx, page in enumerate(pages)] | |
return documents | |
# Initialize vector store | |
def initialize_vectorstore(documents, api_key): | |
os.environ["OPENAI_API_KEY"] = api_key | |
embeddings = OpenAIEmbeddings() | |
vectorstore = Chroma.from_documents(documents, embedding=embeddings) | |
return vectorstore | |
# RAG retrieval and LLM chain | |
def rag_from_pdf(question, pdf_file, api_key): | |
documents = pdf_to_documents(pdf_file) | |
vectorstore = initialize_vectorstore(documents, api_key) | |
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 2}) # Retrieve top 2 relevant sections | |
# Initialize the LLM | |
llm = ChatOpenAI(model="gpt-3.5-turbo") | |
# Create a prompt template for combining context and question | |
prompt_template = """ | |
Answer this question using the provided context only. | |
{question} | |
Context: | |
{context} | |
""" | |
prompt = ChatPromptTemplate.from_messages([("human", prompt_template)]) | |
# Create a RAG chain combining retriever and LLM | |
rag_chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm | |
# Perform retrieval and return LLM's answer | |
response = rag_chain.invoke(question) | |
return response.content | |
# Gradio interface | |
with gr.Blocks() as app: | |
gr.Markdown("## PDF-based Question Answering with RAG") | |
# Input for OpenAI API Key | |
api_key_input = gr.Textbox(label="Enter your OpenAI API Key", type="password") | |
# File upload for the PDF | |
pdf_file_input = gr.File(label="Upload your PDF document") | |
# Question input | |
question_input = gr.Textbox(label="Ask a question related to the PDF") | |
# Output for the RAG response | |
rag_output = gr.Textbox(label="Generated Response", lines=10) | |
# Button to run RAG chain | |
rag_button = gr.Button("Ask Question") | |
# Functionality for the RAG chain | |
rag_button.click(rag_from_pdf, inputs=[question_input, pdf_file_input, api_key_input], outputs=rag_output) | |
# Launch Gradio app | |
app.launch() |