Update app.py

app.py

@@ -2,175 +2,170 @@ import os
 import shutil
 from typing import List
 
-import torch
 import gradio as gr
 from PIL import Image
 
-# Unstructured for PDF parsing
+# PDF parsing
+from pypdf import PdfReader
 from unstructured.partition.pdf import partition_pdf
 from unstructured.partition.utils.constants import PartitionStrategy
 
-# Vision-language captioning (BLIP)
-from transformers import BlipProcessor, BlipForConditionalGeneration, CLIPProcessor, CLIPModel
-
-# Hugging Face Inference client for LLM
-from huggingface_hub import InferenceClient
+# Text splitting
+from langchain.text_splitter import CharacterTextSplitter
 
-# FAISS vectorstore
+# Vectorstore and embeddings
 from langchain_community.vectorstores import FAISS
-
-# Text embeddings
 from langchain_huggingface import HuggingFaceEmbeddings
 
+# Vision-language captioning (BLIP)
+from transformers import BlipProcessor, BlipForConditionalGeneration
+
+# LLM via HF Inference API
+from huggingface_hub import InferenceClient
+
 # ── Globals ───────────────────────────────────────────────────────────────────
 retriever = None
-current_pdf_name = None
-combined_texts: List[str] = []  # text chunks + captions
-combined_vectors: List[List[float]] = []
 pdf_text: str = ""
 
-# ── Setup ─────────────────────────────────────────────────────────────────────
+# ── Setup directories ──────────────────────────────────────────────────────────
 FIGURES_DIR = "figures"
 if os.path.exists(FIGURES_DIR):
     shutil.rmtree(FIGURES_DIR)
-else:
-    os.makedirs(FIGURES_DIR, exist_ok=True)
+os.makedirs(FIGURES_DIR, exist_ok=True)
+
+# ── Models & Clients ───────────────────────────────────────────────────────────
+hf_client = InferenceClient()  # uses HUGGINGFACEHUB_API_TOKEN
 
-# ── Clients & Models ───────────────────────────────────────────────────────────
-hf = InferenceClient()  # for chat completions
-txt_emb = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+# Embeddings model (local lightweight SBERT)
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
+# BLIP for image captioning
 blip_processor = BlipProcessor.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")
-
+# ── Helper functions ───────────────────────────────────────────────────────────
 
 def generate_caption(image_path: str) -> str:
     image = Image.open(image_path).convert("RGB")
     inputs = blip_processor(image, return_tensors="pt")
-    out = blip_model.generate(**inputs)
-    return blip_processor.decode(out[0], skip_special_tokens=True)
-
-
-def embed_texts(texts: List[str]) -> List[List[float]]:
-    return txt_emb.embed_documents(texts)
-
-
-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
+    outputs = blip_model.generate(**inputs)
+    return blip_processor.decode(outputs[0], skip_special_tokens=True)
 
 
 def process_pdf(pdf_file):
-    global retriever, current_pdf_name, combined_texts, combined_vectors, pdf_text
+    global retriever, pdf_text
     if pdf_file is None:
-        return None, "❌ Please upload a PDF file.", gr.update(interactive=False)
+        return None, "❌ Please upload a PDF.", gr.update(interactive=False)
 
-    current_pdf_name = os.path.basename(pdf_file.name)
-    # extract full text
-    from pypdf import PdfReader
+    # read full text
     reader = PdfReader(pdf_file.name)
-    pages = [page.extract_text() or "" for page in reader.pages]
+    pages = [p.extract_text() or "" for p in reader.pages]
     pdf_text = "\n\n".join(pages)
 
-    # rich parsing for images
+    # extract elements with images via unstructured
     try:
-        els = partition_pdf(
+        elements = 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,
         )
-        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"))]
+        text_elems = [e.text for e in elements if e.category not in ["Image","Table"] and e.text]
+        image_files = [os.path.join(FIGURES_DIR, f) for f in os.listdir(FIGURES_DIR)
+                       if f.lower().endswith((".png",".jpg",".jpeg"))]
     except:
-        texts = pages
-        imgs = []
+        text_elems = pages
+        image_files = []
+
+    # generate captions
+    captions = [generate_caption(img) for img in image_files]
 
-    # split text chunks
-    from langchain.text_splitter import CharacterTextSplitter
+    # split text into chunks
     splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
     chunks = []
-    for t in texts:
+    for t in text_elems:
         chunks.extend(splitter.split_text(t))
-    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
-
-    # Build FAISS index from precomputed multimodal vectors:
-    index = FAISS.from_embeddings(
-        combined_vectors,
-        combined_texts,
-    )
+
+    # combine text chunks and image captions
+    docs = chunks + captions
+
+    # embed and index
+    vectors = embeddings.embed_documents(docs)
+    index = FAISS.from_embeddings(vectors, docs)
     retriever = index.as_retriever(search_kwargs={"k":2})
-    status = f"✅ Indexed '{current_pdf_name}' — {len(chunks)} text chunks + {len(imgs)} images"
-    return current_pdf_name, status, gr.update(interactive=True)
+
+    status = f"✅ Indexed — {len(chunks)} text chunks + {len(captions)} captions"
+    return os.path.basename(pdf_file.name), status, gr.update(interactive=True)
 
 
-def ask_question(pdf_name,question):
-    global retriever
+def ask_question(pdf_name, question):
     if retriever is None:
-        return "❌ Please process a PDF first."
-    if not question.strip():
-        return "❌ Enter a question."
+        return "❌ Please upload + index a PDF first."
+    if not question:
+        return "❌ Please ask something."
+
     docs = retriever.get_relevant_documents(question)
-    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)
+    context = "\n\n".join(d.page_content for d in docs)
+    prompt = f"Use the following excerpts to answer:\n{context}\nQuestion: {question}\nAnswer:"
+
+    res = hf_client.chat_completion(
+        model="google/gemma-3-27b-it",
+        messages=[{"role":"user","content":prompt}],
+        max_tokens=128,
+        temperature=0.5,
+    )
     return res["choices"][0]["message"]["content"].strip()
 
 
-def generate_summary(): return ask_question(None,"Summarize:\n"+pdf_text[:2000])
+def generate_summary():
+    if not pdf_text:
+        return "❌ Please index a PDF first."
+    return ask_question(None, f"Summarize concisely:\n{pdf_text[:2000]}")
+
+
+def extract_keywords():
+    if not pdf_text:
+        return "❌ Please index first."
+    return ask_question(None, f"Extract 10–15 key terms:\n{pdf_text[:2000]}")
 
-def extract_keywords(): return ask_question(None,"Extract keywords:\n"+pdf_text[:2000])
 
-def clear_interface():
-    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)
+def clear_all():
+    global retriever, pdf_text
+    retriever = None
+    pdf_text = ""
+    shutil.rmtree(FIGURES_DIR, ignore_errors=True)
+    os.makedirs(FIGURES_DIR, exist_ok=True)
     return None, "", gr.update(interactive=False)
 
-# UI
-theme=gr.themes.Soft(primary_hue="indigo",secondary_hue="blue")
+# ── Gradio UI ────────────────────────────────────────────────────────────────
+theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")
 with gr.Blocks(theme=theme) as demo:
-    gr.Markdown("# DocQueryAI (True Multimodal RAG)")
+    gr.Markdown("# Multimodal RAG with HF & LangChain")
     with gr.Row():
         with gr.Column():
-            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)
+            pdf_disp = gr.Textbox(label="Active PDF", interactive=False)
+            pdf_file = gr.File(label="Upload PDF", type="filepath", file_types=[".pdf"])
+            btn_proc = gr.Button("📄 Process PDF")
+            status  = gr.Textbox(label="Status", interactive=False)
+
         with gr.Column():
-            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()
+            q_in   = gr.Textbox(label="Your question", interactive=False)
+            btn_ask = gr.Button("❓ Ask", interactive=False)
+            ans    = gr.Textbox(label="Answer", interactive=False)
+
+    with gr.Row():
+        btn_sum = gr.Button("📋 Summary", interactive=False)
+        sum_out = gr.Textbox(interactive=False)
+        btn_key = gr.Button("🏷️ Keywords", interactive=False)
+        key_out = gr.Textbox(interactive=False)
+
+    btn_clear = gr.Button("🗑️ Clear All")
+
+    btn_proc.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, [], sum_out)
+    btn_key.click(extract_keywords, [], key_out)
+    btn_clear.click(clear_all, [], [pdf_disp, status, q_in])
+
+if __name__ == "__main__":
+    demo.launch(debug=True)