File size: 866 Bytes
5e7d2c0 271505c 5e7d2c0 fa810c7 ed6d964 5e7d2c0 ed6d964 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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}
|