Update app.py

app.py

@@ -1,10 +1,11 @@
 import torch
 import torch.nn as nn
-import torch.optim as optim
 import pandas as pd
-from torch.utils.data import Dataset, DataLoader
-from flask import Flask, request, jsonify
+from torch.utils.data import Dataset
 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
@@ -41,7 +42,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 inference but useful for training)
 class TextDataset(Dataset):
     def __init__(self, data, tokenizer, max_len=200):
         self.data = data
@@ -92,49 +93,32 @@ 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()
-    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
-
+    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())
 
+# FastAPI app
+app = FastAPI()
 
-# Flask App
-app = Flask(__name__)
+class Query(BaseModel):
+    query: str
 
-@app.route("/")
-def home():
+@app.get("/")
+async def root():
     return {"message": "Transformer-based Response Generator API is running!"}
 
-@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()
+@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}