frankmorales2020/Mistral-7B-v0dot3-text-to-sql-flash-attention-2

#FlashAttention only supports Ampere GPUs or newer. #NEED A100 IN GOOGLE COLAB

pip install -q https://github.com/lesj0610/flash-attention/releases/download/v2.8.3-cu12-torch2.10-cp312/flash_attn-2.8.3%2Bcu12torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

import os
import json
import torch
import random
import numpy as np
import warnings
from transformers import AutoModelForCausalLM, AutoTokenizer, logging
from peft import PeftModel, LoraConfig
import huggingface_hub

# --- Environment & Silence ---
os.environ["PEFT_DISABLE_TORCHAO"] = "1"
warnings.filterwarnings("ignore")
logging.set_verbosity_error()
huggingface_hub.logging.set_verbosity_error()

# --- Reproducibility ---
def set_reproducibility(seed=123):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    print(f"🔐 H2E Determinism Locked | Seed: {seed}")

set_reproducibility(123)

model_id = "mistralai/Mistral-7B-Instruct-v0.3"
peft_model_id = "frankmorales2020/Mistral-7B-v0dot3-text-to-sql-flash-attention-2"

# --- Tokenizer ---
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Add special tokens BEFORE resizing embeddings
special_tokens = {"additional_special_tokens": ["<EXTRA_0>", "<EXTRA_1>"]}
tokenizer.add_special_tokens(special_tokens)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# --- Base Model ---
print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2" 
)

# Crucial: Resize to match the tokenizer with the new special tokens
print(f"Resizing embeddings to {len(tokenizer)}...")
base_model.resize_token_embeddings(len(tokenizer))

# --- PEFT Adapter ---
print("Loading PEFT adapter...")
# We let PeftModel handle the config loading automatically from the hub/path.
# If you get a 'unexpected keyword argument' error, PEFT usually ignores extra keys anyway.
model = PeftModel.from_pretrained(
    base_model, 
    peft_model_id,
    is_trainable=False
)

model.eval()
print("✓ Model and Adapter loaded successfully!\n")

# --- Quick Test ---
def generate_sql(prompt_text):
    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

print("Ready for queries.")

def generate_sql(question, schema):
    # 1. Format prompt
    prompt = f"<EXTRA_0> {schema} <EXTRA_1> {question}"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # 2. Generate
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.1,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    # 3. Decode 
    generated_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    
    # 4. Extract only the SQL statement
    # We split by the separator and look for the first segment that contains 'SELECT'
    segments = generated_text.split("/******/")
    for segment in segments:
        segment = segment.strip()
        if "SELECT" in segment.upper():
            # Slice the string to start exactly at the SELECT command 
            # (ignoring any 'List the names...' text before it)
            start_index = segment.upper().find("SELECT")
            return segment[start_index:].strip()
            
    return "SQL not found."

# Usage
schema = "TABLE: Employees (id, name, department_id, salary)"
question = "Find the names of all employees who work in department 5."
print(f"Generated SQL: {generate_sql(question, schema)}")

@software{vonwerra2020trl,
  title   = {{TRL: Transformers Reinforcement Learning}},
  author  = {von Werra, Leandro and Belkada, Younes and Tunstall, Lewis and Beeching, Edward and Thrush, Tristan and Lambert, Nathan and Huang, Shengyi and Rasul, Kashif and Gallouédec, Quentin},
  license = {Apache-2.0},
  url     = {https://github.com/huggingface/trl},
  year    = {2020}
}