Update app.py

app.py

@@ -1,55 +1,33 @@
-# from fastapi import FastAPI
-
-# app = FastAPI()
-
-# @app.get("/")
-# def greet_json():
-#     return {"Hello": "World!"}
-
-
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import sqlite3
 import torch
 
 app = FastAPI()
 
-# Load the DeepSeek model and tokenizer
+# Load Model & Tokenizer
 MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cpu")  # Use "cuda" if available
-
-
-class ChatRequest(BaseModel):
-    message: str
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-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=100)
-    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return sql_query
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to(device)
 
+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=100, 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: ChatRequest):
-    user_input = request.message
-    
+def chat(request: dict):
+    user_input = request.get("message", "")
+    if not user_input:
+        raise HTTPException(status_code=400, detail="Message cannot be empty")
+
     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"}
-
-
-# Run the API
-# if __name__ == "__main__":
-#     import uvicorn
-#     uvicorn.run(app, host="0.0.0.0", port=8000)
-