Update README.md

README.md

@@ -16,6 +16,14 @@ base_model:
 
 ## Model Details
 
+To load the model:
+```
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("kristiannordby/llama3-sqlcoder-ft")
+model = AutoModelForCausalLM.from_pretrained("kristiannordby/llama3-sqlcoder-ft")
+```
 Prompt:
 ```
 ### Task
@@ -30,6 +38,32 @@ Given the database schema, here is the SQL query that [QUESTION]{user_question}[
 [SQL]
 ```
 
+To prompt the model for generation:
+```
+def build_prompt(user_question, create_table_statements):
+    return (
+        "### Task\n"
+        f"Generate a SQL query to answer [QUESTION]{user_question}[/QUESTION]\n\n"
+        "### Database Schema\n"
+        "The query will run on a database with the following schema:\n"
+        f"{create_table_statements}\n\n"
+        "### Answer\n"
+        f"Given the database schema, here is the SQL query that [QUESTION]{user_question}[/QUESTION]\n"
+        "[SQL]\n"
+    )
+
+def build_output(sql):
+    # Add a newline at end; if the data has a closing "[/SQL]", add it here!
+    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])
+```
+
+
 
 ### Model Description