| 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) |
