Update app.py

app.py

@@ -1,33 +1,26 @@
 import gradio as gr
 import os
 import json
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from ppt_parser import transfer_to_structure
-from functools import lru_cache
+from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText
 
-# ✅ Get Hugging Face token from Space Secrets
+# ✅ Hugging Face Token for gated model access
 hf_token = os.getenv("HF_TOKEN")
 
-# ✅ Load summarization model (BART)
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-
-# ✅ Load Mistral model (memoized to avoid reloading)
-@lru_cache(maxsize=1)
-def load_mistral():
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
-    model = AutoModelForCausalLM.from_pretrained(
-        "mistralai/Mistral-7B-Instruct-v0.1",
-        torch_dtype=torch.float16,
-        device_map="auto",
-        token=hf_token
-    )
-    return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
-
-mistral_pipe = load_mistral()
+# ✅ Load Llama-4-Scout model and processor
+processor = AutoProcessor.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct", token=hf_token)
+model = AutoModelForImageTextToText.from_pretrained(
+    "meta-llama/Llama-4-Scout-17B-16E-Instruct",
+    torch_dtype=torch.float16,
+    device_map="auto",
+    token=hf_token
+)
 
-# ✅ Global variable to hold extracted content
+# ✅ Extracted data storage
 extracted_text = ""
+slide_images = []
 
 def extract_text_from_pptx_json(parsed_json: dict) -> str:
     text = ""
@@ -45,47 +38,58 @@ def extract_text_from_pptx_json(parsed_json: dict) -> str:
                         text += para.get("text", "") + "\n"
     return text.strip()
 
+# ✅ Handle uploaded .pptx
 def handle_pptx_upload(pptx_file):
-    global extracted_text
+    global extracted_text, slide_images
     tmp_path = pptx_file.name
-    parsed_json_str, _ = transfer_to_structure(tmp_path, "images")
+    parsed_json_str, image_paths = transfer_to_structure(tmp_path, "images")
     parsed_json = json.loads(parsed_json_str)
     extracted_text = extract_text_from_pptx_json(parsed_json)
+    slide_images = image_paths
     return extracted_text or "No readable text found in slides."
 
-def summarize_text():
-    global extracted_text
-    if not extracted_text:
-        return "Please upload and extract text from a PPTX file first."
-    summary = summarizer(extracted_text, max_length=200, min_length=50, do_sample=False)[0]['summary_text']
-    return summary
-
-def clarify_concept(question):
-    global extracted_text
-    if not extracted_text:
-        return "Please upload and extract text from a PPTX file first."
-    prompt = f"[INST] Use the following context to answer the question:\n\n{extracted_text}\n\nQuestion: {question} [/INST]"
-    response = mistral_pipe(prompt)[0]["generated_text"]
-    return response.replace(prompt, "").strip()
+# ✅ Ask a question using Llama 4 Scout
+def ask_llama(question):
+    global extracted_text, slide_images
+    if not extracted_text and not slide_images:
+        return "Please upload a PPTX file first."
+
+    inputs = {
+        "role": "user",
+        "content": []
+    }
+
+    # Add first image only (multimodal models may limit batch input size)
+    if slide_images:
+        image = Image.open(slide_images[0])
+        inputs["content"].append({"type": "image", "image": image})
+
+    # Add contextual text + question
+    context = f"{extracted_text}\n\nQuestion: {question}"
+    inputs["content"].append({"type": "text", "text": context})
+
+    outputs = processor(text=[inputs], return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**outputs, max_new_tokens=512)
+    result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return result
 
 # ✅ Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 AI-Powered Study Assistant for PowerPoint Lectures (Mistral 7B)")
+    gr.Markdown("## 🧠 Llama 4 Scout: PPTX-Based Multimodal Study Assistant")
 
     pptx_input = gr.File(label="📂 Upload PPTX File", file_types=[".pptx"])
-    extract_btn = gr.Button("📜 Extract & Summarize")
+    extract_btn = gr.Button("📜 Extract Text + Slides")
 
     extracted_output = gr.Textbox(label="📄 Extracted Text", lines=10, interactive=False)
-    summary_output = gr.Textbox(label="📝 Summary", interactive=False)
 
     extract_btn.click(handle_pptx_upload, inputs=[pptx_input], outputs=[extracted_output])
-    extract_btn.click(summarize_text, outputs=[summary_output])
 
     question = gr.Textbox(label="❓ Ask a Question")
-    ask_btn = gr.Button("💬 Ask Mistral")
-    ai_answer = gr.Textbox(label="🤖 Mistral Answer", lines=4)
+    ask_btn = gr.Button("💬 Ask Llama 4 Scout")
+    ai_answer = gr.Textbox(label="🤖 Llama Answer", lines=4)
 
-    ask_btn.click(clarify_concept, inputs=[question], outputs=[ai_answer])
+    ask_btn.click(ask_llama, inputs=[question], outputs=[ai_answer])
 
 if __name__ == "__main__":
     demo.launch()