reverted back to basic code

app.py

@@ -1,33 +1,39 @@
 from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import sqlite3
 import torch
 
 app = FastAPI()
 
-# Load Model & Tokenizer
+# Load the DeepSeek model and tokenizer
 MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to(device)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cpu")  # Use "cuda" if available
 
-def generate_sql_query(user_input):
-    """ Convert natural language input into an SQL query """
-    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True).to(device)
-    outputs = model.generate(**inputs, max_length=600, do_sample=False, num_beams=2)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
-@app.post("/chat")
-def chat(request: dict):
-    user_input = request.get("message", "")
-    if not user_input:
-        raise HTTPException(status_code=400, detail="Message cannot be empty")
+class ChatRequest(BaseModel):
+    message: str
+
+def generate_sql_query(user_input: str) -> str:
+    """
+    Generate an SQL query from a natural language query using the DeepSeek model.
+    """
+    inputs = tokenizer(user_input, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=600)
+    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return sql_query
 
+
+@app.post("/chat")
+def chat(request: ChatRequest):
+    user_input = request.message
+    
     sql_query = generate_sql_query(user_input)
     print(f"Generated SQL Query: {sql_query}")
+
     return {"response": sql_query}
 
 @app.get("/")
 def home():
-    return {"message": "DeepSeek SQL Query API is running"}
+    return {"message": "DeepSeek SQL Query API is running"}