Model Card for llama3-sql2plan

This model is a fine-tuned version of meta-llama/Llama-3.2-3B using LoRA (Low-Rank Adaptation) for the task of generating PostgreSQL execution plans from SQL queries. The model takes SQL queries as input and outputs the corresponding PostgreSQL execution plan in JSON format.

Model Details

Model Description

This model is specifically designed to convert SQL queries into PostgreSQL execution plans. It was fine-tuned using Parameter-Efficient Fine-Tuning (PEFT) with LoRA adapters, allowing efficient training while maintaining the base model's capabilities.

• Developed by: Anirudh Bharadwaj
• Model type: Causal Language Model (Decoder-only)
• Language(s) (NLP): English (SQL and JSON)
• License: Apache 2.0
• Finetuned from model: meta-llama/Llama-3.2-3B

Model Sources

• Repository: abharadwaj123/llama3-sql2plan
• Base Model: meta-llama/Llama-3.2-3B

Uses

Direct Use

This model can be used directly to generate PostgreSQL execution plans from SQL queries. It is intended for:

• Database query optimization analysis
• Understanding query execution strategies
• Educational purposes for learning PostgreSQL query planning
• Database performance analysis tools

Example Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "abharadwaj123/llama3-sql2plan"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

sql_query = "SELECT * FROM users WHERE age > 25;"
prompt = (
    "Generate the PostgreSQL execution plan in JSON format for the SQL query.\n\n"
    "[QUERY]\n" + sql_query + "\n\n[PLAN]\n"
)

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.7,
        do_sample=True,
    )

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
plan = generated_text.split("[PLAN]\n")[-1].strip()
print(plan)

Out-of-Scope Use

This model should not be used for:

• Generating actual executable SQL queries (it generates execution plans, not queries)
• Real-time database query execution
• Production database systems without proper validation
• Any use case requiring guaranteed accuracy of execution plans

Bias, Risks, and Limitations

Limitations

• Accuracy: The model generates execution plans based on training data patterns and may not always produce accurate or optimal plans for all SQL queries.
• PostgreSQL-specific: The model is trained specifically for PostgreSQL execution plans and may not be suitable for other database systems.
• Training Data Scope: The model was trained on a subset of Stack Overflow data (10,000 samples from ~16,332 available), which may not cover all SQL query patterns.
• No Database Context: The model does not have access to actual database schema, indexes, or statistics, which are crucial for accurate execution plan generation.

Recommendations

Users should:

• Validate generated execution plans against actual PostgreSQL EXPLAIN output
• Not rely solely on this model for critical database optimization decisions
• Use this model as a tool for understanding and learning, not as a replacement for actual database query planning
• Be aware that execution plans may vary based on database configuration, schema, and data distribution

How to Get Started with the Model

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer
model_name = "abharadwaj123/llama3-sql2plan"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Prepare input
sql_query = "SELECT * FROM users WHERE age > 25;"
prompt = (
    "Generate the PostgreSQL execution plan in JSON format for the SQL query.\n\n"
    "[QUERY]\n" + sql_query + "\n\n[PLAN]\n"
)

# Generate plan
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    max_new_tokens=256,
    temperature=0.7,
    do_sample=True,
)

# Extract plan
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
plan = generated_text.split("[PLAN]\n")[-1].strip()

Training Details

Training Data

The model was trained on a dataset derived from Stack Overflow data (stackoverflow_n18147.csv), containing SQL queries and their corresponding PostgreSQL execution plans in JSON format.

• Total samples in dataset: 18,147
• Training samples used: 10,000 (sampled from first 90% of dataset, ~16,332 samples)
• Sampling method: Random sampling with random_state=42
• Data format: SQL query text paired with PostgreSQL execution plan JSON

The training data was filtered to remove rows with missing sql_text or plan_json values.

Training Procedure

Preprocessing

1. Data Loading: Loaded CSV file and filtered out rows with missing SQL text or plan JSON
2. Data Splitting: Used first 90% of dataset as training pool, then randomly sampled 10,000 examples
3. Formatting: Each example was formatted with a prompt template:

   Generate the PostgreSQL execution plan in JSON format for the SQL query.
   
   [QUERY]
   {sql_text}
   
   [PLAN]
   {plan_json}
   

4. Tokenization:
   • Maximum sequence length: 2048 tokens
   • Input prompt tokens were masked in labels (set to -100) to only train on plan generation
   • Padding to max_length

Training Hyperparameters

• Training regime: FP16 mixed precision training
• LoRA Configuration:
  • r: 64
  • lora_alpha: 32
  • lora_dropout: 0.05
  • target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]
• Training Arguments:
  • per_device_train_batch_size: 2
  • gradient_accumulation_steps: 8
  • effective_batch_size: 16
  • learning_rate: 1e-5
  • warmup_ratio: 0.03
  • num_train_epochs: 2
  • gradient_checkpointing: True
  • logging_steps: 20
  • save_strategy: "epoch"

Testing Data

The remaining 10% of the original dataset (~1,815 samples) was held out and could be used for evaluation.

Model Examination

The model uses LoRA (Low-Rank Adaptation) fine-tuning, which allows efficient training by only updating a small number of parameters (low-rank matrices) while keeping the base model weights frozen. This approach:

• Reduces memory requirements during training
• Enables faster training compared to full fine-tuning
• Maintains the base model's general capabilities
• Allows easy merging of adapters with the base model

Technical Specifications

Model Architecture and Objective

• Architecture: Transformer-based decoder-only language model (Llama-3.2-3B)
• Objective: Causal language modeling with masked input tokens
• Fine-tuning Method: LoRA (Low-Rank Adaptation) via PEFT
• Base Model Parameters: ~3 billion
• Trainable Parameters: Significantly reduced via LoRA (exact count depends on LoRA rank and target modules)

Model Card Authors

Anirudh Bharadwaj

Model Card Contact

For questions or issues, please contact through the Hugging Face model repository: abharadwaj123/llama3-sql2plan