Yatheshr's picture
Update app.py
e8fde11 verified
import os
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
import google.generativeai as genai
import tempfile
import shutil
# Local vector store path
INDEX_PATH = "faiss_index"
# Load or initialize vector DB
def load_vectorstore(pdf_path):
loader = PyPDFLoader(pdf_path)
docs = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
chunks = splitter.split_documents(docs)
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(chunks, embeddings)
# Save locally
vectorstore.save_local(INDEX_PATH)
return vectorstore
def load_existing_vectorstore():
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
return FAISS.load_local(INDEX_PATH, embeddings)
# Gemini: Ask question using optimized context
def ask_question(api_key, question, pdf_file):
if not api_key or not api_key.startswith("AI"):
return "❌ Invalid Gemini API Key. Please enter a valid key starting with 'AI'."
# Setup Gemini
genai.configure(api_key=api_key)
model = genai.GenerativeModel("models/gemini-1.5-flash")
# Get uploaded file path
pdf_path = pdf_file.name
# Load or create vectorstore
if os.path.exists(INDEX_PATH):
vs = load_existing_vectorstore()
else:
vs = load_vectorstore(pdf_path)
# Search for relevant chunks (context optimization)
top_docs = vs.similarity_search(question, k=3)
context = "\n\n".join([doc.page_content for doc in top_docs])
# Construct the prompt
prompt = f"""Use the context below to answer the question:\n\nContext:\n{context}\n\nQuestion: {question}"""
try:
response = model.generate_content(prompt)
return response.text
except Exception as e:
return f"❌ Gemini Error: {str(e)}"
# Gradio UI
with gr.Blocks(title="RAG Q&A with Gemini") as demo:
gr.Markdown("## πŸ“„ Upload PDF | πŸ” Ask Questions | πŸ”Œ Powered by Gemini")
api_key = gr.Textbox(label="πŸ” Gemini API Key", placeholder="Enter your Gemini API key")
pdf_file = gr.File(label="πŸ“‚ Upload PDF", file_types=[".pdf"])
question = gr.Textbox(label="❓ Your Question")
answer = gr.Textbox(label="πŸ’¬ Answer", lines=8)
ask_btn = gr.Button("πŸš€ Ask Gemini")
ask_btn.click(fn=ask_question, inputs=[api_key, question, pdf_file], outputs=answer)
# Run the app
if __name__ == "__main__":
demo.launch()