updated

app.py

@@ -1,18 +1,14 @@
 import gradio as gr
-import os
-import tempfile
-import textwrap
 from datetime import datetime
-from typing import List, Dict, Any, Optional
-
+from pathlib import Path
 from src.model_loader import load_local_model
 from src.conversation import ConversationMemory
 from src.chatbot import LocalChatbot
 
-# ----------------------
-# Model setup
-# ----------------------
-MODEL_PATH = "models/gemma-2-2b-it-Q4_K_M"  # quantized 2B
+# ----------------------------
+# Model & Memory
+# ----------------------------
+MODEL_PATH = "togethercomputer/Gemini-2-2B"  # public HF model
 llm = load_local_model(MODEL_PATH, device=-1)  # CPU
 memory = ConversationMemory(max_len=60)
 bot = LocalChatbot(llm, memory)
@@ -27,49 +23,10 @@ INTENT_TEMPLATES = {
 def now_ts():
     return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
-# ----------------------
-# Export TXT/PDF
-# ----------------------
-def export_chat_files(history: List[Dict[str, Any]]) -> Dict[str, Optional[str]]:
-    tmpdir = tempfile.gettempdir()
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    txt_path = os.path.join(tmpdir, f"chat_history_{timestamp}.txt")
-
-    with open(txt_path, "w", encoding="utf-8") as f:
-        for msg in history:
-            content = msg.get("content", "")
-            lines = content.splitlines()
-            lines = [l.replace("USER:", "").replace("ASSISTANT:", "").strip() for l in lines]
-            f.write("\n".join(lines).strip() + "\n")
-            f.write("-" * 60 + "\n")
-
-    pdf_path = None
-    try:
-        from reportlab.lib.pagesizes import A4
-        from reportlab.pdfgen import canvas
-
-        pdf_path = os.path.join(tmpdir, f"chat_history_{timestamp}.pdf")
-        c = canvas.Canvas(pdf_path, pagesize=A4)
-        width, height = A4
-        margin = 40
-        textobject = c.beginText(margin, height - margin)
-        textobject.setFont("Helvetica", 10)
-        with open(txt_path, "r", encoding="utf-8") as fh:
-            for line in fh:
-                for wrapped in textwrap.wrap(line.rstrip(), 100):
-                    textobject.textLine(wrapped)
-        c.drawText(textobject)
-        c.showPage()
-        c.save()
-    except:
-        pdf_path = None
-
-    return {"txt": txt_path, "pdf": pdf_path}
-
-# ----------------------
-# Core chat function
-# ----------------------
-def generate_reply(user_msg: str, history: Optional[List[Dict[str, Any]]]):
+# ----------------------------
+# Chat Function
+# ----------------------------
+def generate_reply(user_msg, history=None):
     if history is None:
         history = []
 
@@ -93,6 +50,7 @@ def generate_reply(user_msg: str, history: Optional[List[Dict[str, Any]]]):
 
     history.append({"role": "user", "content": user_msg})
     history.append({"role": "assistant", "content": bot_reply_ts})
+
     try:
         memory.add(user_msg, bot_reply)
     except:
@@ -100,26 +58,16 @@ def generate_reply(user_msg: str, history: Optional[List[Dict[str, Any]]]):
 
     return history
 
-# ----------------------
+# ----------------------------
 # Gradio UI
-# ----------------------
+# ----------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    with gr.Row():
-        with gr.Column(scale=1, min_width=220):
-            gr.Markdown("### ⚡ Tools & Export")
-            new_chat_btn = gr.Button("➕ New Chat")
-            export_btn = gr.Button("📥 Export TXT/PDF")
-
-        with gr.Column(scale=3):
-            gr.Markdown("<h3>Smart Learning Assistant - Tayyab</h3>")
-            chatbot = gr.Chatbot(height=480, type="messages")
+    gr.Markdown("## ⚡ Smart Learning Assistant - Tayyab")
 
-            with gr.Row():
-                msg = gr.Textbox(placeholder="Type a message", show_label=False, lines=3)
-                send_btn = gr.Button("Send")
-
-    file_txt = gr.File(visible=False)
-    file_pdf = gr.File(visible=False)
+    chatbot = gr.Chatbot(height=480)
+    msg = gr.Textbox(placeholder="Type your message...", lines=3)
+    send_btn = gr.Button("Send")
+    new_chat_btn = gr.Button("➕ New Chat")
 
     send_btn.click(generate_reply, inputs=[msg, chatbot], outputs=[chatbot])
     msg.submit(generate_reply, inputs=[msg, chatbot], outputs=[chatbot])
@@ -129,13 +77,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         return []
     new_chat_btn.click(new_chat, outputs=[chatbot])
 
-    def export_handler(history):
-        files = export_chat_files(history or [])
-        return (
-            gr.update(value=files.get("txt"), visible=True),
-            gr.update(value=files.get("pdf"), visible=bool(files.get("pdf")))
-        )
-    export_btn.click(export_handler, inputs=[chatbot], outputs=[file_txt, file_pdf])
-
 if __name__ == "__main__":
     demo.launch()

src/model_loader.py

@@ -1,18 +1,17 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-def load_local_model(model_path, device=0):
+def load_local_model(model_path: str, device: int = -1):
     """
-    Loads a local quantized model for CPU or GPU.
-    device=-1 => CPU, device>=0 => GPU
+    Loads a Hugging Face model on CPU/GPU.
+    device=-1 → CPU, device=0 → first GPU
     """
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     model = AutoModelForCausalLM.from_pretrained(model_path)
-    
-    # Use pipeline for text generation
+
     generator = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        device=device
+        device=device  # -1=CPU
     )
     return generator