sushruthsam's picture
Update main.py
27f1194 verified
raw
history blame contribute delete
526 Bytes
from ctransformers import AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel
# Model loading with the new model name
llm = AutoModelForCausalLM.from_pretrained("sqlcoder-7b.Q4_K_S.gguf")
class Validation(BaseModel):
prompt: str # Assuming this includes both user_question and table_metadata_string
app = FastAPI()
@app.post("/generate_sql")
async def generate_sql(item: Validation):
# Format the actual prompt using item.prompt
completion = llm(item.prompt)
return completion