app.py

from fastapi import FastAPI, HTTPException 
from pydantic import BaseModel 
from transformers import pipeline 
 
app = FastAPI() 
 
# Initialize the text generation pipeline 
pipe = pipeline("text-generation", model="defog/sqlcoder-7b-2", pad_token_id=2) 
 
class QueryRequest(BaseModel): 
 text: str 
 
@app.get("/") 
def home(): 
 return {"message": "SQL Generation Server is running"} 
 
 
@app.post("/generate") 
def generate(request: QueryRequest): 
 try: 
 text = request.text 
 prompt = f"Generate a valid SQL query for the following request. Only return the SQL query, nothing else:\n\n{text}\n\nSQL query:" 
 output = pipe(prompt, max_new_tokens=100) 
 
 generated_text = output[0]['generated_text'] 
 sql_query = generated_text.split("SQL query:")[-1].strip() 
 
 # Basic validation 
 if not sql_query.lower().startswith(('select', 'show', 'describe')): 
 raise ValueError("Generated text is not a valid SQL query") 
 
 return {"output": sql_query} 
 except Exception as e: 
 raise HTTPException(status_code=500, detail=str(e)) 
 
if __name__ == "__main__": 
 import uvicorn 
 uvicorn.run(app, host="0.0.0.0", port=7860)