.

app.py

@@ -37,11 +37,18 @@ def predict_product(num1, num2):
     eos_token_id = tokenizer.eos_token_id
     past_key_values = None
     for _ in range(100):  # Set a maximum limit to prevent infinite loops
-        outputs = model.generate(generated_ids, max_new_tokens=1, do_sample=False, past_key_values=past_key_values, return_dict_in_generate=True, use_cache=True)
-        generated_ids = torch.cat((generated_ids, outputs.sequences[:, -1:]), dim=-1)
+        outputs = model(
+            input_ids=generated_ids,
+            past_key_values=past_key_values,
+            use_cache=True
+        )
+        logits = outputs.logits
         past_key_values = outputs.past_key_values
         
-        if outputs.sequences[0, -1].item() == eos_token_id:
+        next_token_id = torch.argmax(logits[:, -1, :], dim=-1)
+        generated_ids = torch.cat((generated_ids, next_token_id.unsqueeze(-1)), dim=-1)
+        
+        if next_token_id.item() == eos_token_id:
             break
         
         output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)