AhmedSSoliman's picture
Update README.md
a7b762f
---
library_name: peft
base_model: mistralai/Mistral-7B-Instruct-v0.1
pipeline_tag: text-generation
datasets:
- bugdaryan/sql-create-context-instruction
tags:
- Mistral
- PEFT
- LoRA
- SQL
---
### Model Description
<!-- Provide a longer summary of what this model is. -->
SQL Generation model which is fine-tuned on the Mistral-7B-Instruct-v0.1.
Inspired from https://huggingface.co/kanxxyc/Mistral-7B-SQLTuned
### Code
```py
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
peft_model_id = "AhmedSSoliman/Mistral-Instruct-SQL-Generation"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True, return_dict=True, load_in_4bit=True, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)
def predict_SQL(table, question):
pipe = pipeline('text-generation', model = base_model, tokenizer = tokenizer)
prompt = f"[INST] Write SQL query to answer the following question given the database schema. Please wrap your code answer using ```: Schema: {table} Question: {question} [/INST] Here is the SQL query to answer to the question: {question}: ``` "
#prompt = f"### Schema: {table} ### Question: {question} # "
ans = pipe(prompt, max_new_tokens=200)
generatedSql = ans[0]['generated_text'].split('```')[2]
return generatedSql
table = "CREATE TABLE Employee (name VARCHAR, salary INTEGER);"
question = 'Show names for all employees with salary more than the average.'
generatedSql=predict_SQL(table, question)
print(generatedSql)
```