Update app.py

app.py

@@ -56,18 +56,31 @@ class ChatModel:
         # Format prompt
         prompt = config["template"].format(message=message)
         
-        # Tokenize input
-        inputs = self.tokenizers[model_name](prompt, return_tensors="pt").to("cuda")
+        # Tokenize input with proper max_length handling
+        inputs = self.tokenizers[model_name](
+            prompt,
+            return_tensors="pt",
+            max_length=2048,
+            truncation=True
+        ).to("cuda")
         
-        # Generate response
-        outputs = self.models[model_name].generate(
-            **inputs,
-            max_new_tokens=384,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=self.tokenizers[model_name].eos_token_id
-        )
+        # Generation parameters
+        generation_kwargs = {
+            "inputs": inputs.input_ids,
+            "max_new_tokens": 384,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "do_sample": True,
+            "pad_token_id": self.tokenizers[model_name].eos_token_id
+        }
+        
+        # Phi-3 specific workaround
+        if "phi-3" in model_name:
+            generation_kwargs["attention_mask"] = inputs.attention_mask
+            generation_kwargs.pop("inputs")
+            generation_kwargs["input_ids"] = inputs.input_ids
+        
+        outputs = self.models[model_name].generate(**generation_kwargs)
         
         # Decode response
         response = self.tokenizers[model_name].decode(
@@ -77,7 +90,7 @@ class ChatModel:
         
         # Calculate metrics
         elapsed_time = time.time() - start_time
-        tokens = outputs[0].shape[0] - inputs.input_ids.shape[-1]
+        tokens = outputs[0].shape[-1] - inputs.input_ids.shape[-1]
         tokens_per_sec = tokens / elapsed_time if elapsed_time > 0 else 0
         
         return response, elapsed_time, tokens_per_sec