Update app.py

app.py

@@ -43,18 +43,23 @@ Description: """
     # Generate response
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,
+            input_ids=inputs['input_ids'],  # Explicitly use input_ids
+            attention_mask=inputs['attention_mask'],  # Add attention mask
             max_new_tokens=max_length,
             temperature=temperature,
             top_p=top_p,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            bos_token_id=tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            use_cache=True,  # Enable caching
+            return_dict_in_generate=True,  # Return as dict
+            output_scores=True  # Get scores
         )
     
     # Decode and return the response
-    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return full_text.split("Description: ")[-1].strip()
+    generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+    return generated_text.split("Description: ")[-1].strip()
 
 def create_demo(model_id):
     # Load model and tokenizer