初始化项目，由ModelHub XC社区提供模型
Model: ruohuaw/deepquery-1.5b-rl Source: Original Platform
This commit is contained in:
ModelHub XC
207
train.py
Normal file
207
train.py
Normal file
@@ -0,0 +1,207 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user