Update app.py

app.py

@@ -95,14 +95,17 @@ load_model(model)
 # Generate Response
 def generate_response(model, query, max_length=200):
     model.eval()
-    src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
-    tgt = torch.tensor([[1]]).to(device)  # <SOS>
-    for _ in range(max_length):
-        output = model(src, tgt)
-        next_word = output.argmax(-1)[:, -1].unsqueeze(1)
-        tgt = torch.cat([tgt, next_word], dim=1)
-        if next_word.item() == 2:  # <EOS>
-            break
+    with torch.no_grad():  # Disable gradient tracking
+        src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
+        tgt = torch.tensor([[1]]).to(device)  # <SOS>
+
+        for _ in range(max_length):
+            output = model(src, tgt)
+            next_token = output[:, -1, :].argmax(dim=-1, keepdim=True)
+            tgt = torch.cat([tgt, next_token], dim=1)
+            if next_token.item() == 2:  # <EOS>
+                break
+
     return tokenizer.decode(tgt.squeeze(0).tolist())
 
 # FastAPI app