import os
os.environ['HF_HOME'] = '/tmp/hf_home'
os.environ['HF_DATASETS_CACHE'] = '/tmp/hf_datasets_cache'
os.environ['TRANSFORMERS_CACHE'] = '/tmp/cache'

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import uvicorn

MODEL_NAME = "16pramodh/t2s_model"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

app = FastAPI()

class QueryRequest(BaseModel):
    text: str

@app.get("/")
def read_root():
    return {"status": "running"}

@app.post("/predict")
def predict(request: QueryRequest):
    inputs = tokenizer(request.text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=256)
    return {"sql": tokenizer.decode(outputs[0], skip_special_tokens=True)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)