Update README.md

README.md

@@ -53,30 +53,66 @@ Always validate and test generated queries before execution in a production data
 ## How to Get Started with the Model
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
+from unsloth import FastLanguageModel
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name="mervp/SQLGenie",
+    max_seq_length=2048,
+    dtype=None,
+)
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    "unsloth/llama-3.2-3b-unsloth-bnb-4bit",
-    device_map="auto",
-    trust_remote_code=True
+prompt = """ You are an text to SQL query translator.
+             Users will ask you questions in English
+             and you will generate a SQL query based on their question
+             SQL has to be simple, The schema context has been provided to you.
+
+
+### User Question:
+{}
+
+### Sql Context:
+{}
+
+### Sql Query:
+{}
+"""
+
+question = "List the names of customers who have an account balance greater than 6000."
+schema = """
+CREATE TABLE socially_responsible_lending (
+    customer_id INT,
+    name VARCHAR(50),
+    account_balance DECIMAL(10, 2)
+);
+
+INSERT INTO socially_responsible_lending VALUES
+    (1, 'james Chad', 5000),
+    (2, 'Jane Rajesh', 7000),
+    (3, 'Alia Kapoor', 6000),
+    (4, 'Fatima Patil', 8000);
+"""
+
+inputs = tokenizer(
+    [prompt.format(question, schema, "")],
+    return_tensors="pt",
+    padding=True,
+    truncation=True
+).to("cuda")
+
+output = model.generate(
+    **inputs,
+    max_new_tokens=256, 
+    temperature=0.2,  
+    top_p=0.9,          
+    top_k=50,        
+    do_sample=True     
 )
 
-tokenizer = AutoTokenizer.from_pretrained("unsloth/llama-3.2-3b-unsloth-bnb-4bit")
-model = PeftModel.from_pretrained(base_model, "mervp/SQLGenie")
+decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
 
-prompt = "List the customers from Canada."
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-outputs = model.generate(**inputs, max_new_tokens=100)
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+if "### Sql Query:" in decoded_output:
+    sql_query = decoded_output.split("### Sql Query:")[-1].strip()
+else:
+    sql_query = decoded_output.strip()
 
-#
-OR
-#
-from unsloth import FastLanguageModel
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="mervp/SQLGenie",
-    max_seq_length=2048,
-    dtype=None,                         
-    # load_in_4bit=True,
-)
+print(sql_query)