Update app.py

app.py

@@ -1,16 +1,20 @@
 import torch
 import torch.nn as nn
 import pandas as pd
+from torch.utils.data import Dataset
 from sklearn.model_selection import train_test_split
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Optional
-import uvicorn
+from fastapi.responses import JSONResponse
 import os
 
-# --- Tokenizer ---
+# Load data
+url = "https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY"
+df = pd.read_csv(url)
+
+# Tokenizer
 class ScratchTokenizer:
-    def __init__(self):
+    def _init_(self):
         self.word2idx = {"<PAD>": 0, "<SOS>": 1, "<EOS>": 2, "<UNK>": 3}
         self.idx2word = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}
         self.vocab_size = 4
@@ -31,18 +35,34 @@ class ScratchTokenizer:
     def decode(self, tokens):
         return " ".join([self.idx2word.get(idx, "<UNK>") for idx in tokens if idx > 0])
 
-# --- Load and Prepare Data ---
-url = "https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY"
-df = pd.read_csv(url)
-train_data, _ = train_test_split(df, test_size=0.2, random_state=42)
+# Train-Test Split
+train_data, test_data = train_test_split(df, test_size=0.2, random_state=42)
 
+# Initialize Tokenizer
 tokenizer = ScratchTokenizer()
 tokenizer.build_vocab(train_data["instruction"].tolist() + train_data["response"].tolist())
 
-# --- Model ---
+# Dataset Class (not used in inference but useful for training)
+class TextDataset(Dataset):
+    def _init_(self, data, tokenizer, max_len=200):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+
+    def _len_(self):
+        return len(self.data)
+
+    def _getitem_(self, idx):
+        src_text = self.data.iloc[idx]["instruction"]
+        tgt_text = self.data.iloc[idx]["response"]
+        src = torch.tensor(self.tokenizer.encode(src_text), dtype=torch.long)
+        tgt = torch.tensor(self.tokenizer.encode(tgt_text), dtype=torch.long)
+        return src, tgt
+
+# Model
 class GPTModel(nn.Module):
-    def __init__(self, vocab_size, embed_size=256, num_heads=8, num_layers=6, max_len=200):
-        super(GPTModel, self).__init__()
+    def _init_(self, vocab_size, embed_size=256, num_heads=8, num_layers=6, max_len=200):
+        super(GPTModel, self)._init_()
         self.embedding = nn.Embedding(vocab_size, embed_size)
         self.pos_embedding = nn.Parameter(torch.randn(1, max_len, embed_size))
         self.transformer = nn.TransformerDecoder(
@@ -58,7 +78,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 ---
+# Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = GPTModel(tokenizer.vocab_size).to(device)
 
@@ -72,34 +92,32 @@ def load_model(model, path="gpt_model.pth"):
 
 load_model(model)
 
-# --- Inference ---
-def generate_response(query, max_length=200):
+# Generate Response
+def generate_response(model, query, max_length=200):
     model.eval()
-    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_word = output.argmax(-1)[:, -1].unsqueeze(1)
-            tgt = torch.cat([tgt, next_word], dim=1)
-            if next_word.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 ---
+# FastAPI app
 app = FastAPI()
 
-class QueryRequest(BaseModel):
+class Query(BaseModel):
     query: str
 
 @app.get("/")
-def root():
+async def root():
     return {"message": "Transformer-based Response Generator API is running!"}
 
 @app.post("/query")
-def query_model(data: QueryRequest):
-    query = data.query.strip()
-    if not query:
-        return {"error": "Query cannot be empty"}
-    response = generate_response(query)
-    return {"query": query, "response": response}
+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}