Hugging Face's logo

ucalyptus
/
prem-1B-SQL-grpo

Text Generation
Transformers
Safetensors
llama
Generated from Trainer
trl
grpo
conversational
text-generation-inference
Model card Files
xet
 Community
1
prem-1B-SQL-grpo / train_premsql.py
ucalyptus's picture
ucalyptus
Create train_premsql.py
f7a61b6 verified
Raw History Blame Contribute Delete
15 kB
#!/usr/bin/env python3
import os
# below line clears the console
os.system('clear')
# below line calls gpustat if the gpu id is 0
if os.environ.get('RANK') == '0':
os.system('gpustat')
import re
import torch
import logging
from typing import List, Dict, Any, Optional
from datasets import load_dataset, Dataset
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer
from accelerate import Accelerator
import torch.distributed as dist
from premsql.executors import SQLiteExecutor
from premsql.datasets import BirdDataset
from premsql.evaluator import Text2SQLEvaluator
import pandas as pd
import json
from pathlib import Path
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Constants
SYSTEM_PROMPT = """You are an expert SQL analyst. For the given question and database schema, generate a SQL query.
First analyze the schema and requirements, then write the query.
The database schema and tables are:
{schema}
Respond in the following format:
<reasoning>
[Step by step analysis of the requirements and schema]
</reasoning>
<sql>
[Your SQL query]
</sql>"""
class PremSQLv2Trainer:
def __init__(
self,
model_name: str,
output_dir: str,
dataset_folder: str = "./data",
hf_token: Optional[str] = None
):
self.model_name = model_name
self.output_dir = output_dir
self.dataset_folder = Path(dataset_folder)
self.hf_token = hf_token
self.device_count = torch.cuda.device_count()
self.executor = SQLiteExecutor()
self.evaluator = Text2SQLEvaluator(
executor=self.executor,
experiment_path=output_dir
)
# Create dataset directory if it doesn't exist
self.dataset_folder.mkdir(parents=True, exist_ok=True)
self.setup_environment()
def setup_environment(self):
"""Setup training environment and distributed training"""
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
if self.device_count > 1:
if os.environ.get('RANK') is not None and os.environ.get('WORLD_SIZE') is not None:
if not dist.is_initialized():
dist.init_process_group(backend='nccl')
else:
logger.info("Multiple GPUs detected but distributed training environment variables not set. Running on single GPU.")
@staticmethod
def extract_reasoning(text: str) -> str:
"""Extract reasoning from model output"""
try:
reasoning = text.split("<reasoning>")[-1]
reasoning = reasoning.split("</reasoning>")[0]
return reasoning.strip()
except IndexError:
return ""
@staticmethod
def extract_sql_query(text: str) -> str:
"""Extract SQL query from model output"""
try:
sql = text.split("<sql>")[-1]
sql = sql.split("</sql>")[0]
return sql.strip()
except IndexError:
return ""
def reasoning_quality_reward(self, text: str) -> float:
"""Evaluate the quality of reasoning"""
reasoning = self.extract_reasoning(text)
if not reasoning:
return 0.0
reward = 0.0
# Check for key reasoning components
if "schema" in reasoning.lower(): reward += 0.2
if "table" in reasoning.lower(): reward += 0.2
if "join" in reasoning.lower(): reward += 0.2
if "where" in reasoning.lower(): reward += 0.2
if "condition" in reasoning.lower(): reward += 0.2
# Check for step-by-step analysis
steps = [s for s in reasoning.split('\n') if s.strip()]
if len(steps) >= 3:
reward += 0.3
# Check for schema understanding
if any(word in reasoning.lower() for word in ["primary key", "foreign key", "relationship"]):
reward += 0.2
return min(reward, 1.5) # Cap at 1.5
def show_gpu_usage(self):
"""Show GPU usage if on main process"""
if os.environ.get('RANK') == '0':
os.system('gpustat')
def get_training_dataset(self) -> Dataset:
"""Load and prepare BirdBench dataset"""
if os.environ.get('RANK') == '0':
logger.info("Loading BirdBench dataset...")
# Initialize BirdDataset without forcing download
bird_dataset = BirdDataset(
split="train",
dataset_folder=self.dataset_folder,
hf_token=self.hf_token,
force_download=False
)
# Format the examples with schema information
formatted_examples = []
for example in bird_dataset.dataset:
# Get database schema
db_path = f"{self.dataset_folder}/train_databases/{example['db_id']}/{example['db_id']}.sqlite"
schema = self.get_database_schema(db_path)
# Format prompt with schema
prompt = [
{'role': 'system', 'content': SYSTEM_PROMPT.format(schema=schema)},
{'role': 'user', 'content': f"Database: {example['db_id']}\nQuestion: {example['question']}"}
]
formatted_examples.append({
'prompt': prompt,
'ground_truth': example['SQL'],
'db_id': example['db_id']
})
# Convert to Dataset
dataset = Dataset.from_list(formatted_examples)
# Add tokenized inputs
tokenizer = AutoTokenizer.from_pretrained(
self.model_name,
padding_side="right",
token=self.hf_token
)
def tokenize_function(examples):
# Apply chat template
texts = [tokenizer.apply_chat_template(p, tokenize=False) for p in examples['prompt']]
tokenized = tokenizer(texts, padding=True, truncation=True)
# Keep original fields
return {
'input_ids': tokenized['input_ids'],
'attention_mask': tokenized['attention_mask'],
'prompt': examples['prompt'],
'ground_truth': examples['ground_truth'],
'db_id': examples['db_id']
}
# Tokenize while keeping original fields
dataset = dataset.map(
tokenize_function,
batched=True,
remove_columns=None
)
if os.environ.get('RANK') == '0':
logger.info(f"Dataset loaded with {len(dataset)} examples")
self.show_gpu_usage()
return dataset
def get_database_schema(self, db_path: str) -> str:
"""Extract database schema from SQLite database"""
try:
import sqlite3
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
# Get all tables
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = cursor.fetchall()
schema_info = []
for table in tables:
table_name = table[0]
# Get columns for each table
cursor.execute(f"PRAGMA table_info({table_name});")
columns = cursor.fetchall()
# Format table schema
column_info = [f"{col[1]} ({col[2]})" for col in columns]
schema_info.append(f"Table: {table_name}\nColumns: {', '.join(column_info)}\n")
conn.close()
return "\n".join(schema_info)
except Exception as e:
if os.environ.get('RANK') == '0':
logger.warning(f"Failed to get schema for {db_path}: {str(e)}")
return "Schema information unavailable"
def execution_accuracy_reward(
self,
prompts: List,
completions: List,
**kwargs
) -> List[float]:
"""Reward function based on SQL execution accuracy"""
responses = [completion[0]['content'] for completion in completions]
rewards = []
# Get ground truth from dataset
dataset = kwargs.get('train_dataset')
batch_size = len(prompts)
current_idx = kwargs.get('current_step', 0) * batch_size
for i, response in enumerate(responses):
query = self.extract_sql_query(response)
reward = 0.0
try:
if query:
# Get corresponding ground truth
example = dataset[current_idx + i]
db_path = f"{self.dataset_folder}/train_databases/{example['db_id']}/{example['db_id']}.sqlite"
# Debug information
if os.environ.get('RANK') == '0':
debug_info = {
'query': query,
'ground_truth': example['ground_truth'],
'db_path': db_path
}
print("\n=== Debug Info ===")
print(f"Full Response:\n{response}")
print(f"\nExtracted Query:\n{query}")
print(f"\nDB Path: {db_path}")
print(f"DB Path exists: {os.path.exists(db_path)}")
print(f"\nExample Data:")
print(json.dumps(example, indent=2))
print("\nEvaluator Input:")
print(json.dumps(debug_info, indent=2))
print("================\n")
result = self.evaluator.execute(
metric_name="accuracy",
model_responses=[{
'query': query,
'ground_truth': example['ground_truth'],
'db_path': db_path
}],
filter_by="db_id",
meta_time_out=10
)
# Debug evaluator result
if os.environ.get('RANK') == '0':
print("Evaluator Result:")
print(json.dumps(result, indent=2))
# Reward based on accuracy
reward = result['accuracy'] * 2.0
# Add small rewards for SQL components
reward += 0.1 * query.upper().count("SELECT")
reward += 0.1 * query.upper().count("FROM")
reward += 0.2 * query.upper().count("JOIN")
except Exception as e:
if os.environ.get('RANK') == '0':
logger.warning(f"Query execution error: {str(e)}")
logger.warning(f"Error occurred with query: {query}")
logger.warning(f"Full response was: {response}")
if all(k in query.upper() for k in ["SELECT", "FROM"]):
reward = 0.1
rewards.append(reward)
return rewards
def setup_training_args(self) -> GRPOConfig:
"""Configure training arguments"""
return GRPOConfig(
output_dir=self.output_dir,
run_name=f"{self.model_name.split('/')[-1]}-GRPO-v2",
learning_rate=5e-6,
adam_beta1=0.9,
adam_beta2=0.99,
weight_decay=0.1,
warmup_ratio=0.1,
lr_scheduler_type='cosine',
logging_steps=1,
bf16=True,
per_device_train_batch_size=2,
gradient_accumulation_steps=2,
num_generations=16,
max_prompt_length=512,
max_completion_length=256,
num_train_epochs=1,
save_steps=50,
max_grad_norm=0.1,
log_on_each_node=True,
use_vllm=False,
report_to="wandb", # Using wandb for tracking
ddp_find_unused_parameters=False,
push_to_hub=False,
hub_model_id=None,
hub_token=None,
)
def train(self):
"""Main training function"""
try:
# Load dataset
dataset = self.get_training_dataset()
if os.environ.get('RANK') == '0':
logger.info(f"Loaded dataset with {len(dataset)} examples")
# Initialize model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.bfloat16,
).cuda()
tokenizer = AutoTokenizer.from_pretrained(self.model_name)
tokenizer.pad_token = tokenizer.eos_token
# Setup training arguments
training_args = self.setup_training_args()
class GPUMonitorCallback(transformers.TrainerCallback):
def __init__(self, trainer):
self.trainer = trainer
def on_save(self, args, state, control, **kwargs):
if os.environ.get('RANK') == '0' and state.global_step % 50 == 0:
self.trainer.show_gpu_usage()
# Initialize trainer
trainer = GRPOTrainer(
model=model,
processing_class=tokenizer,
reward_funcs=[self.execution_accuracy_reward],
args=training_args,
train_dataset=dataset,
callbacks=[GPUMonitorCallback(self)]
)
# Start training
trainer.train()
# Save model
if dist.get_rank() == 0:
logger.info(f"Saving model to {self.output_dir}")
trainer.save_model()
self.show_gpu_usage() # Show final GPU usage
# Cleanup distributed training
if dist.is_initialized():
dist.destroy_process_group()
except Exception as e:
if os.environ.get('RANK') == '0':
logger.error(f"Training failed with error: {str(e)}")
if dist.is_initialized():
dist.destroy_process_group()
raise
def main():
"""Main entry point"""
model_name = "premai-io/prem-1B-SQL"
output_dir = "outputs/Prem-1B-SQL-GRPO-v2"
dataset_folder = "./data/birdbench"
hf_token = "hf_"
trainer = PremSQLv2Trainer(
model_name=model_name,
output_dir=output_dir,
dataset_folder=dataset_folder,
hf_token=hf_token
)
trainer.train()
if __name__ == "__main__":
main()