Update app.py

app.py

@@ -91,17 +91,32 @@ def load_model(model, path="gpt_model.pth"):
 
 load_model(model)
 
-# Generate Response
+# # 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
+#     return tokenizer.decode(tgt.squeeze(0).tolist())
+
 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())
 
 # Flask App