.

app.py

@@ -35,11 +35,13 @@ def predict_product(num1, num2):
         valid_input = False
 
     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)
-        generated_ids = torch.cat((generated_ids, outputs[:, -1:]), dim=-1)
+        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)
+        past_key_values = outputs.past_key_values
         
-        if outputs[0, -1].item() == eos_token_id:
+        if outputs.sequences[0, -1].item() == eos_token_id:
             break
         
         output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)