207 lines
7.7 KiB
Python
207 lines
7.7 KiB
Python
from unsloth import FastLanguageModel, PatchFastRL
|
|
PatchFastRL("GRPO", FastLanguageModel)
|
|
from unsloth import is_bfloat16_supported
|
|
import torch
|
|
#---v4
|
|
#pip install diffusers
|
|
#!pip install "unsloth==2025.2.4" vllm
|
|
#!pip install --upgrade pillow
|
|
|
|
num_generations = 8
|
|
per_device_train_batch_size = 8
|
|
gradient_accumulation_steps = 4
|
|
num_train_epochs = 3
|
|
beta = 0.01
|
|
learning_rate = 1e-4
|
|
max_seq_length = 1650
|
|
lora_rank = 8
|
|
lora_alpha = lora_rank * 2
|
|
dataset_dir = "cot-qa-distill.csv"
|
|
model_name = "./models/deepquery4-sft-1"
|
|
output_dir = "./result/deepquery4-grpo-1"
|
|
SYS="""You are DeepQuery, a data science expert. Below, you are presented with a database schema, a question and a hint.Your task is to read the schema with annotations of the columns, understand the question and the hint, and generate a valid SQL query to answer the question. You should reason step by step, and includes your reasonings between <think> and </think>."""
|
|
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name = model_name,
|
|
max_seq_length = max_seq_length,
|
|
load_in_4bit = True, # False for LoRA 16bit
|
|
fast_inference = True, # Enable vLLM fast inference
|
|
max_lora_rank = lora_rank,
|
|
gpu_memory_utilization = 0.4, # Reduce if out of memory
|
|
)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
|
target_modules = [
|
|
"gate_proj","up_proj", "down_proj","q_proj", "k_proj","o_proj", "v_proj"
|
|
#"q_proj", "k_proj", "up_proj", "down_proj",
|
|
], # Remove QKVO if out of memory
|
|
lora_alpha = lora_alpha,
|
|
use_gradient_checkpointing = "unsloth", # Enable long context finetuning
|
|
random_state = 3407,
|
|
)
|
|
def extract_answer(text: str) -> str:
|
|
|
|
start_tag = 'My final answer is: \n```sql\n'
|
|
end_tag = '\n```'
|
|
start_tag_index = text.find(start_tag)
|
|
if start_tag_index != -1:
|
|
start_index = start_tag_index + len(start_tag)
|
|
end_index = text.find(end_tag, start_index)
|
|
if end_index != -1:
|
|
return text[start_index:end_index].strip()
|
|
|
|
return ""
|
|
|
|
from typing import List, Dict
|
|
|
|
|
|
def sqlparser(left, right):
|
|
def format_sql(sql):
|
|
keywords = ['SELECT', 'FROM', 'WHERE', 'GROUP BY', 'ORDER BY', 'JOIN', 'INNER JOIN', 'HAVING']
|
|
formatted_sql = sql.upper()
|
|
for keyword in keywords:
|
|
formatted_sql = re.sub(rf'\b{keyword}\b', keyword, formatted_sql, flags=re.IGNORECASE)
|
|
return formatted_sql
|
|
|
|
left = format_sql(left)
|
|
right = format_sql(right)
|
|
|
|
components = {
|
|
'SELECT': r'SELECT\s+(.+?)\s+FROM',
|
|
'FROM': r'FROM\s+(.+?)(?=\s+(?:WHERE|GROUP BY|ORDER BY|JOIN|INNER JOIN|HAVING|$))',
|
|
'JOIN': r'(?:INNER )?JOIN\s+(.+?)\s+ON\s+(.+?)(?=\s*(?:WHERE|GROUP BY|ORDER BY|JOIN|INNER JOIN|HAVING|$))',
|
|
'WHERE': r'WHERE\s+(.+?)(?=\s+(?:GROUP BY|ORDER BY|HAVING|$))',
|
|
'GROUP BY': r'GROUP BY\s+(.+?)(?=\s+(?:ORDER BY|HAVING|$))',
|
|
'HAVING': r'HAVING\s+(.+?)(?=\s+(?:ORDER BY|$))',
|
|
'ORDER BY': r'ORDER BY\s+(.+)$'
|
|
}
|
|
|
|
def parse_component(sql, component, pattern):
|
|
if component == 'SELECT':
|
|
match = re.search(pattern, sql, re.IGNORECASE)
|
|
if match:
|
|
elements = match.group(1).split(',')
|
|
return set(element.strip() for element in elements)
|
|
return set()
|
|
elif component == 'JOIN':
|
|
joins = []
|
|
for match in re.finditer(
|
|
r'(?:INNER )?JOIN\s+(.+?)\s+ON\s+(.+?)(?=\s*(?:WHERE|GROUP BY|ORDER BY|JOIN|INNER JOIN|HAVING|$))',
|
|
sql,
|
|
re.IGNORECASE
|
|
):
|
|
joins.append((match.group(1).strip(), match.group(2).strip()))
|
|
return joins
|
|
elif component in ['GROUP BY', 'ORDER BY']:
|
|
match = re.search(pattern, sql, re.IGNORECASE)
|
|
if match:
|
|
elements = match.group(1).split(',')
|
|
return set(element.strip() for element in elements)
|
|
return set()
|
|
else:
|
|
match = re.search(pattern, sql, re.IGNORECASE)
|
|
return match.group(1).strip() if match else ''
|
|
|
|
left_components = {}
|
|
right_components = {}
|
|
for component, pattern in components.items():
|
|
left_components[component] = parse_component(left, component, pattern)
|
|
right_components[component] = parse_component(right, component, pattern)
|
|
|
|
score = 0
|
|
total = 0
|
|
for component in components:
|
|
lc = left_components[component]
|
|
rc = right_components[component]
|
|
if lc or rc:
|
|
total += 1
|
|
if lc == rc:
|
|
score += 1
|
|
return score / total if total != 0 else 0
|
|
def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
|
|
responses = [completion[0]['content'] for completion in completions]
|
|
extracted_responses = [extract_answer(r) for r in responses]
|
|
return [sqlparser(a, r) for r, a in zip(extracted_responses, answer)]
|
|
def bonus_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
|
|
responses = [completion[0]['content'] for completion in completions]
|
|
extracted_responses = [extract_answer(r) for r in responses]
|
|
return [1.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]
|
|
|
|
def strict_format_reward_func(completions, **kwargs) -> list[float]:
|
|
"""Reward function that checks if the completion has a specific format and minimum content length."""
|
|
pattern = r"^<think>(.*?)</think>.*"
|
|
responses = [completion[0]["content"] for completion in completions]
|
|
rewards = []
|
|
for r in responses:
|
|
match = re.match(pattern, r, re.DOTALL)
|
|
if match:
|
|
reasoning = match.group(1)
|
|
if len(reasoning.strip()) < 300:
|
|
rewards.append(0.0)
|
|
else:
|
|
rewards.append(0.1)
|
|
else:
|
|
rewards.append(0.0)
|
|
return rewards
|
|
from trl import GRPOConfig, GRPOTrainer
|
|
training_args = GRPOConfig(
|
|
use_vllm = True, # use vLLM for fast inference
|
|
learning_rate = learning_rate,
|
|
adam_beta1 = 0.9,
|
|
adam_beta2 = 0.999,
|
|
weight_decay = 0.1,
|
|
warmup_ratio = 0.1,
|
|
beta=beta,
|
|
lr_scheduler_type = "cosine",
|
|
optim = "paged_adamw_8bit",
|
|
logging_steps = 1,
|
|
bf16 = is_bfloat16_supported(),
|
|
fp16 = not is_bfloat16_supported(),
|
|
per_device_train_batch_size = per_device_train_batch_size,
|
|
gradient_accumulation_steps = gradient_accumulation_steps, # Increase to 4 for smoother training
|
|
num_generations = num_generations, # Decrease if out of memory
|
|
max_prompt_length = max_seq_length,
|
|
max_completion_length = max_seq_length//2,
|
|
num_train_epochs = num_train_epochs, # Set to 1 for a full training run
|
|
save_steps = 100,
|
|
max_grad_norm = 0.5,
|
|
report_to = "wandb", # Can use Weights & Biases
|
|
output_dir = output_dir,
|
|
)
|
|
import re
|
|
from datasets import load_dataset, Dataset
|
|
import pandas as pd
|
|
def dataset_process(name: str = dataset_dir) -> List[Dict]:
|
|
data = pd.read_csv(name)
|
|
processed_data = data.apply(lambda row: {
|
|
'prompt': [
|
|
{'role': 'system', 'content': SYS},
|
|
{'role': 'user', 'content': row['query']}
|
|
],
|
|
'answer': row['answer']
|
|
}, axis=1).tolist()
|
|
return processed_data
|
|
|
|
dataset = dataset_process()
|
|
|
|
from swanlab.integration.transformers import SwanLabCallback
|
|
swanlab_callback = SwanLabCallback(
|
|
project = "deepquery4-grpo",
|
|
experiment_name = "deepquery4-grpo-1",
|
|
)
|
|
|
|
trainer = GRPOTrainer(
|
|
model = model,
|
|
processing_class = tokenizer,
|
|
reward_funcs = [
|
|
strict_format_reward_func,
|
|
correctness_reward_func,
|
|
bonus_reward_func
|
|
],
|
|
args = training_args,
|
|
train_dataset = dataset,
|
|
callbacks = [swanlab_callback]
|
|
)
|
|
trainer.train(resume_from_checkpoint = True) |