Update app.py

app.py

@@ -1,7 +1,7 @@
 import torch
 import time
 import re
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
 # Load tokenizer and model
@@ -9,12 +9,23 @@ model_path = "supp03/E3L-AI-CHATBOT"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto")
 
+# Enable use_cache explicitly for faster generation
+model.config.use_cache = True
+
 # Simple context memory
 context_memory = {
     "last_question": None,
     "last_answer": None
 }
 
+# Cache the static prompt prefix tokenized once to save time
+static_prompt = (
+    "You are a warm, thoughtful, and emotionally intelligent mentor. "
+    "Your job is to provide deeply personalized, practical, and supportive guidance. "
+    "Avoid sounding robotic or generic. Respond with empathy, clarity, and human warmth.\n\n"
+)
+cached_static_prompt_ids = tokenizer(static_prompt, return_tensors="pt").input_ids.to(model.device)
+
 def generate_response(user_input):
     try:
         if not user_input.strip():
@@ -44,25 +55,35 @@ def generate_response(user_input):
         else:
             style_instruction = "Write in a natural, conversational paragraph form."
 
-        prompt = (
-            "You are a warm, thoughtful, and emotionally intelligent mentor. "
-            "Your job is to provide deeply personalized, practical, and supportive guidance. "
-            "Avoid sounding robotic or generic. Respond with empathy, clarity, and human warmth.\n\n"
+        # Combine full prompt text
+        full_prompt = (
+            f"{static_prompt}"
             f"{style_instruction}\n\n"
             f"{context}Question: {user_input.strip()}\n\nAnswer:"
         )
 
-        # Tokenize and generate
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+        # Tokenize only the user-dependent part to append to cached prompt tokens
+        user_prompt_ids = tokenizer(
+            full_prompt[len(static_prompt):],  # remove static part
+            return_tensors="pt"
+        ).input_ids.to(model.device)
+
+        # Concatenate cached static prompt tokens with user prompt tokens (skip overlapping BOS token if any)
+        # Avoid duplicate BOS token by slicing user_prompt_ids if it starts with tokenizer.bos_token_id
+        if user_prompt_ids[0, 0] == tokenizer.bos_token_id:
+            user_prompt_ids = user_prompt_ids[:, 1:]
+
+        input_ids = torch.cat([cached_static_prompt_ids, user_prompt_ids], dim=-1)
 
         start_time = time.time()
-        with torch.no_grad():
+        with torch.inference_mode():
             output_ids = model.generate(
                 input_ids,
-                max_new_tokens=300,    
-                do_sample=False,   
-                use_cache=True,        
-                pad_token_id=tokenizer.eos_token_id
+                max_new_tokens=100,
+                do_sample=False,
+                num_beams=1,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
             )
         elapsed_time = time.time() - start_time
 
@@ -82,7 +103,7 @@ def generate_response(user_input):
         context_memory["last_question"] = user_input
         context_memory["last_answer"] = model_answer
 
-        # Timing
+        # Timing display
         minutes = int(elapsed_time // 60)
         seconds = round(elapsed_time % 60, 1)
         time_display = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"
@@ -96,13 +117,13 @@ def generate_response(user_input):
 interface = gr.Interface(
     fn=generate_response,
     inputs=gr.Textbox(lines=5, placeholder="Ask your question here...", label="Ask your question"),
-    outputs=gr.Textbox(lines=1, label="Response"),
+    outputs=gr.Textbox(lines=2, label="Response"),
     title="Gemma Mentor Chatbot",
     description="Ask questions and get warm, supportive, emotionally intelligent responses. Type 'exit' to reset.",
     allow_flagging="never",
     theme="default",
-    live=False,          
-    api_name="/predict"  
+    live=False,
+    api_name="/predict"
 )
 
 if __name__ == "__main__":