Upload app.py

app.py

@@ -2,26 +2,36 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
 
-# Create a new FastAPI app instance
 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": "Hello World"}
+    return {"message": "SQL Generation Server is running"}
 
-# Define a function to handle the POST request at '/generate'
 @app.post("/generate")
-def generate(request: dict):
+def generate(request: QueryRequest):
     try:
-        text = request.get('text')
-        if not text:
-            raise HTTPException(status_code=400, detail="Text field is required")
+        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()
         
-        output = pipe(text, max_new_tokens=50)
-        # Return the generated text in JSON response
-        return {"output": output[0]['generated_text']}
+        # 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))
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)