Update app.py

app.py

@@ -1,78 +1,99 @@
-import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import PyPDF2
-
-# ----------------- Model Setup -----------------
-model_name = "ibm-granite/granite-3.2-2b-instruct"
+import gradio as gr
+from PyPDF2 import PdfReader
 
-# Some models require trust_remote_code
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# -----------------------------
+# Load model with fallback
+# -----------------------------
+MODEL_NAME = "distilgpt2"  # fallback if granite not available
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto" if torch.cuda.is_available() else None,
-    trust_remote_code=True
-)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True,
+        torch_dtype=torch.float32,
+        device_map="auto"
+    )
+except Exception as e:
+    print(f"⚠️ Model load failed: {e}")
+    print("Falling back to distilgpt2...")
+    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+    model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
-# Ensure pad token is set
+# Ensure pad token exists
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
 
-# ----------------- Core Functions -----------------
-def generate_response(prompt, max_length=512):
+# -----------------------------
+# PDF text extraction
+# -----------------------------
+def extract_text_from_pdf(pdf_file):
+    text = ""
     try:
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256)
-
-        if torch.cuda.is_available():
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_length=max_length,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
-
-        # Slice only the generated continuation
-        response_ids = outputs[0][inputs["input_ids"].shape[-1]:]
-        response = tokenizer.decode(response_ids, skip_special_tokens=True)
-        return response.strip()
+        # pdf_file is a dict in Gradio (with 'name' key)
+        file_path = pdf_file.name if hasattr(pdf_file, "name") else pdf_file
+        reader = PdfReader(file_path)
+        for page in reader.pages:
+            text += page.extract_text() or ""
     except Exception as e:
-        return f"⚠️ Error generating response: {str(e)}"
+        text = f"❌ Error reading PDF: {e}"
+    return text
 
 
-def extract_text_from_pdf(pdf_file):
-    if pdf_file is None:
-        return ""
-    try:
-        pdf_reader = PyPDF2.PdfReader(pdf_file)
-        text = ""
-        for page in pdf_reader.pages:
-            page_text = page.extract_text()
-            if page_text:
-                text += page_text + "\n"
-        return text.strip()
-    except Exception as e:
-        return f"⚠️ Error reading PDF: {str(e)}"
+# -----------------------------
+# Response generation
+# -----------------------------
+def generate_response(pdf_file, user_query):
+    pdf_text = extract_text_from_pdf(pdf_file)
+    if not pdf_text.strip():
+        return "⚠️ Could not extract text from the PDF."
 
+    # Build prompt
+    prompt = f"""
+You are a Sustainable Smart City Assistant. 
+Here is context from a PDF document:
 
-def eco_tips_generator(problem_keywords):
-    if not problem_keywords.strip():
-        return "⚠️ Please enter some keywords."
-    prompt = (
-        f"Generate practical and actionable eco-friendly tips for sustainable living "
-        f"related to: {problem_keywords}. Provide specific solutions and suggestions:"
-    )
-    return generate_response(prompt, max_length=400)
+{pdf_text[:2000]}  # limit context length to avoid overflow
+
+User query: {user_query}
+
+Answer clearly and concisely:
+"""
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+with gr.Blocks() as app:
+    gr.Markdown("## 🌍 Sustainable Smart City Assistant")
+
+    with gr.Row():
+        pdf_file = gr.File(label="Upload PDF", type="file")
+        user_query = gr.Textbox(label="Ask your question")
+
+    output = gr.Textbox(label="Assistant Response")
 
+    submit = gr.Button("Get Answer")
+    submit.click(fn=generate_response, inputs=[pdf_file, user_query], outputs=output)
 
-def policy_summarization(pdf_file, policy_text):
-    if pdf_file is not None:
-        content = extract_text_from_pdf(pdf_file)
-    else:
-        content = polic
+# -----------------------------
+# Run app
+# -----------------------------
+if __name__ == "__main__":
+    app.launch()