Update app.py

app.py

@@ -31,18 +31,18 @@ combined_texts: List[str] = [] # Combined text + image captions corpus
 FIGURES_DIR = "figures"
 if os.path.exists(FIGURES_DIR):
     shutil.rmtree(FIGURES_DIR)
-os.makedirs(FIGURES_DIR, exist_ok=True)
+ os.makedirs(FIGURES_DIR, exist_ok=True)
 
 # ── Clients & Models ───────────────────────────────────────────────────────────
-hf = InferenceClient()  # will use HUGGINGFACEHUB_API_TOKEN from env
+hf = InferenceClient()  # uses HUGGINGFACEHUB_API_TOKEN env var
 
-# BLIP captioner (small local model download)
+# BLIP captioner
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-blip_model     = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
 
 def generate_caption(image_path: str) -> str:
-    """Ask BLIP to caption a local image."""
+    """Generate caption for image via BLIP."""
     image = Image.open(image_path).convert("RGB")
     inputs = blip_processor(image, return_tensors="pt")
     out = blip_model.generate(**inputs)
@@ -50,23 +50,17 @@ def generate_caption(image_path: str) -> str:
 
 
 def embed_texts(texts: List[str]) -> List[List[float]]:
-    """
-    Call the HF embeddings endpoint using google/Gemma-Embeddings-v1.0.
-    """
-    resp = hf.embeddings(
-        model="google/Gemma-Embeddings-v1.0",
-        inputs=texts,
-    )
+    """Call HF inference embeddings endpoint."""
+    resp = hf.embeddings(model="google/Gemma-Embeddings-v1.0", inputs=texts)
     return resp["embeddings"]
 
 
 def process_pdf(pdf_file) -> str:
     """
-    Parse the PDF, caption images, combine text+captions, embed remotely,
-    build FAISS index, and prepare retriever. Falls back to text-only if poppler is missing.
+    Parse PDF, extract text and images, caption images,
+    embed all chunks remotely, build FAISS index.
     """
-    from pdf2image.exceptions import PDFInfoNotInstalledError
-    global current_pdf_name, retriever, combined_texts
+    global retriever, current_pdf_name, combined_texts
 
     if pdf_file is None:
         return "❌ Please upload a PDF file."
@@ -74,106 +68,52 @@ def process_pdf(pdf_file) -> str:
     pdf_path = pdf_file.name
     current_pdf_name = os.path.basename(pdf_path)
 
-    # Try rich parsing; fallback if poppler/pdfinfo is unavailable
+    # Attempt rich parsing
     try:
+        from pdf2image.exceptions import PDFInfoNotInstalledError
         elements = partition_pdf(
             filename=pdf_path,
             strategy=PartitionStrategy.HI_RES,
-            extract_image_block_types=["Image", "Table"],
+            extract_image_block_types=["Image","Table"],
             extract_image_block_output_dir=FIGURES_DIR,
         )
         text_elements = [el.text for el in elements if el.category not in ["Image","Table"] and el.text]
         image_files = [os.path.join(FIGURES_DIR, f) for f in os.listdir(FIGURES_DIR)
                        if f.lower().endswith((".png",".jpg",".jpeg"))]
-    except PDFInfoNotInstalledError:
-        # Fallback: text-only extraction
-        from PyPDF2 import PdfReader
+    except Exception:
+        # Fallback to text-only
+        from pypdf import PdfReader
         reader = PdfReader(pdf_path)
         text_elements = [page.extract_text() or "" for page in reader.pages]
         image_files = []
 
-    # Caption images if any
     captions = [generate_caption(img) for img in image_files]
-
     combined_texts = text_elements + captions
-    vectors = embed_texts(combined_texts)
-    index = FAISS.from_embeddings(texts=combined_texts, embeddings=vectors)
-    retriever = index.as_retriever(search_kwargs={"k": 2})
-
-    return f"✅ Indexed '{current_pdf_name}' — {len(text_elements)} text blocks + {len(captions)} image captions"
-    """
-    Parse the PDF, caption images, combine text+captions, embed remotely,
-    build FAISS index, and prepare retriever.
-    """
-    global current_pdf_name, retriever, combined_texts
-
-    if pdf_file is None:
-        return "❌ Please upload a PDF file."
-
-    # Save and name
-    pdf_path = pdf_file.name
-    current_pdf_name = os.path.basename(pdf_path)
 
-    # Extract blocks
-    elements = partition_pdf(
-        filename=pdf_path,
-        strategy=PartitionStrategy.HI_RES,
-        extract_image_block_types=["Image", "Table"],
-        extract_image_block_output_dir=FIGURES_DIR,
-    )
-
-    # Split text vs. images
-    text_elements = [
-        el.text for el in elements
-        if el.category not in ["Image", "Table"] and el.text
-    ]
-    image_files = [
-        os.path.join(FIGURES_DIR, f)
-        for f in os.listdir(FIGURES_DIR)
-        if f.lower().endswith((".png", ".jpg", ".jpeg"))
-    ]
-
-    # Caption images
-    captions = [generate_caption(img) for img in image_files]
-
-    # Combine
-    combined_texts = text_elements + captions
-
-    # Remote embeddings
     vectors = embed_texts(combined_texts)
-
-    # Build FAISS
-    index = FAISS.from_embeddings(
-        texts=combined_texts,
-        embeddings=vectors,
-    )
-    retriever = index.as_retriever(search_kwargs={"k": 2})
+    index = FAISS.from_embeddings(texts=combined_texts, embeddings=vectors)
+    retriever = index.as_retriever(search_kwargs={"k":2})
 
     return f"✅ Indexed '{current_pdf_name}' — {len(text_elements)} text blocks + {len(captions)} image captions"
 
 
 def ask_question(question: str) -> str:
-    """
-    Retrieve top-k chunks from FAISS and call chat_completions endpoint.
-    """
+    """Retrieve from FAISS and call chat completion."""
     global retriever
     if retriever is None:
-        return "❌ Please upload and process a PDF first."
+        return "❌ Please process a PDF first."
     if not question.strip():
         return "❌ Please enter a question."
 
     docs = retriever.get_relevant_documents(question)
     context = "\n\n".join(doc.page_content for doc in docs)
-
     prompt = (
-        "Use the following document excerpts to answer the question.\n\n"
-        f"{context}\n\n"
-        f"Question: {question}\n"
-        "Answer:"
+        "Use the following excerpts to answer the question:\n\n"
+        f"{context}\n\nQuestion: {question}\nAnswer:"
     )
     response = hf.chat_completion(
         model="google/gemma-3-27b-it",
-        messages=[{"role": "user", "content": prompt}],
+        messages=[{"role":"user","content":prompt}],
         max_tokens=128,
         temperature=0.5,
     )
@@ -181,7 +121,7 @@ def ask_question(question: str) -> str:
 
 
 def clear_interface():
-    """Reset state and clear extracted images."""
+    """Reset all state and clear extracted images."""
     global retriever, current_pdf_name, combined_texts
     retriever = None
     current_pdf_name = None
@@ -190,45 +130,24 @@ def clear_interface():
     os.makedirs(FIGURES_DIR, exist_ok=True)
     return ""
 
-# ── Gradio UI ────────────────────────────────────────────────────────────────
-theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")
-with gr.Blocks(theme=theme, css="""
-    .container { border-radius: 10px; padding: 15px; }
-    .pdf-active { border-left: 3px solid #6366f1;
-                  padding-left: 10px;
-                  background-color: rgba(99,102,241,0.1); }
-    .footer { text-align: center; margin-top: 30px;
-              font-size: 0.8em; color: #666; }
-    .main-title { text-align: center; font-size: 64px;
-                  font-weight: bold; margin-bottom: 20px; }
-""") as demo:
-    gr.Markdown("<div class='main-title'>DocQueryAI (Remote‐RAG)</div>")
-
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")) as demo:
+    gr.Markdown("# DocQueryAI (Remote‐RAG)")
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## 📄 Document Input")
-            pdf_file    = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath")
-            process_btn = gr.Button("📤 Process Document", variant="primary")
-            status_box  = gr.Textbox(label="Status", interactive=False)
-
+            pdf_file = gr.File(file_types=[".pdf"], type="filepath")
+            process_btn = gr.Button("Process PDF")
+            status_box = gr.Textbox(interactive=False)
         with gr.Column():
-            gr.Markdown("## ❓ Ask Questions")
-            question_input = gr.Textbox(lines=3,
-                                        placeholder="Enter your question here…")
-            ask_btn        = gr.Button("🔍 Ask Question", variant="primary")
-            answer_output  = gr.Textbox(label="Answer", lines=8, interactive=False)
-
-    clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-    gr.Markdown("<div class='footer'>Powered by HF Inference + BLIP + FAISS | Gradio</div>")
-
-    process_btn.click(fn=process_pdf,
-                      inputs=[pdf_file],
-                      outputs=[status_box])
-    ask_btn.click(fn=ask_question,
-                  inputs=[question_input],
-                  outputs=[answer_output])
-    clear_btn.click(fn=clear_interface,
-                    outputs=[status_box, answer_output])
+            question_input = gr.Textbox(lines=3)
+            ask_btn = gr.Button("Ask")
+            answer_output = gr.Textbox(interactive=False)
+    clear_btn = gr.Button("Clear All")
+
+    process_btn.click(fn=process_pdf, inputs=[pdf_file], outputs=[status_box])
+    ask_btn.click(fn=ask_question, inputs=[question_input], outputs=[answer_output])
+    clear_btn.click(fn=clear_interface, outputs=[status_box, answer_output])
 
 if __name__ == "__main__":
-    demo.launch(debug=True, share=True)
+    demo.launch()
+