ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2025-12-01 15:23:43 +00:00
commit 727da805d9
42 changed files with 2394 additions and 2394 deletions
Download patch file Download diff file Expand all files Collapse all files

60
tests/utils/aime_eval.py
View file

@ -24,7 +24,7 @@ def download_and_combine_aime_datasets(data_dir: str = "./data/aime") -> str:
"test2025-II": "https://raw.githubusercontent.com/GAIR-NLP/AIME-Preview/main/eval/data/aime/test2025-II.jsonl",
}
os.makedirs(data_dir, exist_ok = True)
os.makedirs(data_dir, exist_ok=True)
combined_filepath = os.path.join(data_dir, "aime.jsonl")
# Check if combined file already exists
@ -67,9 +67,9 @@ def download_and_combine_aime_datasets(data_dir: str = "./data/aime") -> str:
# Write combined dataset
if all_problems:
with open(combined_filepath, "w", encoding = "utf-8") as f:
with open(combined_filepath, "w", encoding="utf-8") as f:
for problem in all_problems:
f.write(json.dumps(problem, ensure_ascii = False) + "\n")
f.write(json.dumps(problem, ensure_ascii=False) + "\n")
print(f"✅ Combined {len(all_problems)} problems from {len(datasets)} datasets")
print(f" Saved to: {combined_filepath}")
@ -92,7 +92,7 @@ def load_aime_dataset(data_dir: str = "./data/aime") -> List[Dict[str, Any]]:
filepath = download_and_combine_aime_datasets(data_dir)
examples = []
with open(filepath, "r", encoding = "utf-8") as f:
with open(filepath, "r", encoding="utf-8") as f:
for line_num, line in enumerate(f):
line = line.strip()
if line:
@ -188,20 +188,20 @@ def get_num_tokens(text, tokenizer_instance):
"""Count tokens in text"""
if not text:
return 0
encoding = tokenizer_instance(text, return_tensors = "pt")
encoding = tokenizer_instance(text, return_tensors="pt")
return len(encoding["input_ids"][0])
def evaluate_model_aime(
model,
tokenizer,
model_type = "base",
lora_request = None,
temperature = 0.3,
n_sampling = 8,
max_tokens = 32768,
top_p = 0.95,
seed = 0,
model_type="base",
lora_request=None,
temperature=0.3,
n_sampling=8,
max_tokens=32768,
top_p=0.95,
seed=0,
):
"""Evaluate model on combined AIME dataset with official configuration"""
@ -237,11 +237,11 @@ def evaluate_model_aime(
# Setup sampling parameters (AIME configuration)
sampling_params = SamplingParams(
temperature = temperature,
top_p = top_p,
max_tokens = max_tokens,
n = n_sampling, # Multiple samples per question
seed = seed,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
n=n_sampling, # Multiple samples per question
seed=seed,
)
print(f"\n🔧 Configuration:")
@ -272,13 +272,13 @@ def evaluate_model_aime(
# Main evaluation loop
with tqdm(
total = len(eval_dataset), desc = "Processing AIME problems", unit = "problem"
total=len(eval_dataset), desc="Processing AIME problems", unit="problem"
) as pbar:
for task_id, item in enumerate(eval_dataset):
try:
# Prepare prompt
prompt_text = tokenizer.apply_chat_template(
item["prompt"], add_generation_prompt = True, tokenize = False
item["prompt"], add_generation_prompt=True, tokenize=False
)
input_tokens.append(get_num_tokens(prompt_text, tokenizer))
@ -286,9 +286,9 @@ def evaluate_model_aime(
# Generate multiple responses
outputs = model.fast_generate(
[prompt_text],
sampling_params = sampling_params,
lora_request = lora_request,
use_tqdm = False,
sampling_params=sampling_params,
lora_request=lora_request,
use_tqdm=False,
)[0].outputs
# Process all generated responses
@ -413,8 +413,8 @@ def evaluate_model_aime(
# Save results
filename = f"aime_eval_combined_{model_type}_t{temperature}_n{n_sampling}.json"
with open(filename, "w", encoding = "utf-8") as f:
json.dump({"results": results, "records": records}, f, indent = 4)
with open(filename, "w", encoding="utf-8") as f:
json.dump({"results": results, "records": records}, f, indent=4)
# Print comprehensive summary
print(f"\n{'='*70}")
@ -517,27 +517,27 @@ def compare_aime_results(all_results):
if all_results and "source_accuracies" in all_results[0]:
datasets = list(all_results[0]["source_accuracies"].keys())
print(f"{'Model':<15}", end = "")
print(f"{'Model':<15}", end="")
for dataset in datasets:
print(f"{dataset:<15}", end = "")
print(f"{dataset:<15}", end="")
print()
print("-" * (15 + 15 * len(datasets)))
for result in all_results:
print(f"{result['model_type']:<15}", end = "")
print(f"{result['model_type']:<15}", end="")
for dataset in datasets:
accuracy = result["source_accuracies"].get(dataset, 0)
print(f"{accuracy:<15.1f}", end = "")
print(f"{accuracy:<15.1f}", end="")
print()
# Save comparison
comparison_data = {
"summary": all_results,
"best_model": max(all_results, key = lambda x: x["accuracy"]),
"best_model": max(all_results, key=lambda x: x["accuracy"]),
}
with open("aime_model_comparison.json", "w") as f:
json.dump(comparison_data, f, indent = 4)
json.dump(comparison_data, f, indent=4)
print(
f"\nBest performing model: {comparison_data['best_model']['model_type']} "

112
tests/utils/hf_utils.py
View file

@ -79,27 +79,27 @@ def generate_responses(
skip_special_tokens: bool = True,
dtype: torch.dtype = None,
):
inputs = [tokenizer(prompt, return_tensors = "pt") for _ in range(num_generations)]
inputs = [tokenizer(prompt, return_tensors="pt") for _ in range(num_generations)]
keys = inputs[0].keys()
batched_inputs = {
key: torch.cat([input[key] for input in inputs], dim = 0).to(model.device)
key: torch.cat([input[key] for input in inputs], dim=0).to(model.device)
for key in keys
}
if dtype is not None:
inference_context = torch.autocast(device_type = "cuda", dtype = dtype)
inference_context = torch.autocast(device_type="cuda", dtype=dtype)
else:
inference_context = nullcontext()
with inference_context:
outputs = model.generate(
**batched_inputs,
max_new_tokens = max_new_tokens,
do_sample = do_sample,
temperature = temperature,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
temperature=temperature,
)
responses = tokenizer.batch_decode(outputs, skip_special_tokens = skip_special_tokens)
responses = tokenizer.batch_decode(outputs, skip_special_tokens=skip_special_tokens)
return responses
@ -117,11 +117,11 @@ def sample_responses(
model,
tokenizer,
prompt,
temperature = temperature,
num_generations = num_generations,
max_new_tokens = max_new_tokens,
skip_special_tokens = skip_special_tokens,
dtype = dtype,
temperature=temperature,
num_generations=num_generations,
max_new_tokens=max_new_tokens,
skip_special_tokens=skip_special_tokens,
dtype=dtype,
)
return responses
@ -136,32 +136,32 @@ def setup_tokenizer(model_name, fixup_funcs: list[Callable] = []):
def setup_model(
model_name,
quantize: bool = True,
dtype = torch.bfloat16,
peft_config = None,
dtype=torch.bfloat16,
peft_config=None,
autocast_adapter: bool = True,
):
if quantize:
bnb_config = BitsAndBytesConfig(
load_in_4bit = True,
bnb_4bit_use_double_quant = True,
bnb_4bit_quant_type = "nf4",
bnb_4bit_compute_dtype = dtype,
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=dtype,
)
else:
bnb_config = None
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map = "cuda:0",
attn_implementation = "sdpa",
quantization_config = bnb_config,
torch_dtype = dtype,
device_map="cuda:0",
attn_implementation="sdpa",
quantization_config=bnb_config,
torch_dtype=dtype,
)
model = prepare_model_for_kbit_training(model) if quantize else model
if peft_config is not None:
model = get_peft_model(
model, peft_config, autocast_adapter_dtype = autocast_adapter
model, peft_config, autocast_adapter_dtype=autocast_adapter
)
return model
@ -169,19 +169,19 @@ def setup_model(
def get_peft_config(
lora_rank,
lora_alpha = None,
lora_dropout = 0.0,
bias = "none",
target_modules = "all-linear",
lora_alpha=None,
lora_dropout=0.0,
bias="none",
target_modules="all-linear",
):
lora_alpha = lora_alpha or 2 * lora_rank
peft_config = LoraConfig(
lora_alpha = lora_alpha,
lora_dropout = lora_dropout,
r = lora_rank,
bias = bias,
target_modules = target_modules,
task_type = "CAUSAL_LM",
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
r=lora_rank,
bias=bias,
target_modules=target_modules,
task_type="CAUSAL_LM",
)
return peft_config
@ -191,18 +191,18 @@ def setup_trainer(
tokenizer,
dataset,
train_args,
peft_config = None,
formatting_func = None,
collator = None,
peft_config=None,
formatting_func=None,
collator=None,
):
return SFTTrainer(
model = model,
peft_config = peft_config,
train_dataset = dataset,
processing_class = tokenizer,
formatting_func = formatting_func,
data_collator = collator,
args = train_args,
model=model,
peft_config=peft_config,
train_dataset=dataset,
processing_class=tokenizer,
formatting_func=formatting_func,
data_collator=collator,
args=train_args,
)
@ -212,17 +212,17 @@ def setup_lora(
dataset,
peft_config,
train_args,
formatting_func = None,
collator = None,
formatting_func=None,
collator=None,
):
return LoraConfig(
model = model,
peft_config = peft_config,
train_dataset = dataset,
processing_class = tokenizer,
formatting_func = formatting_func,
data_collator = collator,
args = train_args,
model=model,
peft_config=peft_config,
train_dataset=dataset,
processing_class=tokenizer,
formatting_func=formatting_func,
data_collator=collator,
args=train_args,
)
@ -236,7 +236,7 @@ def convert_weights_back_to_dtype(model, dtype):
param.data = param.data.to(dtype)
def fix_llama3_tokenizer(tokenizer, padding_side = "right"):
def fix_llama3_tokenizer(tokenizer, padding_side="right"):
tokenizer.padding_side = padding_side
added_vocab = tokenizer.get_added_vocab()
pad_token = [w for w in added_vocab if "pad" in w]
@ -276,12 +276,12 @@ def _convert_lora_to_linear(module: LoraLayer, adapter_name: str = "default"):
w_dq = w_dq.to(original_dtype)
new_module = torch.nn.Linear(
w_dq.shape[1], w_dq.shape[0], bias = module.base_layer.bias is not None
w_dq.shape[1], w_dq.shape[0], bias=module.base_layer.bias is not None
)
new_module.weight.data = torch.nn.Parameter(w_dq, requires_grad = False)
new_module.weight.data = torch.nn.Parameter(w_dq, requires_grad=False)
if module.lora_bias[adapter_name]:
bias_data = module.base_layer.bias.data + module.lora_B[adapter_name].bias
new_module.bias.data = torch.nn.Parameter(bias_data, requires_grad = False)
new_module.bias.data = torch.nn.Parameter(bias_data, requires_grad=False)
return new_module