GTimothee/sql-code-llama-4bits

Model Card for Model ID

4 bit general purpose text-to-SQL model.

Takes 5677MiB of GPU memory.

Model Details

Model Description

Provide the CREATE statement of the target table(s) in the context of your prompt and ask a question to your database. The model outputs a query to answer the question.

Data used for fine tuning: https://huggingface.co/datasets/b-mc2/sql-create-context

Uses

This model can be coupled with a chat model like llama2-chat to convert the output into a text answer.

Direct Use

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch


# load model
base_model = "GTimothee/sql-code-llama-4bits"
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    load_in_4bit=True,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.eval()

# load tokenizer
tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-hf")

eval_prompt = """You are a powerful text-to-SQL model. Your job is to answer questions about a database. You are given a question and context regarding one or more tables.

You must output the SQL query that answers the question.
### Input:
Which Class has a Frequency MHz larger than 91.5, and a City of license of hyannis, nebraska?

### Context:
CREATE TABLE table_name_12 (class VARCHAR, frequency_mhz VARCHAR, city_of_license VARCHAR)

### Response:
"""

model_input = tokenizer(eval_prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
    print(tokenizer.decode(model.generate(**model_input, max_new_tokens=100)[0], skip_special_tokens=True))

Outputs:

### Response:
SELECT class FROM table_name_12 WHERE frequency_mhz > 91.5 AND city_of_license = "hyannis, nebraska"

Bias, Risks, and Limitations

• potential security issues if there is a malicious use. If you execute blindly the SQL queries that are being generated by end users you could lose data, leak information etc.
• may be mistaken depending on the way the prompt has been written.

Recommendations

• Make sure that you check the generated SQL before applying it if the model is used by end users directly.
• The model works well when used on simple tables and simple queries. If possible, try to break a complex query into multiple simple queries.