Update app.py

app.py

@@ -1,143 +1,122 @@
-import warnings
-warnings.filterwarnings("ignore")
-
-import uuid
-import gradio as gr
 import os
+import tempfile
+import gradio as gr
 
-from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from huggingface_hub import InferenceClient
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.chains import RetrievalQA
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from pypdf import PdfReader
 
 
-# -------------------------
-# LLM (Hugging Face Router API)
-# -------------------------
+# -----------------------------
+# 1. LLM CLIENT (Router API)
+# -----------------------------
+HF_MODEL = "Qwen/Qwen2.5-7B-Instruct"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-repo_id = "Qwen/Qwen2.5-7B-Instruct"
-
-llm_endpoint = HuggingFaceEndpoint(
-    repo_id=repo_id,
-    task="text-generation",
-    max_new_tokens=512,
-    temperature=0.1,
-    huggingfacehub_api_token=os.getenv("HF_TOKEN"),
-    provider="hf-inference",   # forces router.huggingface.co
-    model_kwargs={"wait_for_model": True},
+client = InferenceClient(
+    model=HF_MODEL,
+    token=HF_TOKEN
 )
 
-llm = ChatHuggingFace(llm=llm_endpoint)
-
 
-# -------------------------
-# Document loader
-# -------------------------
+def call_llm(prompt: str) -> str:
+    """Call the HF Router API using the new chat.completions endpoint."""
+    response = client.chat.completions.create(
+        model=HF_MODEL,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=512,
+        temperature=0.1,
+    )
+    return response.choices[0].message["content"]
 
-def document_loader(file_path: str):
-    loader = PyPDFLoader(file_path)
-    return loader.load()
 
+# -----------------------------
+# 2. PDF → TEXT
+# -----------------------------
+def extract_text_from_pdf(pdf_file):
+    reader = PdfReader(pdf_file)
+    text = ""
+    for page in reader.pages:
+        text += page.extract_text() or ""
+    return text
 
-# -------------------------
-# Text splitter
-# -------------------------
 
-def text_splitter(data):
+# -----------------------------
+# 3. TEXT → CHUNKS → EMBEDDINGS → CHROMA
+# -----------------------------
+def build_vectorstore(text: str):
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=50,
-        length_function=len,
+        chunk_size=800,
+        chunk_overlap=150
     )
-    return splitter.split_documents(data)
+    chunks = splitter.split_text(text)
 
-
-# -------------------------
-# Embedding model
-# -------------------------
-
-def embedding_model():
-    return HuggingFaceEmbeddings(
+    embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
 
+    vectordb = Chroma.from_texts(
+        texts=chunks,
+        embedding=embeddings
+    )
 
-# -------------------------
-# Vector DB
-# -------------------------
+    return vectordb
 
-def vector_database(chunks):
-    chunks = [
-        c for c in chunks
-        if c.page_content and c.page_content.strip() and len(c.page_content.strip()) > 10
-    ]
 
-    vectordb = Chroma.from_documents(
-        documents=chunks,
-        embedding=embedding_model(),
-        collection_name=f"pdf_chat_{uuid.uuid4().hex}",
-    )
-    return vectordb
+# -----------------------------
+# 4. RAG PIPELINE
+# -----------------------------
+def rag_answer(pdf_file, question):
+    if pdf_file is None:
+        return "Please upload a PDF first."
 
+    # Extract text
+    text = extract_text_from_pdf(pdf_file)
 
-# -------------------------
-# Retriever
-# -------------------------
+    # Build vectorstore
+    vectordb = build_vectorstore(text)
 
-def build_retriever(file_path: str):
-    docs = document_loader(file_path)
-    chunks = text_splitter(docs)
-    vectordb = vector_database(chunks)
-    return vectordb.as_retriever()
+    # Retrieve context
+    docs = vectordb.similarity_search(question, k=4)
+    context = "\n\n".join([d.page_content for d in docs])
 
+    # Build prompt
+    prompt = f"""
+You are a helpful assistant. Use ONLY the context below to answer.
 
-# -------------------------
-# QA Chain
-# -------------------------
+Context:
+{context}
 
-def retriever_qa(file, query):
-    retr = build_retriever(file)
+Question: {question}
 
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=retr,
-        return_source_documents=False,
-    )
+Answer clearly and concisely:
+"""
 
-    result = qa.invoke({"query": query})
-    return result["result"]
-
-
-# -------------------------
-# Gradio interface
-# -------------------------
-
-rag_application = gr.Interface(
-    fn=retriever_qa,
-    allow_flagging="never",
-    inputs=[
-        gr.File(
-            label="Upload PDF File",
-            file_count="single",
-            file_types=[".pdf"],
-            type="filepath",
-        ),
-        gr.Textbox(
-            label="Input Query",
-            lines=2,
-            placeholder="Type your question here...",
-        ),
-    ],
-    outputs=gr.Textbox(label="Output"),
-    title="PDF QA Chatbot",
-    description=(
-        "Upload a PDF document and ask any question. "
-        "The chatbot will try to answer using the provided document."
-    ),
-)
+    # LLM call
+    answer = call_llm(prompt)
+    return answer
+
+
+# -----------------------------
+# 5. GRADIO UI
+# -----------------------------
+with gr.Blocks(title="PDF Analyzer") as demo:
+    gr.Markdown("# 📄 PDF Analyzer\nAsk questions about your PDF using RAG + HF Router API.")
+
+    with gr.Row():
+        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
+        question_input = gr.Textbox(label="Your question")
+
+    answer_output = gr.Textbox(label="Answer")
+
+    run_button = gr.Button("Analyze")
+
+    run_button.click(
+        fn=rag_answer,
+        inputs=[pdf_input, question_input],
+        outputs=answer_output
+    )
 
-if __name__ == "__main__":
-    rag_application.launch()
+demo.launch()