Update app.py

app.py

@@ -1,36 +1,27 @@
 import os
 import json
-from PIL import Image
 import torch
 import gradio as gr
-from transformers import (
-    BlipProcessor,
-    AutoTokenizer,
-    Llama4ForConditionalGeneration,
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from ppt_parser import transfer_to_structure
 
-# ✅ Load Hugging Face token and model ID
-hf_token = os.getenv("HF_TOKEN")
-model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
+# ✅ Hugging Face token (optional if public + unauthenticated)
+hf_token = os.getenv("HF_TOKEN", None)
+model_id = "meta-llama/Llama-3.1-8B-Instruct"
 
-# ✅ Load processor, tokenizer, and model
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+# ✅ Load model + tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-
-model = Llama4ForConditionalGeneration.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_token,
-    attn_implementation="flex_attention",
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
 )
+llama_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
 
-# ✅ Global state
+# ✅ Global storage
 extracted_text = ""
-image_paths = []
 
-# ✅ Extract all text from parsed PPTX JSON
 def extract_text_from_pptx_json(parsed_json: dict) -> str:
     text = ""
     for slide in parsed_json.values():
@@ -47,57 +38,36 @@ def extract_text_from_pptx_json(parsed_json: dict) -> str:
                         text += para.get("text", "") + "\n"
     return text.strip()
 
-# ✅ Handle uploaded .pptx files
 def handle_pptx_upload(pptx_file):
-    global extracted_text, image_paths
+    global extracted_text
     tmp_path = pptx_file.name
-    parsed_json_str, image_paths = transfer_to_structure(tmp_path, "images")
+    parsed_json_str, _ = transfer_to_structure(tmp_path, "images")
     parsed_json = json.loads(parsed_json_str)
     extracted_text = extract_text_from_pptx_json(parsed_json)
     return extracted_text or "No readable text found in slides."
 
-# ✅ Ask Llama 4 Scout using image + text
 def ask_llama(question):
-    global extracted_text, image_paths
-
-    if not extracted_text and not image_paths:
-        return "Please upload and extract a PPTX file first."
-
-    # Use the first slide image
-    image = Image.open(image_paths[0]).convert("RGB")
-    vision_inputs = processor(images=image, return_tensors="pt").to(model.device)
-
-    # Prepare text prompt
-    prompt = f"<|user|>\n{extracted_text}\n\nQuestion: {question}<|end|>\n<|assistant|>\n"
-    text_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        output = model.generate(
-            input_ids=text_inputs["input_ids"],
-            pixel_values=vision_inputs["pixel_values"],
-            max_new_tokens=256,
-        )
+    global extracted_text
+    if not extracted_text:
+        return "Please upload a PPTX file first."
 
-    # Decode response
-    response = tokenizer.decode(
-        output[0][text_inputs["input_ids"].shape[-1]:],
-        skip_special_tokens=True
-    )
-    return response.strip()
+    prompt = f"<|user|>\nContext:\n{extracted_text}\n\nQuestion: {question}<|end|>\n<|assistant|>\n"
+    response = llama_pipe(prompt)[0]["generated_text"]
+    return response.replace(prompt, "").strip()
 
 # ✅ Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Llama-4-Scout Multimodal Study Assistant")
+    gr.Markdown("## 🧠 Study Assistant with LLaMA 3.1 8B")
 
     pptx_input = gr.File(label="📂 Upload PPTX File", file_types=[".pptx"])
-    extract_btn = gr.Button("📜 Extract Text + Slides")
+    extract_btn = gr.Button("📜 Extract Slide Text")
 
     extracted_output = gr.Textbox(label="📄 Slide Text", lines=10, interactive=False)
     extract_btn.click(handle_pptx_upload, inputs=[pptx_input], outputs=[extracted_output])
 
     question = gr.Textbox(label="❓ Ask a Question")
-    ask_btn = gr.Button("💬 Ask Scout")
-    ai_answer = gr.Textbox(label="🤖 Llama Answer", lines=6)
+    ask_btn = gr.Button("💬 Ask LLaMA")
+    ai_answer = gr.Textbox(label="🤖 LLaMA Answer", lines=6)
 
     ask_btn.click(ask_llama, inputs=[question], outputs=[ai_answer])