Update app.py

app.py

@@ -60,14 +60,16 @@ def generate_response(message, history, enable_reasoning, temperature, top_p, ma
             "top_p": top_p if top_p < 1 else 0.95,
             "do_sample": True,
             "max_new_tokens": max_tokens,
-            "eos_token_id": tokenizer.eos_token_id
+            "eos_token_id": tokenizer.eos_token_id,
+            "use_cache": False  # Disable cache to avoid the error
         }
     else:
         # Greedy search for non-reasoning
         generation_kwargs = {
             "do_sample": False,
             "max_new_tokens": max_tokens,
-            "eos_token_id": tokenizer.eos_token_id
+            "eos_token_id": tokenizer.eos_token_id,
+            "use_cache": False  # Disable cache to avoid the error
         }
     
     # Generate response
@@ -75,15 +77,8 @@ def generate_response(message, history, enable_reasoning, temperature, top_p, ma
         outputs = model.generate(tokenized_chat, **generation_kwargs)
     
     # Decode and extract the assistant's response
-    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Extract only the new response (after the last user message)
-    # This is a simple approach - you might need to adjust based on the model's output format
-    response_parts = full_response.split(message)
-    if len(response_parts) > 1:
-        response = response_parts[-1].strip()
-    else:
-        response = full_response.strip()
+    generated_tokens = outputs[0][tokenized_chat.shape[-1]:]  # Get only new tokens
+    response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
     
     return response
 
@@ -153,15 +148,19 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             return history
         
         message = history[-1][0]
-        response = generate_response(
-            message, 
-            history[:-1], 
-            enable_reasoning, 
-            temperature, 
-            top_p, 
-            max_tokens
-        )
-        history[-1][1] = response
+        try:
+            response = generate_response(
+                message, 
+                history[:-1], 
+                enable_reasoning, 
+                temperature, 
+                top_p, 
+                max_tokens
+            )
+            history[-1][1] = response
+        except Exception as e:
+            history[-1][1] = f"Error generating response: {str(e)}"
+        
         return history
     
     msg.submit(