Bahar110's picture
Update app.py
281a072 verified
Raw
History Blame Contribute Delete
3.43 kB
import os
import numpy as np
import faiss
import gradio as gr
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from groq import Groq
# ---------- Setup ----------
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
embedder = SentenceTransformer("all-MiniLM-L6-v2")
state = {"index": None, "chunks": None}
# ---------- Helper functions ----------
def extract_text(pdf_path):
reader = PdfReader(pdf_path)
text = ""
for page in reader.pages:
page_text = page.extract_text()
if page_text:
text += page_text + "\n"
return text
def chunk_text(text, chunk_size=300, overlap=50):
words = text.split()
chunks = []
i = 0
while i < len(words):
chunk = " ".join(words[i:i + chunk_size])
chunks.append(chunk)
i += chunk_size - overlap
return chunks
def build_index(chunks):
embeddings = embedder.encode(chunks, show_progress_bar=False)
embeddings = np.array(embeddings).astype("float32")
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)
return index
def retrieve(question, index, chunks, k=4):
q_emb = embedder.encode([question]).astype("float32")
distances, indices = index.search(q_emb, k)
return [chunks[i] for i in indices[0] if i < len(chunks)]
def generate_answer(question, context_chunks):
context = "\n\n".join(context_chunks)
prompt = f"""Answer the question based only on the context below.
If the answer isn't in the context, say you don't know.
Context:
{context}
Question: {question}
Answer:"""
response = client.chat.completions.create(
model="llama-3.3-70b-versatile",
messages=[{"role": "user", "content": prompt}],
temperature=0.2,
)
return response.choices[0].message.content
# ---------- App functions ----------
def process_pdf(pdf_file):
if pdf_file is None:
return "Please upload a PDF first."
text = extract_text(pdf_file.name)
if not text.strip():
return "Couldn't read any text from this PDF (it may be scanned images)."
chunks = chunk_text(text)
index = build_index(chunks)
state["index"] = index
state["chunks"] = chunks
return f"✅ PDF processed into {len(chunks)} chunks. You can ask questions now."
def answer_question(question, history):
if state["index"] is None:
return history + [[question, "Please upload and process a PDF first."]]
if not question or not question.strip():
return history
relevant_chunks = retrieve(question, state["index"], state["chunks"])
answer = generate_answer(question, relevant_chunks)
return history + [[question, answer]]
# ---------- UI ----------
with gr.Blocks(title="Chat with your PDF") as demo:
gr.Markdown("# 📄 Chat with your PDF\nUpload a PDF, then ask questions about it.")
with gr.Row():
pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
upload_btn = gr.Button("Process PDF", variant="primary")
status = gr.Textbox(label="Status", interactive=False)
chatbot = gr.Chatbot(label="Conversation")
question = gr.Textbox(label="Ask a question", placeholder="Type your question and press Enter")
upload_btn.click(process_pdf, inputs=pdf_input, outputs=status)
question.submit(answer_question, inputs=[question, chatbot], outputs=chatbot)
question.submit(lambda: "", None, question)
demo.launch()