Update app.py

app.py

@@ -11,17 +11,34 @@ trained_model.to(device)
 untrained_model.to(device)
 
 def generate(commentary_text, max_length, temperature):
-    if temperature <= 0:
-        return "Error: Temperature must be a strictly positive float.", "Error: Temperature must be a strictly positive float."
-    
-    # Generate text using the finetuned model
-    input_ids = trained_tokenizer(commentary_text, return_tensors="pt").input_ids.to(device)
-    trained_output = trained_model.generate(input_ids, max_length=max_length, num_beams=5, do_sample=True, temperature=temperature)
+# Generate text using the finetuned model
+    inputs = trained_tokenizer(commentary_text, return_tensors="pt", padding=True)
+    input_ids = inputs.input_ids.to(device)
+    attention_mask = inputs.attention_mask.to(device)
+    trained_output = trained_model.generate(
+        input_ids, 
+        max_length=max_length, 
+        num_beams=5, 
+        do_sample=True, 
+        temperature=temperature, 
+        attention_mask=attention_mask,
+        pad_token_id=trained_tokenizer.eos_token_id
+    )
     trained_text = trained_tokenizer.decode(trained_output[0], skip_special_tokens=True)
     
     # Generate text using the base model
-    input_ids = untrained_tokenizer(commentary_text, return_tensors="pt").input_ids.to(device)
-    untrained_output = untrained_model.generate(input_ids, max_length=max_length, num_beams=5, do_sample=True,temperature=temperature)
+    inputs = untrained_tokenizer(commentary_text, return_tensors="pt", padding=True)
+    input_ids = inputs.input_ids.to(device)
+    attention_mask = inputs.attention_mask.to(device)
+    untrained_output = untrained_model.generate(
+        input_ids, 
+        max_length=max_length, 
+        num_beams=5, 
+        do_sample=True, 
+        temperature=temperature, 
+        attention_mask=attention_mask,
+        pad_token_id=untrained_tokenizer.eos_token_id
+    )
     untrained_text = untrained_tokenizer.decode(untrained_output[0], skip_special_tokens=True)
     
     return trained_text, untrained_text