Update README.md

README.md

@@ -58,10 +58,15 @@ def build_output(sql):
     return f"{sql.strip()}\n"
 
 create_table_statements = "YOUR TABLE SCHEMA HERE"
-inputs = build_prompt("YOUR PROMPT HERE", create_table_statements)
-input_ids = tokenizer(inputs, return_tensors="pt", padding = True, truncation = True, max_length = 512).input_ids.to(model.device)
-outputs = model.generate(input_ids, max_new_tokens=100)
-output = tokenizer.decode(outputs[0])
+
+def sqllamma(question):
+    input_ids = tokenizer(build_prompt(question, create_table_statements), return_tensors="pt", padding = True, truncation = True, max_length = 512).input_ids.to(model.device)
+    outputs = model.generate(input_ids, max_new_tokens=100)
+    output = tokenizer.decode(outputs[0])
+    sql = output.split("###")[3].split("[SQL]")[1].strip()    
+    return sql
+
+sqllama("YOUR QUESTION HERE")
 ```