krogoldAI/QueryRefiner-0.5B-v0.1-GRPO
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: krogoldAI/QueryRefiner-0.5B-v0.1-GRPO
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-10-22 21:45:35 +08:00
commit 6f9e8a2556
18 changed files with 153491 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
Codes/01 - GRPO Training.py Normal file
View File

@@ -0,0 +1,247 @@
import torch
import wandb
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset, Dataset
from huggingface_hub import login
from common.constants import SYSTEM_PROMPT
from common.rewards import SmoothXMLRewardEvaluator
##########
# CONFIG #
##########
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET = "krogoldAI/rag-ambiguous-queries"
OUTPUT_MODEL = "model/repo"
HF_TOKEN = "<MY_HF_TOKEN>"
# Training hyperparameters
NUM_EPOCHS = 1
PER_DEVICE_BATCH = 8
GRADIENT_ACCUMULATION = 1 # 2
LEARNING_RATE = 5e-6 # Low for GRPO stability
NUM_GENERATIONS = 2
MAX_NEW_TOKENS = 512
WARMUP_STEPS = 50
#########
# LOGIN #
#########
print("Logging into Hugging Face...")
login(HF_TOKEN)
##################
# LOAD TOKENIZER #
##################
print(f"Loading tokenizer from {MODEL_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.padding_side = "left"
# Ensure pad token is set
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
################
# LOAD DATASET #
################
print(f"Loading dataset {DATASET}...")
raw_dataset = load_dataset(DATASET, "default", split="train")
# Extract the list of questions
questions_list = raw_dataset["questions"]
print(f"Dataset size: {len(questions_list)} examples")
# Create a proper dataset from the list
dataset = Dataset.from_dict({"question": questions_list})
dataset = dataset.shuffle(seed=42)
# Take on 3000 examples
dataset = dataset.select(range(3000))
###########################
# PREPROCESSING FUNCTION #
###########################
def preprocess(example):
"""
Apply chat template to create proper prompt format.
"""
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": example["question"]},
]
# Apply chat template without tokenization to get the prompt string
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
return {
"prompt": prompt,
"query": example["question"] # Keep for reference
}
print("Preprocessing dataset with chat template...")
dataset = dataset.map(preprocess, remove_columns=dataset.column_names)
##############
# LOAD MODEL #
##############
print(f"\nLoading model {MODEL_NAME}...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16,
device_map="auto",
attn_implementation="flash_attention_2",
use_cache=False,
)
# If Liger Kernel available, apply now
try:
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
try:
apply_liger_kernel_to_qwen2(model)
except TypeError:
apply_liger_kernel_to_qwen2()
print("Liger Kernel applied successfully for Qwen2 optimization")
except Exception as e:
print(f"Liger Kernel not available or failed to apply: {e}")
print(f"Model loaded. Parameters: {model.num_parameters() / 1e9:.3f}B")
model.generation_config.max_new_tokens = MAX_NEW_TOKENS
model.generation_config.temperature = 0.7
model.generation_config.do_sample = True
###################
# REWARD FUNCTION #
###################
# Initialize reward evaluator
evaluator = SmoothXMLRewardEvaluator()
def reward_function(completions, prompts, **kwargs):
rewards = []
for i, (completion, prompt) in enumerate(zip(completions, prompts)):
# Extract generated portion
if completion.startswith(prompt):
response = completion[len(prompt):].strip()
else:
response = completion
try:
analysis = evaluator.get_detailed_analysis(response)
score = analysis["final_score"]
rewards.append(score)
except Exception as e:
print(f"ERROR in reward evaluation: {e}")
print(f"Response that failed: {response[:200]}")
rewards.append(0.0)
return rewards
##############
# MONITORING #
##############
# Monitoring the LLM from WandB
wandb.login(key = "<MY_WANDB_KEY>")
run = wandb.init(project='QueryRefiner-GRPO-Phase1', job_type="training", anonymous="allow")
######################
# TRAINING ARGUMENTS #
######################
print("Configuring training arguments...")
training_args = GRPOConfig(
output_dir="./qwen_grpo_checkpoints",
run_name="qwen-grpo-query-analysis",
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=PER_DEVICE_BATCH,
gradient_accumulation_steps=GRADIENT_ACCUMULATION,
learning_rate=LEARNING_RATE,
num_generations=NUM_GENERATIONS,
optim="adamw_torch_fused",
warmup_steps=WARMUP_STEPS,
lr_scheduler_type="cosine",
weight_decay=0.01,
max_grad_norm=1.0,
bf16=True,
gradient_checkpointing=True,
use_liger_kernel=True,
logging_steps=10,
save_steps=200,
save_total_limit=3,
push_to_hub=True,
hub_model_id=OUTPUT_MODEL,
hub_private_repo=True,
hub_strategy="checkpoint",
dataloader_num_workers=4,
dataloader_pin_memory=True,
remove_unused_columns=False,
report_to="wandb",
seed=42,
)
###########
# TRAINER #
###########
print("Initializing GRPOTrainer...")
trainer = GRPOTrainer(
model=model,
reward_funcs=reward_function,
args=training_args,
train_dataset=dataset, # Uses by default the "prompt" column, all other columns are ignored
)
#########
# TRAIN #
#########
print("=" * 50)
print("Starting GRPO training...")
print(f"Model: {MODEL_NAME}")
print(f"Dataset size: {len(dataset)}")
print(f"Epochs: {NUM_EPOCHS}")
print(f"Effective batch size: {PER_DEVICE_BATCH * GRADIENT_ACCUMULATION}")
print(f"Generations per prompt: {NUM_GENERATIONS}")
print("=" * 50)
trainer.train()
####################
# SAVE FINAL MODEL #
####################
print("\nSaving final model...")
final_model = trainer.model
final_model.config.use_cache = True # Re-enable cache for inference
final_model.save_pretrained("./final_grpo_model")
tokenizer.save_pretrained("./final_grpo_model")
###############
# PUSH TO HUB #
###############
print(f"\nPushing final model to {OUTPUT_MODEL}...")
try:
final_model.push_to_hub(OUTPUT_MODEL, private=True)
tokenizer.push_to_hub(OUTPUT_MODEL, private=True)
print("Successfully pushed model and tokenizer to Hub!")
except Exception as e:
print(f"Warning: Failed to push to hub: {e}")
print("Model saved locally at ./final_grpo_model")
print("\n" + "=" * 50)
print("GRPO Training Complete!")
print("=" * 50)

318
Codes/02 - Breadth-focused full SFT.py Normal file
View File

@@ -0,0 +1,318 @@
import os
import re
import numpy as np
import torch
import wandb
from datasets import load_dataset, Dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling,
)
from huggingface_hub import login
os.environ["TOKENIZERS_PARALLELISM"] = "false"
##########
# CONFIG #
##########
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET = "krogoldAI/rag-query-analysis"
OUTPUT_MODEL = "model/repo"
# Training hyperparams (tune as needed)
NUM_EPOCHS = 2
PER_DEVICE_BATCH = 4
GRADIENT_ACCUMULATION = 4
LEARNING_RATE = 2e-5
WEIGHT_DECAY = 0.01
WARMUP_STEPS = 100
BF16 = True
TORCH_COMPILE = False
#########
# LOGIN #
#########
login("<MY_HF_TOKEN>")
##################
# LOAD TOKENIZER #
##################
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.padding_side = "right"
################
# LOAD DATASET #
################
raw_ds = load_dataset(DATASET, "default", split="train")
raw_ds = raw_ds.shuffle(seed=42)
# Apply Qwen chat template and clean
def format_and_clean(batch):
out = []
for conv in batch["text"]:
s = tokenizer.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)
s = re.sub(r"```[a-zA-Z0-9]*\n?", "", s).replace("```", "").strip()
out.append(s)
return {"text": out}
ds = raw_ds.map(format_and_clean, batched=True, batch_size=128)
########################
# CUSTOM DATA COLLATOR #
########################
class Qwen25DataCollator(DataCollatorForLanguageModeling):
def __init__(self, tokenizer, mlm=False):
super().__init__(tokenizer=tokenizer, mlm=mlm)
# get token ids robustly (some tokenizers might return [] for encode if token missing)
try:
self.im_start_token = tokenizer.encode("<|im_start|>", add_special_tokens=False)[0]
except Exception:
self.im_start_token = None
try:
self.im_end_token = tokenizer.encode("<|im_end|>", add_special_tokens=False)[0]
except Exception:
self.im_end_token = None
# "assistant" token sequence (may be multiple tokens)
try:
self.assistant_text = tokenizer.encode("assistant", add_special_tokens=False)
except Exception:
self.assistant_text = []
# Provide both __call__ and torch_call for compatibility
def __call__(self, features):
return self.torch_call(features)
def torch_call(self, examples):
"""
examples: list of dicts returned by tokenization (each example contains 'input_ids', 'attention_mask', etc.)
We'll leverage the parent to create initial batch and then mask labels for assistant responses only.
"""
batch = super().torch_call(examples) # returns input_ids, attention_mask, labels (for MLM)
input_ids = batch["input_ids"]
labels = batch["labels"]
# If special tokens are not present, return default batch unchanged
if self.im_start_token is None or self.im_end_token is None or len(self.assistant_text) == 0:
return batch
# Iterate examples in batch to mask labels: only assistant response tokens should be supervised
for i, ids in enumerate(input_ids):
# Find positions of <|im_start|> and <|im_end|>
im_start_positions = torch.where(ids == self.im_start_token)[0]
im_end_positions = torch.where(ids == self.im_end_token)[0]
if im_start_positions.numel() == 0 or im_end_positions.numel() == 0:
# no recognized chat markers: leave labels as-is (or continue)
continue
last_assistant_start = None
# Find last im_start that is followed by "assistant"
for start_pos in im_start_positions:
# check if tokens following start_pos match "assistant"
as_len = len(self.assistant_text)
candidate_end = start_pos + 1 + as_len
if candidate_end <= len(ids):
segment = ids[start_pos + 1:start_pos + 1 + as_len]
if torch.equal(segment, torch.tensor(self.assistant_text, device=ids.device)):
last_assistant_start = int(start_pos)
if last_assistant_start is None:
continue
# Find first im_end after last_assistant_start
assistant_end_positions = im_end_positions[im_end_positions > last_assistant_start]
if assistant_end_positions.numel() == 0:
continue
assistant_end = int(assistant_end_positions[0])
# Response text is between (last_assistant_start + 1 + len("assistant")) and assistant_end - 1 (inclusive),
# but because template may include a newline or an extra token, we set response_start carefully.
response_start = last_assistant_start + 1 + len(self.assistant_text)
# If there's a newline token or separator, skip it if present in input_ids
# (this is conservative: we do not assume an extra token, but we keep it if present)
if response_start < len(ids) and ids[response_start] == tokenizer.encode("\n", add_special_tokens=False)[0]:
response_start += 1
# Apply masking:
# Set everything before response_start to -100 (ignored), preserve response tokens, set rest to -100
labels[i, :] = -100
if response_start < len(ids):
# labels slice up to assistant_end inclusive
end_idx = min(assistant_end + 1, ids.shape[0])
labels[i, response_start:end_idx] = ids[response_start:end_idx]
# assign modified labels back
batch["labels"] = labels
return batch
collator = Qwen25DataCollator(tokenizer=tokenizer, mlm=False)
###############################################
# ANALYZE DATASET LENGTHS TO SET `max_length` #
###############################################
# We analyze the dataset to optimize the choice of `max_length`
print("Analyzing dataset to determine max_length (sample up to 1000)...")
assistant_lengths = []
full_lengths = []
sample_limit = min(1000, len(ds))
for example in ds["text"][:sample_limit]:
full_tokens = tokenizer(example, truncation=False, add_special_tokens=True)
full_lengths.append(len(full_tokens["input_ids"]))
# extract the last assistant response via regex pattern
pattern = r"<\|im_start\|>assistant\n(.*?)<\|im_end\|>"
matches = re.findall(pattern, example, re.DOTALL)
if matches:
last_response = matches[-1]
resp_tokens = tokenizer(last_response, truncation=False, add_special_tokens=False)
assistant_lengths.append(len(resp_tokens["input_ids"]))
# Basic statistics (guard for empty lists)
def safe_stat(arr):
if len(arr) == 0:
return 0.0, 0.0, 0.0, 0.0
return np.mean(arr), np.median(arr), np.percentile(arr, 95), np.percentile(arr, 99)
mean_ass, med_ass, p95_ass, p99_ass = safe_stat(assistant_lengths)
mean_full, _, p95_full, _ = safe_stat(full_lengths)
print(f"Assistant response mean={mean_ass:.1f}, median={med_ass:.1f}, 95%={p95_ass:.1f}, 99%={p99_ass:.1f}")
print(f"Full conversation mean={mean_full:.1f}, 95%={p95_full:.1f}")
# Round up to nearest power of two but don't exceed tokenizer.model_max_length
def next_power_of_2(x):
if x <= 1:
return 1
return 2 ** int(np.ceil(np.log2(x)))
target_length = int(min(p95_full if p95_full > 0 else tokenizer.model_max_length, tokenizer.model_max_length))
MAX_LENGTH = next_power_of_2(target_length)
if MAX_LENGTH > tokenizer.model_max_length:
MAX_LENGTH = tokenizer.model_max_length
print(f"Using MAX_LENGTH = {MAX_LENGTH}")
####################
# TOKENIZE DATASET #
####################
def tokenize_function(examples):
return tokenizer(examples["text"], truncation=True, max_length=MAX_LENGTH, padding=False)
tokenized_ds = ds.map(tokenize_function, batched=True, remove_columns=ds.column_names)
##############
# LOAD MODEL #
##############
# Load model in bfloat16
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16 if BF16 else None,
device_map="auto",
attn_implementation="flash_attention_2",
use_cache=False,
)
# If Liger Kernel available, apply now
try:
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
try:
apply_liger_kernel_to_qwen2(model)
except TypeError:
apply_liger_kernel_to_qwen2()
print("Liger Kernel applied successfully for Qwen2 optimization")
except Exception:
print("Liger Kernel not available or failed to apply; continuing without it.")
print(f"Model loaded. Parameters: {model.num_parameters() / 1e9:.3f}B")
##############
# MONITORING #
##############
# Monitoring the LLM from WandB
wandb.login(key = "<MY_WANDB_KEY>")
run = wandb.init(project='QueryRefiner-GRPO-Phase2', job_type="training", anonymous="allow")
######################
# TRAINING ARGUMENTS #
######################
training_args = TrainingArguments(
output_dir="./qwen_rephraser_checkpoints",
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=PER_DEVICE_BATCH,
gradient_accumulation_steps=GRADIENT_ACCUMULATION,
learning_rate=LEARNING_RATE,
weight_decay=WEIGHT_DECAY,
warmup_steps=WARMUP_STEPS,
lr_scheduler_type="cosine",
logging_steps=10,
save_steps=500,
save_total_limit=2,
bf16=BF16,
optim="adamw_torch_fused",
gradient_checkpointing=True,
report_to="wandb",
push_to_hub=False, # we'll push manually at the end
hub_model_id=OUTPUT_MODEL,
hub_private_repo=True,
dataloader_num_workers=4,
dataloader_pin_memory=True,
ddp_find_unused_parameters=False,
torch_compile=TORCH_COMPILE,
)
###########
# TRAINER #
###########
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_ds,
data_collator=collator,
)
#########
# TRAIN #
#########
print("Starting training...")
trainer.train()
####################
# SAVE FINAL MODEL #
####################
print("Saving model to ./final_model ...")
model.config.use_cache = True
trainer.save_model("./final_model")
tokenizer.save_pretrained("./final_model")
##################
# PUSHING TO HUB #
##################
try:
print(f"Pushing model and tokenizer to the hub as {OUTPUT_MODEL} (private)...")
model.push_to_hub(OUTPUT_MODEL, private=True)
tokenizer.push_to_hub(OUTPUT_MODEL, private=True)
print("Push completed.")
except Exception as e:
print("Warning: push_to_hub failed:", e)
print("Training complete!")

339
Codes/03 - Depth-focused full SFT.py Normal file
View File

@@ -0,0 +1,339 @@
import os
import re
import numpy as np
import torch
import wandb
from datasets import load_dataset, Dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling,
)
from huggingface_hub import login, hf_hub_download
os.environ["TOKENIZERS_PARALLELISM"] = "false"
##########
# CONFIG #
##########
MODEL_NAME = "krogoldAI/QueryRefiner-0.5B-v0.1-GRPO-SFT1"
DATASET = "krogoldAI/rag-query-analysis"
OUTPUT_MODEL = "model/repo"
# Training hyperparams
NUM_EPOCHS = 3
PER_DEVICE_BATCH = 4
GRADIENT_ACCUMULATION = 4
LEARNING_RATE = 2e-5
WEIGHT_DECAY = 0.01
WARMUP_STEPS = 50
BF16 = True
TORCH_COMPILE = False
#########
# LOGIN #
#########
login("<MY_HF_TOKEN>")
##################
# LOAD TOKENIZER #
##################
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.padding_side = "right"
################
# LOAD DATASET #
################
raw_ds = load_dataset(DATASET, "default", split="train")
raw_ds = raw_ds.shuffle(seed=42)
# Metrics to filter on
metrics = [
"ambiguity_assessment",
"domain_accuracy",
"follows_guidelines",
"intent_accuracy",
"intent_preservation",
"rephrasing_quality",
]
# Filter on metrics
def check(metadata):
if metadata.get("critical_issue") is not None or metadata.get("usable") == False:
return False
for metric in metrics:
if metadata.get(metric) != 5:
return False
return True
filtered_ds = raw_ds.filter(lambda x: check(x["metadata"]))
# Apply Qwen chat template and clean
def format_and_clean(batch):
out = []
for conv in batch["text"]:
s = tokenizer.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)
s = re.sub(r"```[a-zA-Z0-9]*\n?", "", s).replace("```", "").strip()
out.append(s)
return {"text": out}
ds = filtered_ds.map(format_and_clean, batched=True, batch_size=128)
########################
# CUSTOM DATA COLLATOR #
########################
class Qwen25DataCollator(DataCollatorForLanguageModeling):
def __init__(self, tokenizer, mlm=False):
super().__init__(tokenizer=tokenizer, mlm=mlm)
# get token ids robustly (some tokenizers might return [] for encode if token missing)
try:
self.im_start_token = tokenizer.encode("<|im_start|>", add_special_tokens=False)[0]
except Exception:
self.im_start_token = None
try:
self.im_end_token = tokenizer.encode("<|im_end|>", add_special_tokens=False)[0]
except Exception:
self.im_end_token = None
# "assistant" token sequence (may be multiple tokens)
try:
self.assistant_text = tokenizer.encode("assistant", add_special_tokens=False)
except Exception:
self.assistant_text = []
# Provide both __call__ and torch_call for compatibility
def __call__(self, features):
return self.torch_call(features)
def torch_call(self, examples):
"""
examples: list of dicts returned by tokenization (each example contains 'input_ids', 'attention_mask', etc.)
We'll leverage the parent to create initial batch and then mask labels for assistant responses only.
"""
batch = super().torch_call(examples) # returns input_ids, attention_mask, labels (for MLM)
input_ids = batch["input_ids"]
labels = batch["labels"]
# If special tokens are not present, return default batch unchanged
if self.im_start_token is None or self.im_end_token is None or len(self.assistant_text) == 0:
return batch
# Iterate examples in batch to mask labels: only assistant response tokens should be supervised
for i, ids in enumerate(input_ids):
# Find positions of <|im_start|> and <|im_end|>
im_start_positions = torch.where(ids == self.im_start_token)[0]
im_end_positions = torch.where(ids == self.im_end_token)[0]
if im_start_positions.numel() == 0 or im_end_positions.numel() == 0:
# no recognized chat markers: leave labels as-is (or continue)
continue
last_assistant_start = None
# Find last im_start that is followed by "assistant"
for start_pos in im_start_positions:
# check if tokens following start_pos match "assistant"
as_len = len(self.assistant_text)
candidate_end = start_pos + 1 + as_len
if candidate_end <= len(ids):
segment = ids[start_pos + 1:start_pos + 1 + as_len]
if torch.equal(segment, torch.tensor(self.assistant_text, device=ids.device)):
last_assistant_start = int(start_pos)
if last_assistant_start is None:
continue
# Find first im_end after last_assistant_start
assistant_end_positions = im_end_positions[im_end_positions > last_assistant_start]
if assistant_end_positions.numel() == 0:
continue
assistant_end = int(assistant_end_positions[0])
# Response text is between (last_assistant_start + 1 + len("assistant")) and assistant_end - 1 (inclusive),
# but because template may include a newline or an extra token, we set response_start carefully.
response_start = last_assistant_start + 1 + len(self.assistant_text)
# If there's a newline token or separator, skip it if present in input_ids
# (this is conservative: we do not assume an extra token, but we keep it if present)
if response_start < len(ids) and ids[response_start] == tokenizer.encode("\n", add_special_tokens=False)[0]:
response_start += 1
# Apply masking:
# Set everything before response_start to -100 (ignored), preserve response tokens, set rest to -100
labels[i, :] = -100
if response_start < len(ids):
# labels slice up to assistant_end inclusive
end_idx = min(assistant_end + 1, ids.shape[0])
labels[i, response_start:end_idx] = ids[response_start:end_idx]
# assign modified labels back
batch["labels"] = labels
return batch
collator = Qwen25DataCollator(tokenizer=tokenizer, mlm=False)
###############################################
# ANALYZE DATASET LENGTHS TO SET `max_length` #
###############################################
# We analyze the dataset to optimize the choice of `max_length`
print("Analyzing dataset to determine max_length (sample up to 1000)...")
assistant_lengths = []
full_lengths = []
sample_limit = min(1000, len(ds))
for example in ds["text"][:sample_limit]:
full_tokens = tokenizer(example, truncation=False, add_special_tokens=True)
full_lengths.append(len(full_tokens["input_ids"]))
# extract the last assistant response via regex pattern
pattern = r"<\|im_start\|>assistant\n(.*?)<\|im_end\|>"
matches = re.findall(pattern, example, re.DOTALL)
if matches:
last_response = matches[-1]
resp_tokens = tokenizer(last_response, truncation=False, add_special_tokens=False)
assistant_lengths.append(len(resp_tokens["input_ids"]))
# Basic statistics (guard for empty lists)
def safe_stat(arr):
if len(arr) == 0:
return 0.0, 0.0, 0.0, 0.0
return np.mean(arr), np.median(arr), np.percentile(arr, 95), np.percentile(arr, 99)
mean_ass, med_ass, p95_ass, p99_ass = safe_stat(assistant_lengths)
mean_full, _, p95_full, _ = safe_stat(full_lengths)
print(f"Assistant response mean={mean_ass:.1f}, median={med_ass:.1f}, 95%={p95_ass:.1f}, 99%={p99_ass:.1f}")
print(f"Full conversation mean={mean_full:.1f}, 95%={p95_full:.1f}")
# Round up to nearest power of two but don't exceed tokenizer.model_max_length
def next_power_of_2(x):
if x <= 1:
return 1
return 2 ** int(np.ceil(np.log2(x)))
target_length = int(min(p95_full if p95_full > 0 else tokenizer.model_max_length, tokenizer.model_max_length))
MAX_LENGTH = next_power_of_2(target_length)
if MAX_LENGTH > tokenizer.model_max_length:
MAX_LENGTH = tokenizer.model_max_length
print(f"Using MAX_LENGTH = {MAX_LENGTH}")
####################
# TOKENIZE DATASET #
####################
def tokenize_function(examples):
return tokenizer(examples["text"], truncation=True, max_length=MAX_LENGTH, padding=False)
tokenized_ds = ds.map(tokenize_function, batched=True, remove_columns=ds.column_names)
##############
# LOAD MODEL #
##############
# Load model in bfloat16
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16 if BF16 else None,
device_map="auto",
attn_implementation="flash_attention_2",
use_cache=False,
)
# If Liger Kernel available, apply now
try:
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
try:
apply_liger_kernel_to_qwen2(model)
except TypeError:
apply_liger_kernel_to_qwen2()
print("Liger Kernel applied successfully for Qwen2 optimization")
except Exception:
print("Liger Kernel not available or failed to apply; continuing without it.")
print(f"Model loaded. Parameters: {model.num_parameters() / 1e9:.3f}B")
##############
# MONITORING #
##############
# Monitoring the LLM from WandB
wandb.login(key = "<MY_WANDB_KEY>")
run = wandb.init(project='QueryRefiner-GRPO-Phase2', job_type="training", anonymous="allow")
######################
# TRAINING ARGUMENTS #
######################
training_args = TrainingArguments(
output_dir="./qwen_rephraser_checkpoints",
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=PER_DEVICE_BATCH,
gradient_accumulation_steps=GRADIENT_ACCUMULATION,
learning_rate=LEARNING_RATE,
weight_decay=WEIGHT_DECAY,
warmup_steps=WARMUP_STEPS,
lr_scheduler_type="cosine",
logging_steps=10,
save_steps=500,
save_total_limit=2,
bf16=BF16,
optim="adamw_torch_fused",
gradient_checkpointing=True,
report_to="wandb",
push_to_hub=False, # we'll push manually at the end
hub_model_id=OUTPUT_MODEL,
hub_private_repo=True,
dataloader_num_workers=4,
dataloader_pin_memory=True,
ddp_find_unused_parameters=False,
torch_compile=TORCH_COMPILE,
)
###########
# TRAINER #
###########
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_ds,
data_collator=collator,
)
#########
# TRAIN #
#########
print("Starting training...")
trainer.train()
####################
# SAVE FINAL MODEL #
####################
print("Saving model to ./final_model ...")
model.config.use_cache = True
trainer.save_model("./final_model")
tokenizer.save_pretrained("./final_model")
##################
# PUSHING TO HUB #
##################
try:
print(f"Pushing model and tokenizer to the hub as {OUTPUT_MODEL} (private)...")
model.push_to_hub(OUTPUT_MODEL, private=True)
tokenizer.push_to_hub(OUTPUT_MODEL, private=True)
print("Push completed.")
except Exception as e:
print("Warning: push_to_hub failed:", e)
print("Training complete!")

74
Codes/common/constants.py Normal file
View File

@@ -0,0 +1,74 @@
SYSTEM_PROMPT = """You are a query analysis and rephraser for a Retrieval-Augmented Generation (RAG) system.
Your sole task is to **analyze user queries** and output a structured XML document.
You must **not answer the query itself**, only analyze and rephrase it.
## RAG Query Optimization
Effective rephrasing should optimize for document retrieval by:
- Using **specific terminology** and domain vocabulary likely to appear in relevant documents
- **Expanding acronyms** when they add context (but not when the acronym itself is the subject)
- **Adding disambiguating context** without over-constraining the search
- **Making implicit references explicit** using placeholders for missing entities (e.g., [PERSON], [COMPANY])
- **Preserving user intent** while improving retrieval precision
Examples: "How do I reset my password?" → "password reset procedure authentication"
"What's their revenue?" → "What's [COMPANY]'s revenue?"
## Analysis Process
Follow this systematic approach to decompose each query:
1. **Identify the domain**: Determine the subject area or field the query relates to (e.g., banking, healthcare, technology, legal). Consider both explicit domain indicators and contextual clues.
2. **Determine the intent**: Classify what the user is trying to accomplish (e.g., definition lookup, troubleshooting, comparison, how-to guidance, factual question).
3. **Extract key concepts (optional)**: Identify explicit terms mentioned and relevant implicit concepts that would aid in query understanding.
4. **Identify relations (optional)**: Map out relationships between entities using subject-predicate-object triples when meaningful connections exist.
5. **Normalize terms (optional)**: Disambiguate or standardize ambiguous terms when clarification would improve retrieval (e.g., "Apple" → "Apple Inc." vs "apple fruit").
6. **Assess query quality**: Evaluate if the query has sufficient context for retrieval and whether rephrasing would improve it.
7. **Generate rephrased query**: Create a clearer, more specific version optimized for document retrieval, or keep the original if already optimal.
## Technical Rules
1. **Never answer the user's question.** Only analyze and rephrase.
2. Always produce valid XML strictly following the schema below.
3. `<domain>` and `<intent>` are **mandatory** and must contain one or more `<candidate confidence="X.X">...</candidate>` entries:
- Confidence scores must always sum to 1.0
- If unambiguous: **exactly one candidate** with `confidence="1.0"` and `ambiguous="false"`
- If ambiguous: multiple candidates with `ambiguous="true"` and confidence distributed proportionally to plausibility:
- Use uniform distribution only when candidates are genuinely equally likely
- Otherwise, weight confidence toward the more probable interpretation
- Examples:
- "What is Mercury's rotation period?" → Astronomy 0.5, Chemistry 0.5 (equally plausible)
- "Jaguar speed in the wild" → Zoology 0.8, Automotive 0.2 (context favors animal)
4. Confidence values must always have one decimal place (e.g., `0.5`, `1.0`).
5. Only `<concepts>`, `<relations>`, and `<normalized_terms>` are optional. **All other elements are mandatory.**
6. `<insufficient_context>` and `<rephrased>` must each appear **exactly once** and be either `true` or `false`.
7. `<rephrased_query>` must always appear, even if identical to the input.
8. **Output only valid XML.** Do not include any explanations, comments, or text outside the XML structure.
9. All elements must appear in the order specified in the schema:
`<domain> → <intent> → <concepts> → <relations> → <normalized_terms> → <insufficient_context> → <rephrased> → <rephrased_query>`.
## Output Schema
<query_analysis>
<domain ambiguous="true|false">
<candidate confidence="X.X">...</candidate>
</domain>
<intent ambiguous="true|false">
<candidate confidence="X.X">...</candidate>
</intent>
<!-- Optional sections -->
<concepts>
<explicit>...</explicit>
<implicit>...</implicit>
</concepts>
<relations>
<relation subject="..." predicate="..." object="..."/>
</relations>
<normalized_terms>
<term original="..." normalized="..."/>
</normalized_terms>
<!-- End optional sections -->
<insufficient_context>true|false</insufficient_context>
<rephrased>true|false</rephrased>
<rephrased_query>...</rephrased_query>
</query_analysis>
"""

204
Codes/common/rewards.py Normal file
View File

@@ -0,0 +1,204 @@
import re
import xml.etree.ElementTree as ET
import math
from typing import Dict, Tuple
class SmoothXMLRewardEvaluator:
"""
Smooth, differentiable XML reward function.
Returns continuous scores between 0.0 and 1.0 for all components.
"""
def __init__(self):
self.composite_weights = {
'structure': 0.30,
'xml_valid': 0.25,
'order': 0.25,
'confidence': 0.18,
'distribution': 0.02
}
def evaluate_structure(self, output: str) -> float:
"""
Check structure (tags present).
"""
required_tags = [
r"<query_analysis>.*</query_analysis>",
r"<domain ambiguous=\"(true|false)\">.*</domain>",
r"<intent ambiguous=\"(true|false)\">.*</intent>",
r"<candidate confidence=\"(?:0\.\d|1\.0)\">.*?</candidate>",
r"<insufficient_context>(true|false)</insufficient_context>",
r"<rephrased>(true|false)</rephrased>",
r"<rephrased_query>.*?</rephrased_query>"
]
hits = sum(bool(re.search(tag, output, re.DOTALL)) for tag in required_tags)
return hits / len(required_tags)
def evaluate_xml_validity(self, output: str) -> float:
"""
Check XML validity.
"""
try:
ET.fromstring(output.strip())
return 1.0
except ET.ParseError:
return 0.0
def evaluate_order(self, output: str) -> float:
"""
Check order of elements (fraction of correct sequence for continuity).
"""
sequence = ["<query_analysis>", "<domain", "<intent", "<insufficient_context>", "<rephrased>", "<rephrased_query>"]
last_index = -1
correct_count = 0
for tag in sequence:
match = re.search(tag, output)
if match:
idx = match.start()
if idx > last_index:
correct_count += 1
last_index = idx
return correct_count / len(sequence)
def evaluate_confidence(self, output: str) -> float:
"""
Check confidence correctness.
"""
score = 0.0
blocks = ['domain', 'intent']
for block_name in blocks:
block_match = re.search(f"<{block_name} ambiguous=\"(true|false)\">.*?</{block_name}>", output, re.DOTALL)
if not block_match:
continue
try:
is_ambiguous = block_match.group(1) == "true"
confidences = [float(c) for c in re.findall(r"<candidate confidence=\"(0\.\d|1\.0)\">", block_match.group(0))]
if not confidences:
continue
if is_ambiguous:
target_sum = 1.0
actual_sum = sum(confidences)
# Continuous score: closer to 1.0 sum → higher reward
score += max(0, 1 - abs(actual_sum - target_sum))
else:
if len(confidences) == 1:
score += 1.0 - abs(confidences[0] - 1.0)
except (ValueError, AttributeError):
continue
return score / len(blocks) if blocks else 0.0
def evaluate_distribution(self, output: str) -> float:
"""
Check confidence distribution (entropy-based, normalized).
"""
total_score = 0.0
blocks_evaluated = 0
for block_name in ['domain', 'intent']:
block_match = re.search(f'<{block_name} ambiguous="(true|false)".*?</{block_name}>', output, re.DOTALL)
if not block_match:
continue
is_ambiguous = block_match.group(1) == "true"
confidences = [float(c) for c in re.findall(r'confidence="([01]\.\d)"', block_match.group(0))]
if not confidences:
continue
blocks_evaluated += 1
if is_ambiguous and len(confidences) > 1:
# entropy normalized
entropy = -sum(p * math.log(p + 1e-8) for p in confidences)
max_entropy = math.log(len(confidences))
total_score += entropy / max_entropy if max_entropy > 0 else 0
else:
total_score += 1.0
return total_score / blocks_evaluated if blocks_evaluated > 0 else 0.0
def structural_penalty(self, output: str) -> float:
"""
Compute structural penalty (soft, subtractive).
"""
penalty = 0.0
try:
root = ET.fromstring(output.strip())
except ET.ParseError:
return 1.0
for tag in root.findall(".//domain") + root.findall(".//intent"):
if "ambiguous" not in tag.attrib:
penalty += 0.1
for cand in root.findall(".//candidate"):
if "confidence" not in cand.attrib:
penalty += 0.05
return min(1.0, penalty)
def answering_penalty(self, output: str) -> float:
"""
Compute answering penalty.
"""
stripped = output.strip()
if stripped.startswith('<query_analysis>') and stripped.endswith('</query_analysis>'):
return 0.0
return 1.0
def evaluate(self, output: str) -> Tuple[float, Dict[str, float]]:
"""
Return a composite evaluation, using intermediary checks.
"""
ap = self.answering_penalty(output)
if ap > 0:
return 0.0, {"answering_penalty": ap}
components = {
"structure": self.evaluate_structure(output),
"xml_valid": self.evaluate_xml_validity(output),
"order": self.evaluate_order(output),
"confidence": self.evaluate_confidence(output),
"distribution": self.evaluate_distribution(output)
}
sp = self.structural_penalty(output)
components["structural_penalty"] = sp
# Subtractive penalty instead of multiplicative
reward = sum(self.composite_weights[k] * v for k, v in components.items() if k != "structural_penalty")
final_score = max(0.0, reward - 0.5 * sp)
components["final_score"] = final_score
return final_score, components
def get_detailed_analysis(self, output: str) -> Dict[str, any]:
"""
Provides detailed analysis of the XML output including all component scores.
"""
final_score, component_scores = self.evaluate(output)
return {
'final_score': final_score,
'component_scores': component_scores,
'weights_used': self.composite_weights,
'recommendations': self._get_recommendations(component_scores)
}
def _get_recommendations(self, component_scores: Dict[str, float]) -> list:
"""
Generate improvement recommendations based on component scores.
"""
recommendations = []
if component_scores.get('structure', 1.0) < 0.85:
recommendations.append("Improve XML structure - ensure all required tags are present")
if component_scores.get('xml_valid', 1.0) < 1.0:
recommendations.append("Fix XML syntax errors - ensure proper tag closing and nesting")
if component_scores.get('order', 1.0) < 0.9:
recommendations.append("Reorder XML elements to match expected structure")
if component_scores.get('confidence', 1.0) < 0.8:
recommendations.append("Fix confidence values - ensure they sum to 1.0 for ambiguous cases")
if component_scores.get('distribution', 1.0) < 0.8:
recommendations.append("Improve confidence distribution balance for ambiguous classifications")
if component_scores.get('structural_penalty', 0.0) > 0.2:
recommendations.append("Address structural issues - missing attributes or malformed tags")
return recommendations