from fastapi import FastAPI
from pydantic import BaseModel
from ctransformers import AutoModelForCausalLM, AutoTokenizer

# Model loading
llm = AutoModelForCausalLM.from_pretrained("sqlcoder-7b.Q4_K_S.gguf")
tokenizer = AutoTokenizer.from_pretrained("sqlcoder-7b.Q4_K_S.gguf")

# Pydantic object for request validation
class Validation(BaseModel):
    prompt: str

# Initialize FastAPI app
app = FastAPI()

# Endpoint for SQL query generation
@app.post("/generate_sql")
async def generate_sql(item: Validation):
    # Tokenize the input prompt
    input_ids = tokenizer.encode(item.prompt, return_tensors="pt")

    # Use the tokenized prompt for model completion
    completion = llm.generate(input_ids)

    # Decode the generated SQL query
    generated_sql = tokenizer.decode(completion[0], skip_special_tokens=True)

    return {"generated_sql": generated_sql}