Update app.py

app.py

@@ -2,44 +2,52 @@ import os
 import shutil
 from typing import List
 
+import torch
 import gradio as gr
 from PIL import Image
 
-# Unstructured for rich PDF parsing
+# Unstructured for PDF parsing
 from unstructured.partition.pdf import partition_pdf
 from unstructured.partition.utils.constants import PartitionStrategy
 
 # Vision-language captioning (BLIP)
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import BlipProcessor, BlipForConditionalGeneration, CLIPProcessor, CLIPModel
 
-# Hugging Face Inference client
+# Hugging Face Inference client for LLM
 from huggingface_hub import InferenceClient
 
 # FAISS vectorstore
-from langchain.vectorstores.faiss import FAISS
+from langchain_community.vectorstores import FAISS
+
+# Text embeddings
+from langchain_huggingface import HuggingFaceEmbeddings
 
 # ── Globals ───────────────────────────────────────────────────────────────────
-retriever = None               # FAISS retriever for multimodal content
-current_pdf_name = None        # Name of the currently loaded PDF
-combined_texts: List[str] = [] # Combined text + image captions corpus
-pdf_text: str = ""           # Full PDF text for summary/keywords
+retriever = None
+current_pdf_name = None
+combined_texts: List[str] = []  # text chunks + captions
+combined_vectors: List[List[float]] = []
+pdf_text: str = ""
 
-# ── Setup: directories ─────────────────────────────────────────────────────────
+# ── Setup ─────────────────────────────────────────────────────────────────────
 FIGURES_DIR = "figures"
 if os.path.exists(FIGURES_DIR):
     shutil.rmtree(FIGURES_DIR)
-os.makedirs(FIGURES_DIR, exist_ok=True)
+else:
+    os.makedirs(FIGURES_DIR, exist_ok=True)
 
 # ── Clients & Models ───────────────────────────────────────────────────────────
-hf = InferenceClient()  # uses HUGGINGFACEHUB_API_TOKEN env var
+hf = InferenceClient()  # for chat completions
+txt_emb = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-# 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")
+
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 
 
 def generate_caption(image_path: str) -> str:
-    """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)
@@ -47,159 +55,118 @@ def generate_caption(image_path: str) -> str:
 
 
 def embed_texts(texts: List[str]) -> List[List[float]]:
-    """Call HF inference embeddings endpoint."""
-    resp = hf.embeddings(model="google/Gemma-Embeddings-v1.0", inputs=texts)
-    return resp["embeddings"]
+    return txt_emb.embed_documents(texts)
 
 
-def process_pdf(pdf_file):
-    """
-    Reads & extracts text and images from the PDF, captions images,
-    splits & embeds chunks, builds FAISS index, and stores full text.
-    Returns filename, status, and enables Q&A box.
-    """
-    global retriever, current_pdf_name, combined_texts, pdf_text
+def embed_images(image_paths: List[str]) -> List[List[float]]:
+    feats = []
+    for p in image_paths:
+        img = Image.open(p).convert("RGB")
+        inputs = clip_processor(images=img, return_tensors="pt")
+        with torch.no_grad():
+            v = clip_model.get_image_features(**inputs)
+        feats.append(v[0].cpu().tolist())
+    return feats
+
 
+def process_pdf(pdf_file):
+    global retriever, current_pdf_name, combined_texts, combined_vectors, pdf_text
     if pdf_file is None:
         return None, "❌ Please upload a PDF file.", gr.update(interactive=False)
 
     current_pdf_name = os.path.basename(pdf_file.name)
-    # extract full text for summary/keywords
+    # extract full text
     from pypdf import PdfReader
     reader = PdfReader(pdf_file.name)
     pages = [page.extract_text() or "" for page in reader.pages]
     pdf_text = "\n\n".join(pages)
 
-    # parse with unstructured for images
+    # rich parsing for images
     try:
-        elements = partition_pdf(
+        els = partition_pdf(
             filename=pdf_file.name,
             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 Exception:
-        text_elements = pages
-        image_files = []
-
-    captions = [generate_caption(img) for img in image_files]
-    # split text elements into chunks
+        texts = [e.text for e in els if e.category not in ["Image","Table"] and e.text]
+        imgs = [os.path.join(FIGURES_DIR,f) for f in os.listdir(FIGURES_DIR)
+                if f.lower().endswith((".png",".jpg",".jpeg"))]
+    except:
+        texts = pages
+        imgs = []
+
+    # split text chunks
     from langchain.text_splitter import CharacterTextSplitter
     splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
     chunks = []
-    for t in text_elements:
+    for t in texts:
         chunks.extend(splitter.split_text(t))
-    combined_texts = chunks + captions
-
-    vectors = embed_texts(combined_texts)
-    index = FAISS.from_embeddings(texts=combined_texts, embeddings=vectors)
+    caps = [generate_caption(i) for i in imgs]
+
+    # embed
+    tvecs = embed_texts(chunks + caps)
+    ivecs = embed_images(imgs)
+    # align dims: captions embedded twice? simple: drop caps embeddings from tvecs
+    text_count = len(chunks)
+    cap_count = len(caps)
+    # use text embeddings for text and clip for images
+    combined_texts = chunks + caps
+    combined_vectors = tvecs[:text_count] + ivecs
+
+    index = FAISS.from_embeddings(texts=combined_texts, embeddings=combined_vectors)
     retriever = index.as_retriever(search_kwargs={"k":2})
-
-    status = f"✅ Indexed '{current_pdf_name}' — {len(chunks)} text chunks + {len(captions)} image captions"
+    status = f"✅ Indexed '{current_pdf_name}' — {len(chunks)} text chunks + {len(imgs)} images"
     return current_pdf_name, status, gr.update(interactive=True)
 
 
-def ask_question(pdf_name, question):
-    """Retrieve relevant chunks and generate answer via remote LLM."""
+def ask_question(pdf_name,question):
     global retriever
     if retriever is None:
-        return "❌ Please upload and index a PDF first."
+        return "❌ Please process a PDF first."
     if not question.strip():
-        return "❌ Please enter a question."
-
+        return "❌ 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\nQuestion: {question}\nAnswer:"
-    )
-    response = hf.chat_completion(
-        model="google/gemma-3-27b-it",
-        messages=[{"role":"user","content":prompt}],
-        max_tokens=128,
-        temperature=0.5,
-    )
-    return response["choices"][0]["message"]["content"].strip()
-
-
-def generate_summary():
-    """Ask remote LLM for concise summary using full text."""
-    if not pdf_text:
-        return "❌ Please upload and index a PDF first."
-    ctx = pdf_text[:2000]
-    resp = hf.chat_completion(
-        model="google/gemma-3-27b-it",
-        messages=[{"role":"user","content":f"Summarize concisely:\n\n{ctx}..."}],
-        max_tokens=150,
-        temperature=0.5,
-    )
-    return resp["choices"][0]["message"]["content"].strip()
-
-
-def extract_keywords():
-    """Ask remote LLM to extract key terms from full text."""
-    if not pdf_text:
-        return "❌ Please upload and index a PDF first."
-    ctx = pdf_text[:2000]
-    resp = hf.chat_completion(
-        model="google/gemma-3-27b-it",
-        messages=[{"role":"user","content":f"Extract 10-15 key terms:\n\n{ctx}..."}],
-        max_tokens=60,
-        temperature=0.5,
-    )
-    return resp["choices"][0]["message"]["content"].strip()
+    ctx = "\n\n".join(d.page_content for d in docs)
+    prompt = f"Use contexts:\n{ctx}\nQuestion:{question}\nAnswer:"
+    res = hf.chat_completion(model="google/gemma-3-27b-it",messages=[{"role":"user","content":prompt}],max_tokens=128)
+    return res["choices"][0]["message"]["content"].strip()
 
 
+def generate_summary(): return ask_question(None,"Summarize:\n"+pdf_text[:2000])
+
+def extract_keywords(): return ask_question(None,"Extract keywords:\n"+pdf_text[:2000])
+
 def clear_interface():
-    """Reset state and clear extracted images."""
-    global retriever, current_pdf_name, combined_texts, pdf_text
-    retriever = None
-    current_pdf_name = None
-    combined_texts = []
-    pdf_text = ""
-    shutil.rmtree(FIGURES_DIR, ignore_errors=True)
-    os.makedirs(FIGURES_DIR, exist_ok=True)
+    global retriever,combined_texts,combined_vectors,pdf_text
+    retriever=None
+    combined_texts=[]
+    combined_vectors=[]
+    pdf_text=""
+    shutil.rmtree(FIGURES_DIR,ignore_errors=True)
+    os.makedirs(FIGURES_DIR,exist_ok=True)
     return None, "", gr.update(interactive=False)
 
-# ── 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 (Multimodal RAG)</div>")
+# UI
+theme=gr.themes.Soft(primary_hue="indigo",secondary_hue="blue")
+with gr.Blocks(theme=theme) as demo:
+    gr.Markdown("# DocQueryAI (True Multimodal RAG)")
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## 📄 Document Input")
-            pdf_display   = gr.Textbox(label="Active Document", interactive=False, elem_classes="pdf-active")
-            pdf_file      = gr.File(file_types=[".pdf"], type="filepath")
-            upload_button = gr.Button("📤 Process Document", variant="primary")
-            status_box    = gr.Textbox(label="Status", interactive=False)
+            pdf_disp=gr.Textbox(label="Active Document",interactive=False)
+            pdf_file=gr.File(file_types=[".pdf"],type="filepath")
+            btn_process=gr.Button("Process PDF")
+            status=gr.Textbox(interactive=False)
         with gr.Column():
-            gr.Markdown("## ❓ Ask Questions")
-            question_input = gr.Textbox(lines=3, placeholder="Enter your question here…", interactive=False)
-            ask_button     = gr.Button("🔍 Ask Question", variant="primary", interactive=False)
-            answer_output  = gr.Textbox(label="Answer", lines=8, interactive=False)
-    with gr.Row():
-        summary_button  = gr.Button("📋 Generate Summary", variant="secondary", interactive=False)
-        summary_output  = gr.Textbox(label="Summary", lines=4, interactive=False)
-        keywords_button = gr.Button("🏷️ Extract Keywords", variant="secondary", interactive=False)
-        keywords_output = gr.Textbox(label="Keywords", lines=4, interactive=False)
-    clear_button = gr.Button("🗑️ Clear All", variant="secondary")
-    gr.Markdown("<div class='footer'>Powered by HF Inference + FAISS + BLIP | Gradio</div>")
-
-    upload_button.click(process_pdf, [pdf_file], [pdf_display, status_box, question_input])
-    ask_button.click(ask_question, [pdf_display, question_input], answer_output)
-    summary_button.click(generate_summary, [], summary_output)
-    keywords_button.click(extract_keywords, [], keywords_output)
-    clear_button.click(clear_interface, [], [pdf_display, status_box, question_input])
-
-if __name__ == "__main__":
-    demo.launch(debug=True)
-
-
+            q_in=gr.Textbox(lines=3,interactive=False)
+            btn_ask=gr.Button("Ask")
+            ans=gr.Textbox(interactive=False)
+    btn_sum=gr.Button("Summary",interactive=False);out_sum=gr.Textbox(interactive=False)
+    btn_key=gr.Button("Keywords",interactive=False);out_key=gr.Textbox(interactive=False)
+    btn_clear=gr.Button("Clear All")
+    btn_process.click(process_pdf,[pdf_file],[pdf_disp,status,q_in])
+    btn_ask.click(ask_question,[pdf_disp,q_in],ans)
+    btn_sum.click(generate_summary,[],out_sum)
+    btn_key.click(extract_keywords,[],out_key)
+    btn_clear.click(clear_interface,[],[pdf_disp,status,q_in])
+if __name__=="__main__": demo.launch()