Update response_1.py

response_1.py

@@ -1,13 +1,16 @@
 import torch
 import torch.nn as nn
-import torch.optim as optim
 import pandas as pd
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
 from flask import Flask, request, jsonify
 from sklearn.model_selection import train_test_split
 import os
 import time
 
+# Enable cudnn benchmark for better GPU performance
+if torch.cuda.is_available():
+    torch.backends.cudnn.benchmark = True
+
 # Load data
 url = "https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY"
 df = pd.read_csv(url)
@@ -42,7 +45,7 @@ train_data, test_data = train_test_split(df, test_size=0.2, random_state=42)
 tokenizer = ScratchTokenizer()
 tokenizer.build_vocab(train_data["instruction"].tolist() + train_data["response"].tolist())
 
-# Dataset Class
+# Dataset Class (not used in this file but kept for completeness)
 class TextDataset(Dataset):
     def __init__(self, data, tokenizer, max_len=200):
         self.data = data
@@ -78,7 +81,7 @@ class GPTModel(nn.Module):
         output = self.transformer(tgt_emb.permute(1, 0, 2), src_emb.permute(1, 0, 2), tgt_mask=tgt_mask)
         return self.fc_out(output.permute(1, 0, 2))
 
-# Load model
+# Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = GPTModel(tokenizer.vocab_size).to(device)
 
@@ -92,35 +95,33 @@ def load_model(model, path="gpt_model.pth"):
 
 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
-#     return tokenizer.decode(tgt.squeeze(0).tolist())
-
-def generate_response(model, query, max_length=200):
+# Generate Response with mixed precision if CUDA is available
+def generate_response(model, query, max_length=100):
     model.eval()
-    with torch.no_grad():  # Disable gradient tracking
+    with torch.no_grad():
         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
+        if device.type == "cuda":
+            scaler = torch.cuda.amp.autocast()
+            with torch.cuda.amp.autocast():
+                for _ in range(max_length):
+                    output = model(src, tgt)
+                    logits = output[:, -1, :]
+                    next_token = torch.argmax(logits, dim=-1, keepdim=True)
+                    tgt = torch.cat([tgt, next_token], dim=1)
+                    if next_token.item() == 2:
+                        break
+        else:
+            for _ in range(max_length):
+                output = model(src, tgt)
+                logits = output[:, -1, :]
+                next_token = torch.argmax(logits, dim=-1, keepdim=True)
+                tgt = torch.cat([tgt, next_token], dim=1)
+                if next_token.item() == 2:
+                    break
 
     return tokenizer.decode(tgt.squeeze(0).tolist())
 
-
 # Flask App
 app = Flask(__name__)
 
@@ -130,7 +131,7 @@ def home():
 
 @app.route("/intent")
 def intents():
-    return jsonify({"intents" :list(set(df['intent'].dropna()))})
+    return jsonify({"intents": list(set(df['intent'].dropna()))})
 
 @app.route("/query", methods=["POST"])
 def query_model():
@@ -141,7 +142,7 @@ def query_model():
     start = time.time()
     response = generate_response(model, query)
     end = time.time()
-    return jsonify({"query": query, "response": response,"response_time":(end-start)})
+    return jsonify({"query": query, "response": response, "response_time": end - start})
 
 if __name__ == "__main__":
     load_model(model)