Update main.py

main.py

@@ -1,9 +1,10 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from ctransformers import AutoModelForCausalLM  # Ensure correct import for your model
+from transformers 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):
@@ -15,6 +16,13 @@ app = FastAPI()
 # Endpoint for SQL query generation
 @app.post("/generate_sql")
 async def generate_sql(item: Validation):
-    # Use the provided prompt directly
-    completion = llm(item.prompt)
-    return completion
+    # 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}