Update app.py

app.py

@@ -1,11 +1,10 @@
 import torch
 import torch.nn as nn
+import torch.optim as optim
 import pandas as pd
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, DataLoader
+from flask import Flask, request, jsonify
 from sklearn.model_selection import train_test_split
-from fastapi import FastAPI
-from pydantic import BaseModel
-from fastapi.responses import JSONResponse
 import os
 
 # Load data
@@ -42,7 +41,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 (not used in inference but useful for training)
+# Dataset Class
 class TextDataset(Dataset):
     def __init__(self, data, tokenizer, max_len=200):
         self.data = data
@@ -93,31 +92,49 @@ 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):
     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
-app = FastAPI()
 
-class Query(BaseModel):
-    query: str
+# Flask App
+app = Flask(__name__)
 
-@app.get("/")
-async def root():
+@app.route("/")
+def home():
     return {"message": "Transformer-based Response Generator API is running!"}
 
-@app.post("/query")
-async def query_model(query: Query):
-    if not query.query.strip():
-        return JSONResponse(status_code=400, content={"error": "Query cannot be empty"})
-    response = generate_response(model, query.query)
-    return {"query": query.query, "response": response}
+@app.route("/query", methods=["POST"])
+def query_model():
+    data = request.get_json()
+    query = data.get("query", "")
+    if not query:
+        return jsonify({"error": "Query cannot be empty"}), 400
+    response = generate_response(model, query)
+    return jsonify({"query": query, "response": response})
+
+# DO NOT ADD app.run()
+