from fastapi import FastAPI | |
from pydantic import BaseModel | |
from ctransformers import AutoModelForCausalLM # Ensure correct import for your model | |
# Model loading | |
llm = AutoModelForCausalLM.from_pretrained("sqlcoder-34b-alpha.Q4_K_S.gguf", | |
max_tokens = 1024, threads = 3) | |
# Pydantic object for request validation | |
class Validation(BaseModel): | |
prompt: str | |
# Initialize FastAPI app | |
app = FastAPI() | |
# Endpoint for SQL query generation | |
async def generate_sql(item: Validation): | |
# Use the provided prompt directly | |
completion = llm(item.prompt) | |
return completion | |