Upload app.py

app.py

@@ -0,0 +1,96 @@
+import os
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers'
+
+from fastapi import FastAPI, Header, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import Optional, Dict, Annotated
+import logging
+import torch
+import os
+
+# Initialize logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load model
+MODEL_NAME = "defog/sqlcoder-7b-2"
+logger.info(f"Loading model: {MODEL_NAME}")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+
+# FastAPI init
+app = FastAPI(title="Text to SQL API")
+
+# CORS for Hugging Face Space
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Request Model
+class RequestModel(BaseModel):
+    entity_urn: str
+    prompt: str
+
+# Response Model
+class ResponseModel(BaseModel):
+    message: str
+    result: str
+    action_type: str
+    entity_urn: str
+    metadata: Optional[Dict] = None
+
+@app.get("/")
+async def root():
+    return {
+        "message": "Text-to-SQL API running",
+        "docs": "/docs",
+        "health": "/health"
+    }
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
+@app.post("/generate", response_model=ResponseModel)
+async def generate_sql(
+    request: RequestModel,
+    x_api_key: Annotated[str, Header()]  # Optional token check
+):
+    try:
+        if not request.prompt.strip():
+            return ResponseModel(
+                message="failure",
+                result="Empty prompt",
+                action_type="text_to_sql",
+                entity_urn=request.entity_urn
+            )
+
+        inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_length=512)
+        sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        return ResponseModel(
+            message="success",
+            result=sql.strip(),
+            action_type="text_to_sql",
+            entity_urn=request.entity_urn,
+            metadata={"tokens": len(inputs["input_ids"][0])}
+        )
+
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        return ResponseModel(
+            message="failure",
+            result=f"Error: {str(e)}",
+            action_type="text_to_sql",
+            entity_urn=request.entity_urn
+        )