Update main.py

main.py

@@ -1,28 +1,31 @@
+from ctransformers import AutoModelForCausalLM
 from fastapi import FastAPI
 from pydantic import BaseModel
-from ctransformers import AutoModelForCausalLM, AutoTokenizer
 
-# Model loading
+# Model loading with the new model name
 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
+    prompt: str  # Assuming this includes both user_question and table_metadata_string
 
-# 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")
+    # Updated system prompt
+    system_prompt = """### Task
+Generate a SQL query to answer the following question:
+`{question}`
 
-    # Use the tokenized prompt for model completion
-    completion = llm.generate(input_ids)
+### Database Schema
+The query will run on a database with the following schema:
+{schema}
 
-    # Decode the generated SQL query
-    generated_sql = tokenizer.decode(completion[0], skip_special_tokens=True)
-
-    return {"generated_sql": generated_sql}
+### Answer
+Given the database schema, here is the SQL query that answers `{question}`:
+```sql
+"""
+    # Format the actual prompt using item.prompt
+    prompt = system_prompt.format(user_question="Your question here", table_metadata_string="Your schema here")
+    completion = llm(prompt)
+    return completion