reroute merge logic language models + comprehensive tests + eval kits (#2673)
This commit is contained in:
Roland Tannous
GitHub
parent
332eabf309
commit
58f3a6e29d
19 changed files with 4760 additions and 7 deletions
|
|
@ -0,0 +1,218 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
"""Load model and compute perplexity in subprocess"""
|
||||
from unsloth import FastLanguageModel
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=load_in_4bit,
|
||||
load_in_8bit=load_in_8bit,
|
||||
)
|
||||
# Set up tokenizer
|
||||
merged_tokenizer = get_chat_template(
|
||||
merged_tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
# Format the dataset
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [merged_tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Compute perplexity using the passed dataset
|
||||
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
||||
|
||||
|
||||
# IMPORTANT: Convert to Python float if it's a tensor
|
||||
if torch.is_tensor(ppl_value):
|
||||
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
||||
elif hasattr(ppl_value, 'item'):
|
||||
ppl_value = ppl_value.item() # Convert numpy or other array types
|
||||
else:
|
||||
ppl_value = float(ppl_value) # Ensure it's a float
|
||||
|
||||
# Return only the perplexity value
|
||||
result_queue.put(ppl_value)
|
||||
|
||||
# Clean up
|
||||
del merged_model
|
||||
del merged_tokenizer
|
||||
del dataset_ppl
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Llama-3.2-3B-Instruct",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=10,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
|
||||
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
# saving and merging the model to local disk
|
||||
print("merge and save to local disk")
|
||||
model.save_pretrained_merged(
|
||||
save_directory='./unsloth_out/merged_llama_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# print("cleaning")
|
||||
# del model
|
||||
# del tokenizer
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
|
||||
print("Computing 8-bit model perplexity in subprocess...")
|
||||
result_queue = mp.Queue()
|
||||
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
ppl_8bit = result_queue.get()
|
||||
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
||||
|
||||
print("Loading merged model in 16 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=False,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
print_model_comparison()
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -0,0 +1,298 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
"""Load model and compute perplexity in subprocess"""
|
||||
from unsloth import FastLanguageModel
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_mistral_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=load_in_4bit,
|
||||
load_in_8bit=load_in_8bit,
|
||||
)
|
||||
# Set up tokenizer
|
||||
# merged_tokenizer = get_chat_template(
|
||||
# merged_tokenizer,
|
||||
# chat_template="llama-3.1",
|
||||
# )
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
||||
|
||||
### Instruction:
|
||||
{}
|
||||
|
||||
### Input:
|
||||
{}
|
||||
|
||||
### Response:
|
||||
{}"""
|
||||
|
||||
EOS_TOKEN = merged_tokenizer.eos_token
|
||||
|
||||
def formatting_prompts_func(examples):
|
||||
instructions = []
|
||||
inputs = []
|
||||
outputs = []
|
||||
texts = []
|
||||
|
||||
for conversation in examples["messages"]:
|
||||
# Extract user message and assistant response
|
||||
user_message = ""
|
||||
assistant_message = ""
|
||||
|
||||
for turn in conversation:
|
||||
if turn["role"] == "user":
|
||||
user_message = turn["content"]
|
||||
elif turn["role"] == "assistant":
|
||||
assistant_message = turn["content"]
|
||||
|
||||
# Store intermediate format
|
||||
instruction = "Complete the statement"
|
||||
instructions.append(instruction)
|
||||
inputs.append(user_message)
|
||||
outputs.append(assistant_message)
|
||||
|
||||
# Create formatted text
|
||||
text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
|
||||
texts.append(text)
|
||||
|
||||
return {
|
||||
"instruction": instructions,
|
||||
"input": inputs,
|
||||
"output": outputs,
|
||||
"text": texts
|
||||
}
|
||||
|
||||
|
||||
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Compute perplexity using the passed dataset
|
||||
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
||||
|
||||
|
||||
# IMPORTANT: Convert to Python float if it's a tensor
|
||||
if torch.is_tensor(ppl_value):
|
||||
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
||||
elif hasattr(ppl_value, 'item'):
|
||||
ppl_value = ppl_value.item() # Convert numpy or other array types
|
||||
else:
|
||||
ppl_value = float(ppl_value) # Ensure it's a float
|
||||
|
||||
# Return only the perplexity value
|
||||
result_queue.put(ppl_value)
|
||||
|
||||
# Clean up
|
||||
del merged_model
|
||||
del merged_tokenizer
|
||||
del dataset_ppl
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/mistral-7b-v0.3",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
|
||||
EOS_TOKEN = tokenizer.eos_token
|
||||
|
||||
|
||||
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
||||
|
||||
### Instruction:
|
||||
{}
|
||||
|
||||
### Input:
|
||||
{}
|
||||
|
||||
### Response:
|
||||
{}"""
|
||||
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
instructions = []
|
||||
inputs = []
|
||||
outputs = []
|
||||
texts = []
|
||||
|
||||
for conversation in examples["messages"]:
|
||||
# Extract user message and assistant response
|
||||
user_message = ""
|
||||
assistant_message = ""
|
||||
|
||||
for turn in conversation:
|
||||
if turn["role"] == "user":
|
||||
user_message = turn["content"]
|
||||
elif turn["role"] == "assistant":
|
||||
assistant_message = turn["content"]
|
||||
|
||||
# Store intermediate format
|
||||
instruction = "Complete the statement"
|
||||
instructions.append(instruction)
|
||||
inputs.append(user_message)
|
||||
outputs.append(assistant_message)
|
||||
|
||||
# Create formatted text
|
||||
text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
|
||||
texts.append(text)
|
||||
|
||||
|
||||
return {
|
||||
"instruction": instructions,
|
||||
"input": inputs,
|
||||
"output": outputs,
|
||||
"text": texts
|
||||
}
|
||||
|
||||
|
||||
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=200,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
# saving and merging the model to local disk
|
||||
print("merge and save to local disk")
|
||||
model.save_pretrained_merged(
|
||||
save_directory='./unsloth_out/merged_mistral_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# print("cleaning")
|
||||
# del model
|
||||
# del tokenizer
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_mistral_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
|
||||
print("Computing 8-bit model perplexity in subprocess...")
|
||||
result_queue = mp.Queue()
|
||||
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
ppl_8bit = result_queue.get()
|
||||
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
||||
|
||||
print("Loading merged model in 16 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_mistral_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=False,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
print_model_comparison()
|
||||
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -0,0 +1,222 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [
|
||||
tokenizer.apply_chat_template(
|
||||
convo, tokenize = False, add_generation_prompt = False
|
||||
)
|
||||
for convo in convos
|
||||
]
|
||||
return { "text" : texts, }
|
||||
|
||||
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
"""Load model and compute perplexity in subprocess"""
|
||||
from unsloth import FastLanguageModel
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_phi4_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=load_in_4bit,
|
||||
load_in_8bit=load_in_8bit,
|
||||
)
|
||||
# Set up tokenizer
|
||||
merged_tokenizer = get_chat_template(
|
||||
merged_tokenizer,
|
||||
chat_template="phi-4",
|
||||
)
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
# Format the dataset
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [merged_tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Compute perplexity using the passed dataset
|
||||
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
||||
|
||||
|
||||
# IMPORTANT: Convert to Python float if it's a tensor
|
||||
if torch.is_tensor(ppl_value):
|
||||
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
||||
elif hasattr(ppl_value, 'item'):
|
||||
ppl_value = ppl_value.item() # Convert numpy or other array types
|
||||
else:
|
||||
ppl_value = float(ppl_value) # Ensure it's a float
|
||||
|
||||
# Return only the perplexity value
|
||||
result_queue.put(ppl_value)
|
||||
|
||||
# Clean up
|
||||
del merged_model
|
||||
del merged_tokenizer
|
||||
del dataset_ppl
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Phi-4",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template="phi-4",
|
||||
)
|
||||
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=200,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part="<|im_start|>user<|im_sep|>\n\n",
|
||||
response_part="<|im_start|>assistant<|im_sep|>\n\n",
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
# saving and merging the model to local disk
|
||||
print("merge and save to local disk")
|
||||
model.save_pretrained_merged(
|
||||
save_directory='./unsloth_out/merged_phi4_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# print("cleaning")
|
||||
# del model
|
||||
# del tokenizer
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_phi4_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
|
||||
print("Computing 8-bit model perplexity in subprocess...")
|
||||
result_queue = mp.Queue()
|
||||
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
ppl_8bit = result_queue.get()
|
||||
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
||||
|
||||
print("Loading merged model in 16 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_phi4_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=False,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
print_model_comparison()
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
"""Load model and compute perplexity in subprocess"""
|
||||
from unsloth import FastLanguageModel
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=load_in_4bit,
|
||||
load_in_8bit=load_in_8bit,
|
||||
)
|
||||
# Set up tokenizer
|
||||
merged_tokenizer = get_chat_template(
|
||||
merged_tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
# Format the dataset
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [merged_tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Compute perplexity using the passed dataset
|
||||
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
||||
|
||||
|
||||
# IMPORTANT: Convert to Python float if it's a tensor
|
||||
if torch.is_tensor(ppl_value):
|
||||
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
||||
elif hasattr(ppl_value, 'item'):
|
||||
ppl_value = ppl_value.item() # Convert numpy or other array types
|
||||
else:
|
||||
ppl_value = float(ppl_value) # Ensure it's a float
|
||||
|
||||
# Return only the perplexity value
|
||||
result_queue.put(ppl_value)
|
||||
|
||||
# Clean up
|
||||
del merged_model
|
||||
del merged_tokenizer
|
||||
del dataset_ppl
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Llama-3.1-8B-Instruct",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
|
||||
print("\n dataset sample [0]")
|
||||
print(dataset_train[0])
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=200,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
|
||||
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
)
|
||||
|
||||
tokenizer.decode(trainer.train_dataset[0]["input_ids"])
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
# saving and merging the model to local disk
|
||||
print("merge and save to local disk")
|
||||
model.save_pretrained_merged(
|
||||
save_directory='./unsloth_out/merged_llama_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# print("cleaning")
|
||||
# del model
|
||||
# del tokenizer
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
|
||||
print("Computing 8-bit model perplexity in subprocess...")
|
||||
result_queue = mp.Queue()
|
||||
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
ppl_8bit = result_queue.get()
|
||||
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
||||
|
||||
print("Loading merged model in 16 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_llama_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=False,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
print_model_comparison()
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -0,0 +1,287 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
|
||||
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
||||
|
||||
### Instruction:
|
||||
{}
|
||||
|
||||
### Input:
|
||||
{}
|
||||
|
||||
### Response:
|
||||
{}"""
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
instructions = []
|
||||
inputs = []
|
||||
outputs = []
|
||||
texts = []
|
||||
|
||||
for conversation in examples["messages"]:
|
||||
# Extract user message and assistant response
|
||||
user_message = ""
|
||||
assistant_message = ""
|
||||
|
||||
for turn in conversation:
|
||||
if turn["role"] == "user":
|
||||
user_message = turn["content"]
|
||||
elif turn["role"] == "assistant":
|
||||
assistant_message = turn["content"]
|
||||
|
||||
# Store intermediate format
|
||||
instruction = "Complete the statement"
|
||||
instructions.append(instruction)
|
||||
inputs.append(user_message)
|
||||
outputs.append(assistant_message)
|
||||
|
||||
# Create formatted text
|
||||
text = alpaca_prompt.format(instruction, user_message, assistant_message)
|
||||
texts.append(text)
|
||||
|
||||
return {
|
||||
"instruction": instructions,
|
||||
"input": inputs,
|
||||
"output": outputs,
|
||||
"text": texts
|
||||
}
|
||||
|
||||
|
||||
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
"""Load model and compute perplexity in subprocess"""
|
||||
from unsloth import FastLanguageModel
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_qwen_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=load_in_4bit,
|
||||
load_in_8bit=load_in_8bit,
|
||||
)
|
||||
# Set up tokenizer
|
||||
# merged_tokenizer = get_chat_template(
|
||||
# merged_tokenizer,
|
||||
# chat_template="llama-3.1",
|
||||
# )
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
||||
|
||||
### Instruction:
|
||||
{}
|
||||
|
||||
### Input:
|
||||
{}
|
||||
|
||||
### Response:
|
||||
{}"""
|
||||
|
||||
def formatting_prompts_func(examples):
|
||||
instructions = []
|
||||
inputs = []
|
||||
outputs = []
|
||||
texts = []
|
||||
|
||||
for conversation in examples["messages"]:
|
||||
# Extract user message and assistant response
|
||||
user_message = ""
|
||||
assistant_message = ""
|
||||
|
||||
for turn in conversation:
|
||||
if turn["role"] == "user":
|
||||
user_message = turn["content"]
|
||||
elif turn["role"] == "assistant":
|
||||
assistant_message = turn["content"]
|
||||
|
||||
# Store intermediate format
|
||||
instruction = "Complete the statement"
|
||||
instructions.append(instruction)
|
||||
inputs.append(user_message)
|
||||
outputs.append(assistant_message)
|
||||
|
||||
# Create formatted text
|
||||
text = alpaca_prompt.format(instruction, user_message, assistant_message)
|
||||
texts.append(text)
|
||||
|
||||
return {
|
||||
"instruction": instructions,
|
||||
"input": inputs,
|
||||
"output": outputs,
|
||||
"text": texts
|
||||
}
|
||||
|
||||
|
||||
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Compute perplexity using the passed dataset
|
||||
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
||||
|
||||
|
||||
# IMPORTANT: Convert to Python float if it's a tensor
|
||||
if torch.is_tensor(ppl_value):
|
||||
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
||||
elif hasattr(ppl_value, 'item'):
|
||||
ppl_value = ppl_value.item() # Convert numpy or other array types
|
||||
else:
|
||||
ppl_value = float(ppl_value) # Ensure it's a float
|
||||
|
||||
# Return only the perplexity value
|
||||
result_queue.put(ppl_value)
|
||||
|
||||
# Clean up
|
||||
# del merged_model
|
||||
# del merged_tokenizer
|
||||
# del dataset_ppl
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Qwen2.5-7B-Instruct",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=200,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
# saving and merging the model to local disk
|
||||
print("merge and save to local disk")
|
||||
model.save_pretrained_merged(
|
||||
save_directory='./unsloth_out/merged_qwen_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# print("cleaning")
|
||||
# del model
|
||||
# del tokenizer
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_qwen_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
|
||||
print("Computing 8-bit model perplexity in subprocess...")
|
||||
result_queue = mp.Queue()
|
||||
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
ppl_8bit = result_queue.get()
|
||||
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
||||
|
||||
print("Loading merged model in 16 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_qwen_text_model",
|
||||
max_seq_length=2048,
|
||||
load_in_4bit=False,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
|
||||
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
||||
|
||||
print_model_comparison()
|
||||
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
178
tests/saving/language_models/test_push_to_hub_merged.py
Normal file
178
tests/saving/language_models/test_push_to_hub_merged.py
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
import os
|
||||
from huggingface_hub import HfFileSystem, hf_hub_download
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Llama-3.2-1B-Instruct",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=30,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
|
||||
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
|
||||
# saving and merging the model to local disk
|
||||
hf_username = os.environ.get("HF_USER", "")
|
||||
if not hf_username:
|
||||
hf_username = input("Please enter your Hugging Face username: ").strip()
|
||||
os.environ["HF_USER"] = hf_username
|
||||
|
||||
hf_token = os.environ.get("HF_TOKEN", "")
|
||||
if not hf_token:
|
||||
hf_token = input("Please enter your Hugging Face token: ").strip()
|
||||
os.environ["HF_TOKEN"] = hf_token
|
||||
|
||||
|
||||
repo_name = f"{hf_username}/merged_llama_text_model"
|
||||
success = {
|
||||
"upload": False,
|
||||
"download": False,
|
||||
}
|
||||
|
||||
# Stage 1: Upload model to Hub
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== UPLOADING MODEL TO HUB ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token)
|
||||
success["upload"] = True
|
||||
print("✅ Model uploaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to upload model: {e}")
|
||||
raise Exception("Model upload failed.")
|
||||
|
||||
t
|
||||
# Stage 2: Test downloading the model (even if cached)
|
||||
safe_remove_directory(f"./{hf_username}")
|
||||
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TESTING MODEL DOWNLOAD ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# Force download even if cached
|
||||
model,tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
|
||||
success["download"] = True
|
||||
print("✅ Model downloaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Download failed: {e}")
|
||||
raise Exception("Model download failed.")
|
||||
|
||||
# Final report
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VALIDATION REPORT ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
for stage, passed in success.items():
|
||||
status = "✓" if passed else "✗"
|
||||
print(f"{status} {stage.replace('_', ' ').title()}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
if all(success.values()):
|
||||
print("\n🎉 All stages completed successfully!")
|
||||
else:
|
||||
raise Exception("Validation failed for one or more stages.")
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
|
@ -0,0 +1,197 @@
|
|||
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
||||
from datasets import load_dataset, Dataset
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Process, Queue
|
||||
import gc
|
||||
import os
|
||||
from huggingface_hub import HfFileSystem, hf_hub_download
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
|
||||
return {"text": texts}
|
||||
|
||||
|
||||
|
||||
if torch.cuda.is_bf16_supported():
|
||||
compute_dtype = torch.bfloat16
|
||||
attn_implementation = 'flash_attention_2'
|
||||
else:
|
||||
compute_dtype = torch.float16
|
||||
attn_implementation = 'sdpa'
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/Llama-3.1-8B-Instruct",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template="llama-3.1",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
||||
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
||||
|
||||
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=16,
|
||||
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=2048,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_ratio=0.1,
|
||||
max_steps=30,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=50,
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
|
||||
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
|
||||
# saving and merging the model to local disk
|
||||
hf_username = os.environ.get("HF_USER", "")
|
||||
if not hf_username:
|
||||
hf_username = input("Please enter your Hugging Face username: ").strip()
|
||||
os.environ["HF_USER"] = hf_username
|
||||
|
||||
hf_token = os.environ.get("HF_TOKEN", "")
|
||||
if not hf_token:
|
||||
hf_token = input("Please enter your Hugging Face token: ").strip()
|
||||
os.environ["HF_TOKEN"] = hf_token
|
||||
|
||||
|
||||
repo_name = f"{hf_username}/merged_llama_text_model"
|
||||
success = {
|
||||
"upload": False,
|
||||
"safetensors_check": False,
|
||||
"download": False,
|
||||
}
|
||||
|
||||
# Stage 1: Upload model to Hub
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== UPLOADING MODEL TO HUB ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token)
|
||||
success["upload"] = True
|
||||
print("✅ Model uploaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to upload model: {e}")
|
||||
raise Exception("Model upload failed.")
|
||||
|
||||
# Stage 2: Verify safetensors.index.json exists
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VERIFYING REPO CONTENTS ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
fs = HfFileSystem(token=hf_token)
|
||||
file_list = fs.ls(repo_name, detail=True)
|
||||
safetensors_found = any(
|
||||
file["name"].endswith("model.safetensors.index.json") for file in file_list
|
||||
)
|
||||
if safetensors_found:
|
||||
success["safetensors_check"] = True
|
||||
print("✅ model.safetensors.index.json found in repo!")
|
||||
else:
|
||||
raise Exception("model.safetensors.index.json not found in repo.")
|
||||
except Exception as e:
|
||||
print(f"❌ Verification failed: {e}")
|
||||
raise Exception("Repo verification failed.")
|
||||
|
||||
# Stage 3: Test downloading the model (even if cached)
|
||||
safe_remove_directory("./RTannous")
|
||||
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TESTING MODEL DOWNLOAD ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# Force download even if cached
|
||||
model,tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
|
||||
success["download"] = True
|
||||
print("✅ Model downloaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Download failed: {e}")
|
||||
raise Exception("Model download failed.")
|
||||
|
||||
# Final report
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VALIDATION REPORT ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
for stage, passed in success.items():
|
||||
status = "✓" if passed else "✗"
|
||||
print(f"{status} {stage.replace('_', ' ').title()}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
if all(success.values()):
|
||||
print("\n🎉 All stages completed successfully!")
|
||||
else:
|
||||
raise Exception("Validation failed for one or more stages.")
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
794
tests/saving/language_models/test_save_merged_grpo_model.py
Normal file
794
tests/saving/language_models/test_save_merged_grpo_model.py
Normal file
|
|
@ -0,0 +1,794 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
"""test_Llama3_1_(3B)_GRPO_LoRA (1).ipynb
|
||||
|
||||
### Unsloth
|
||||
|
||||
"""
|
||||
|
||||
from unsloth import FastLanguageModel
|
||||
import torch
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import multiprocessing as mp
|
||||
import gc
|
||||
from multiprocessing import Queue
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.aime_eval import evaluate_model_aime, compare_aime_results
|
||||
|
||||
|
||||
max_seq_length = 2048 # Can increase for longer reasoning traces
|
||||
lora_rank = 64 # Larger rank = smarter, but slower
|
||||
|
||||
|
||||
def evaluate_merged_model(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
from unsloth import FastLanguageModel
|
||||
from tests.utils.aime_eval import evaluate_model_aime
|
||||
max_seq_length = 2048 # Can increase for longer reasoning traces
|
||||
lora_rank = 64 # Larger rank = smarter, but slower
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "./final_merged_model",
|
||||
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.8, # Reduce if out of memory
|
||||
)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
if load_in_4bit:
|
||||
print("🔍 EVALUATION Merged model: 4 bits load")
|
||||
model_type="merged_model_4bits"
|
||||
elif load_in_8bit:
|
||||
print("🔍 EVALUATION Merged model: 8 bits load")
|
||||
model_type="merged_model_8bits"
|
||||
else:
|
||||
print("🔍 EVALUATION Merged model: 16 bits load")
|
||||
model_type="merged_model_16bits"
|
||||
print(f"{'='*60}")
|
||||
|
||||
evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type=model_type,
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768,
|
||||
top_p=0.95,
|
||||
seed=0
|
||||
)
|
||||
|
||||
result_queue.put(results)
|
||||
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
def training_run(result_queue):
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "meta-llama/Llama-3.2-3B-Instruct",
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False, # False for LoRA 16bit
|
||||
fast_inference = True, # Enable vLLM fast inference
|
||||
max_lora_rank = lora_rank,
|
||||
gpu_memory_utilization = 0.8, # Reduce if out of memory
|
||||
)
|
||||
|
||||
"""### Helper Functions
|
||||
<a name="Data"></a>
|
||||
|
||||
#### Helper functions - Data Prep
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
def extract_hash_answer(text):
|
||||
"""Extract answer from GSM8K format"""
|
||||
if "####" not in text:
|
||||
return None
|
||||
return text.split("####")[1].strip()
|
||||
|
||||
def prepare_gsm8k_dataset(dataset):
|
||||
"""Format GSM8K dataset for training"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
system_prompt = (
|
||||
f"You are given a problem. Think about the problem and reason step by step. "
|
||||
f"Place your thinking process between {reasoning_start} and {reasoning_end}. "
|
||||
f"Then, provide your final numerical solution between {solution_start}{solution_end}"
|
||||
)
|
||||
|
||||
def format_gsm8k(example):
|
||||
return {
|
||||
"prompt": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": example["question"]},
|
||||
],
|
||||
"answer": extract_hash_answer(example["answer"]),
|
||||
}
|
||||
|
||||
return dataset.map(format_gsm8k)
|
||||
|
||||
def prepare_limo_dataset(dataset):
|
||||
"""Format LIMO dataset for SFT training"""
|
||||
if dataset is None:
|
||||
return None
|
||||
|
||||
system_prompt = """You are a helpful reasoning assistant. When given a problem, think through it step by step and provide your answer in the following format:
|
||||
|
||||
<reasoning>
|
||||
[Your detailed step-by-step reasoning and solution process]
|
||||
</reasoning>
|
||||
<answer>
|
||||
[Your final numerical answer]
|
||||
</answer>"""
|
||||
|
||||
def format_limo(example):
|
||||
# Create the assistant response
|
||||
assistant_response = f"<reasoning>\n{example['solution']}\n</reasoning>\n<answer>\n{example['answer']}\n</answer>"
|
||||
|
||||
# Return a DICTIONARY with the conversation in a field
|
||||
return {
|
||||
"prompt": [ # ← This is the key change - wrap in a dict
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": example["question"]},
|
||||
{"role": "assistant", "content": assistant_response}
|
||||
]
|
||||
}
|
||||
|
||||
return dataset.map(format_limo)
|
||||
|
||||
print("\n✅ Dataset preparation functions defined!")
|
||||
|
||||
"""#### Helper functions - Evaluation"""
|
||||
|
||||
def get_max_prompt_length(dataset, tokenizer):
|
||||
"""Calculate maximum and average prompt length in dataset"""
|
||||
print("Analyzing prompt lengths...")
|
||||
|
||||
lengths = dataset.map(
|
||||
lambda x: {
|
||||
"tokens": tokenizer.apply_chat_template(
|
||||
x["prompt"],
|
||||
add_generation_prompt=True,
|
||||
tokenize=True
|
||||
)
|
||||
},
|
||||
batched=True,
|
||||
).map(lambda x: {"length": len(x["tokens"])})["length"]
|
||||
|
||||
max_length = max(lengths)
|
||||
avg_length = sum(lengths) / len(lengths)
|
||||
min_length = min(lengths)
|
||||
|
||||
print(f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}")
|
||||
return max_length, avg_length
|
||||
|
||||
def extract_unsloth_answer(text, start_tag="<SOLUTION>", end_tag="</SOLUTION>"):
|
||||
"""Extract answer from Unsloth SOLUTION tags"""
|
||||
pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
|
||||
matches = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
if matches:
|
||||
answer = matches[-1] # Get the last match
|
||||
answer = re.sub(r"[%$,]", "", answer).strip()
|
||||
return answer
|
||||
return ""
|
||||
|
||||
def find_number(search_string):
|
||||
"""Find the last number in a string"""
|
||||
numbers = re.compile(
|
||||
r"-?[\d,]*\.?\d+",
|
||||
re.MULTILINE | re.DOTALL | re.IGNORECASE,
|
||||
).findall(search_string)
|
||||
|
||||
if numbers:
|
||||
return numbers[-1].replace(",", "").strip()
|
||||
return ""
|
||||
|
||||
def remove_symbols(x: str) -> str:
|
||||
"""Remove commas, percent and dollar symbols"""
|
||||
if not x:
|
||||
return ""
|
||||
return x.replace(",", "").replace("%", "").replace("$", "").strip()
|
||||
|
||||
def get_num_tokens(text, tokenizer_instance):
|
||||
"""Count tokens in text"""
|
||||
if not text:
|
||||
return 0
|
||||
encoding = tokenizer_instance(text, return_tensors="pt")
|
||||
return len(encoding["input_ids"][0])
|
||||
|
||||
def check_format_compliance(text, format_type="unsloth"):
|
||||
"""Check if response follows expected format"""
|
||||
if format_type == "unsloth":
|
||||
reasoning_start = "<start_reasoning>"
|
||||
reasoning_end = "<end_reasoning>"
|
||||
solution_start = "<SOLUTION>"
|
||||
solution_end = "</SOLUTION>"
|
||||
|
||||
pattern = (
|
||||
rf"^[\s]*{re.escape(reasoning_start)}.+?{re.escape(reasoning_end)}.*?"
|
||||
rf"{re.escape(solution_start)}.+?{re.escape(solution_end)}[\s]*$"
|
||||
)
|
||||
else:
|
||||
return False
|
||||
|
||||
return bool(re.match(pattern, text.strip(), re.DOTALL))
|
||||
|
||||
def normalize_answer(answer):
|
||||
"""Normalize answer for comparison"""
|
||||
if not answer:
|
||||
return ""
|
||||
|
||||
normalized = remove_symbols(str(answer))
|
||||
|
||||
try:
|
||||
float_val = float(normalized)
|
||||
if float_val.is_integer():
|
||||
return str(int(float_val))
|
||||
else:
|
||||
return str(float_val)
|
||||
except (ValueError, TypeError):
|
||||
return normalized
|
||||
|
||||
def evaluate_answer_correctness(extracted_answer, ground_truth):
|
||||
"""Evaluate answer correctness with multiple criteria"""
|
||||
if not extracted_answer or not ground_truth:
|
||||
return False, False, 0.0
|
||||
|
||||
norm_extracted = normalize_answer(extracted_answer)
|
||||
norm_ground_truth = normalize_answer(ground_truth)
|
||||
|
||||
if norm_extracted == norm_ground_truth:
|
||||
return True, True, 1.0
|
||||
|
||||
try:
|
||||
extracted_num = float(norm_extracted)
|
||||
ground_truth_num = float(norm_ground_truth)
|
||||
|
||||
if ground_truth_num != 0:
|
||||
relative_error = abs(extracted_num - ground_truth_num) / abs(ground_truth_num)
|
||||
|
||||
if relative_error < 0.01:
|
||||
return True, True, 0.9
|
||||
elif relative_error < 0.05:
|
||||
return False, True, 0.7
|
||||
elif relative_error < 0.10:
|
||||
return False, True, 0.5
|
||||
else:
|
||||
if extracted_num == 0:
|
||||
return True, True, 1.0
|
||||
elif abs(extracted_num) < 0.01:
|
||||
return False, True, 0.7
|
||||
|
||||
except (ValueError, TypeError):
|
||||
if norm_extracted.lower() == norm_ground_truth.lower():
|
||||
return True, True, 1.0
|
||||
|
||||
return False, False, 0.0
|
||||
|
||||
"""#### Reward Functions for GRPO"""
|
||||
|
||||
def match_format_exactly(completions, **kwargs):
|
||||
"""Reward function for exact format matching"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
pattern = (
|
||||
rf"^[\s]*{re.escape(reasoning_start)}.+?{re.escape(reasoning_end)}.*?"
|
||||
rf"{re.escape(solution_start)}.+?{re.escape(solution_end)}[\s]*$"
|
||||
)
|
||||
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
rewards = [3.0 if re.match(pattern, response, re.DOTALL) else 0.0 for response in responses]
|
||||
return rewards
|
||||
|
||||
def match_format_approximately(completions, **kwargs):
|
||||
"""Reward function for approximate format matching"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answerr>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
scores = []
|
||||
for completion in completions:
|
||||
score = 0
|
||||
response = completion[0]["content"]
|
||||
score += 0.5 if response.count(reasoning_start) == 1 else -1.0
|
||||
score += 0.5 if response.count(reasoning_end) == 1 else -1.0
|
||||
score += 0.5 if response.count(solution_start) == 1 else -1.0
|
||||
score += 0.5 if response.count(solution_end) == 1 else -1.0
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
def check_answer_correctness(prompts, completions, answer, **kwargs):
|
||||
"""Reward function for answer correctness"""
|
||||
def extract_solution_answer(text):
|
||||
pattern = r"<answer>(.*?)</answer>"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
return re.sub(r"[%$,]", "", match.group(1)).strip()
|
||||
return ""
|
||||
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
extracted_responses = [extract_solution_answer(r) for r in responses]
|
||||
|
||||
scores = []
|
||||
for guess, true_answer in zip(extracted_responses, answer):
|
||||
score = 0
|
||||
if not guess:
|
||||
scores.append(0)
|
||||
continue
|
||||
|
||||
if guess == true_answer:
|
||||
score += 3.0
|
||||
elif guess.strip() == true_answer.strip():
|
||||
score += 1.5
|
||||
else:
|
||||
try:
|
||||
ratio = float(guess) / float(true_answer)
|
||||
if 0.9 <= ratio <= 1.1:
|
||||
score += 1.0
|
||||
elif 0.8 <= ratio <= 1.2:
|
||||
score += 0.5
|
||||
else:
|
||||
score -= 1.5
|
||||
except:
|
||||
score -= 1.5
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
print("✅ Reward functions defined!")
|
||||
|
||||
"""#### Main Evaluation Function"""
|
||||
|
||||
import gc
|
||||
|
||||
|
||||
|
||||
"""#### Comparison and Memory Management"""
|
||||
|
||||
def compare_model_results(all_results):
|
||||
"""Generate comprehensive comparison of multiple model results"""
|
||||
print(f"\n{'='*80}")
|
||||
print("COMPREHENSIVE MODEL COMPARISON")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Main table
|
||||
print(f"{'Model':<15} {'Format %':<10} {'Exact %':<10} {'Plausible %':<12} {'Confidence':<12}")
|
||||
print("-" * 80)
|
||||
|
||||
for result in all_results:
|
||||
print(f"{result['model_type']:<15} "
|
||||
f"{result['correct_format_pct']:<10.1f} "
|
||||
f"{result['exact_match_pct']:<10.1f} "
|
||||
f"{result['plausible_match_pct']:<12.1f} "
|
||||
f"{result['avg_confidence']:<12.3f}")
|
||||
|
||||
# Improvement analysis
|
||||
if len(all_results) > 1:
|
||||
print(f"\n{'='*50}")
|
||||
print("IMPROVEMENT ANALYSIS")
|
||||
print(f"{'='*50}")
|
||||
|
||||
base_result = all_results[0]
|
||||
for result in all_results[1:]:
|
||||
print(f"\n{result['model_type']} vs {base_result['model_type']}:")
|
||||
format_improvement = result['correct_format_pct'] - base_result['correct_format_pct']
|
||||
exact_improvement = result['exact_match_pct'] - base_result['exact_match_pct']
|
||||
plausible_improvement = result['plausible_match_pct'] - base_result['plausible_match_pct']
|
||||
|
||||
print(f" Format compliance: {format_improvement:+.1f}%")
|
||||
print(f" Exact matches: {exact_improvement:+.1f}%")
|
||||
print(f" Plausible matches: {plausible_improvement:+.1f}%")
|
||||
|
||||
# Save comparison
|
||||
comparison_data = {
|
||||
"summary": all_results,
|
||||
"best_model": max(all_results, key=lambda x: x['exact_match_pct']),
|
||||
}
|
||||
|
||||
with open("model_comparison_comprehensive.json", "w") as f:
|
||||
json.dump(comparison_data, f, indent=4)
|
||||
|
||||
print(f"\nBest performing model: {comparison_data['best_model']['model_type']} "
|
||||
f"({comparison_data['best_model']['exact_match_pct']:.1f}% exact matches)")
|
||||
|
||||
def cleanup_memory():
|
||||
"""Comprehensive memory cleanup"""
|
||||
print("🧹 Cleaning up GPU memory...")
|
||||
for _ in range(10):
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
allocated = torch.cuda.memory_allocated() / 1024**3
|
||||
reserved = torch.cuda.memory_reserved() / 1024**3
|
||||
print(f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
|
||||
|
||||
"""#### Data Loading and Preparation"""
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
|
||||
# Load GSM8K
|
||||
gsm8k_dataset = load_dataset("openai/gsm8k", "main", split="train")
|
||||
|
||||
# Load LIMO (adjust this based on your access method)
|
||||
limo_train = load_dataset("GAIR/LIMO", split="train")
|
||||
|
||||
# Prepare datasets
|
||||
gsm8k_train = prepare_gsm8k_dataset(gsm8k_dataset)
|
||||
limo_train = prepare_limo_dataset(limo_train)
|
||||
|
||||
|
||||
print(f" GSM8K train: {len(gsm8k_train)}")
|
||||
print(f" LIMO train: {len(limo_train) if limo_train else 0}")
|
||||
|
||||
# Store results
|
||||
all_results = []
|
||||
|
||||
# Single temperature evaluation on combined dataset
|
||||
results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="base",
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768,
|
||||
top_p=0.95,
|
||||
seed=0
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template = "llama-3.1",
|
||||
)
|
||||
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["prompt"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
|
||||
return { "text" : texts, }
|
||||
pass
|
||||
|
||||
limo_train = limo_train.map(formatting_prompts_func, batched = True,)
|
||||
|
||||
from trl import SFTTrainer
|
||||
from transformers import DataCollatorForSeq2Seq, TrainingArguments
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
|
||||
print(f"\n{'*'*60}")
|
||||
print("🎯 STAGE 1: Qlora Fine-Tuning on LIMO")
|
||||
print(f"{'*'*60}")
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
target_modules = [
|
||||
"q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",
|
||||
], # Remove QKVO if out of memory
|
||||
lora_alpha = lora_rank,
|
||||
use_gradient_checkpointing = "unsloth", # Enable long context finetuning
|
||||
random_state = 3407,
|
||||
)
|
||||
|
||||
|
||||
if limo_train is not None:
|
||||
trainer = SFTTrainer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
train_dataset = limo_train,
|
||||
dataset_text_field = "text",
|
||||
max_seq_length = max_seq_length,
|
||||
data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
|
||||
dataset_num_proc = 2,
|
||||
packing = False, # Can make training 5x faster for short sequences.
|
||||
args = TrainingArguments(
|
||||
per_device_train_batch_size = 2,
|
||||
gradient_accumulation_steps = 4,
|
||||
warmup_steps = 5,
|
||||
num_train_epochs = 1, # Set this for 1 full training run.
|
||||
#max_steps = 60,
|
||||
learning_rate = 2e-4,
|
||||
fp16 = not is_bfloat16_supported(),
|
||||
bf16 = is_bfloat16_supported(),
|
||||
logging_steps = 1,
|
||||
optim = "adamw_8bit",
|
||||
weight_decay = 0.01,
|
||||
lr_scheduler_type = "linear",
|
||||
seed = 3407,
|
||||
output_dir = "outputs",
|
||||
report_to = "none", # Use this for WandB etc
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part = "<|start_header_id|>user<|end_header_id|>\n\n",
|
||||
response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
|
||||
)
|
||||
|
||||
# Train
|
||||
print(f"🚂 Starting SFT training on {len(limo_train)} examples...")
|
||||
trainer.train()
|
||||
|
||||
# Save checkpoint
|
||||
model.save_pretrained("qlora_checkpoint")
|
||||
tokenizer.save_pretrained("qlora_checkpoint")
|
||||
print("💾 Qlora checkpoint saved!")
|
||||
|
||||
# Cleanup
|
||||
del trainer
|
||||
cleanup_memory()
|
||||
|
||||
print("✅ Qlora training completed!")
|
||||
else:
|
||||
print("⚠️ Skipping Qlora training - no LIMO dataset available")
|
||||
|
||||
# Cleanup
|
||||
cleanup_memory()
|
||||
|
||||
global PRINTED_TIMES
|
||||
PRINTED_TIMES = 0
|
||||
global PRINT_EVERY_STEPS
|
||||
PRINT_EVERY_STEPS = 5
|
||||
|
||||
match_numbers = re.compile(
|
||||
solution_start + r".*?([\d\.\,]{1,})",
|
||||
flags = re.MULTILINE | re.DOTALL
|
||||
)
|
||||
|
||||
def check_numbers(prompts, completions, answer, **kwargs):
|
||||
question = prompts[0][-1]["content"]
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
|
||||
extracted_responses = [
|
||||
guess.group(1)
|
||||
if (guess := match_numbers.search(r)) is not None else None \
|
||||
for r in responses
|
||||
]
|
||||
|
||||
scores = []
|
||||
# Print only every few steps
|
||||
global PRINTED_TIMES
|
||||
global PRINT_EVERY_STEPS
|
||||
if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
|
||||
print('*'*20, f"Question:\n{question}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
|
||||
PRINTED_TIMES += 1
|
||||
|
||||
for guess, true_answer in zip(extracted_responses, answer):
|
||||
if guess is None:
|
||||
scores.append(0)
|
||||
continue
|
||||
# Convert to numbers
|
||||
try:
|
||||
true_answer = float(true_answer.strip())
|
||||
# Remove commas like in 123,456
|
||||
guess = float(guess.strip().replace(",", ""))
|
||||
scores.append(1.5 if guess == true_answer else -0.5)
|
||||
except:
|
||||
scores.append(0)
|
||||
continue
|
||||
return scores
|
||||
|
||||
print(f"\n{'*'*60}")
|
||||
print("🎯 STAGE 2: GRPO Fine-Tuning on GSM8K")
|
||||
print(f"{'*'*60}")
|
||||
|
||||
# Get max prompt length
|
||||
max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
|
||||
max_prompt_length = min(max_prompt_length + 10, 512) # Add buffer, cap at 512
|
||||
|
||||
print(f"Using max_prompt_length: {max_prompt_length}")
|
||||
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
training_args = GRPOConfig(
|
||||
learning_rate = 5e-6,
|
||||
weight_decay = 0.1,
|
||||
warmup_ratio = 0.1,
|
||||
lr_scheduler_type = "cosine",
|
||||
optim = "adamw_torch_fused",
|
||||
logging_steps = 1,
|
||||
per_device_train_batch_size = 1,
|
||||
gradient_accumulation_steps = 4, # Increase to 4 for smoother training
|
||||
num_generations = 8, # Decrease if out of memory
|
||||
max_prompt_length = max_prompt_length,
|
||||
max_completion_length = max_seq_length - max_prompt_length,
|
||||
# num_train_epochs = 1, # Set to 1 for a full training run
|
||||
#max_steps = 250,
|
||||
max_steps = 1000,
|
||||
save_steps = 250,
|
||||
max_grad_norm = 0.1,
|
||||
report_to = "none", # Can use Weights & Biases
|
||||
output_dir = "outputs",
|
||||
)
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
processing_class = tokenizer,
|
||||
reward_funcs = [
|
||||
match_format_exactly,
|
||||
match_format_approximately,
|
||||
check_answer_correctness,
|
||||
check_numbers,
|
||||
],
|
||||
args = training_args,
|
||||
train_dataset = gsm8k_train,
|
||||
)
|
||||
|
||||
|
||||
# Train
|
||||
print(f"🚂 Starting GRPO training on {len(gsm8k_train)} examples...")
|
||||
trainer.train()
|
||||
|
||||
# Save checkpoint
|
||||
model.save_pretrained("grpo_checkpoint")
|
||||
tokenizer.save_pretrained("grpo_checkpoint")
|
||||
print("💾 GRPO checkpoint saved!")
|
||||
|
||||
# Cleanup
|
||||
del trainer
|
||||
del training_args
|
||||
cleanup_memory()
|
||||
|
||||
print("✅ GRPO training completed!")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("🔍 EVALUATION 3: Final GRPO Model")
|
||||
print(f"{'='*60}")
|
||||
|
||||
grpo_results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="grpo",
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768,
|
||||
top_p=0.95,
|
||||
seed=0
|
||||
)
|
||||
|
||||
all_results.append(grpo_results)
|
||||
print("✅ Final model evaluation complete!")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("💾 SAVING FINAL MODEL")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# Save as merged model
|
||||
try:
|
||||
model.save_pretrained_merged("final_merged_model", tokenizer, save_method="merged_16bit")
|
||||
print("✅ Merged model saved to: final_merged_model/")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not save merged model: {e}")
|
||||
print("Final model saved as LoRA adapter only")
|
||||
|
||||
print("💾 Model saving complete!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
result_queue.put(results)
|
||||
|
||||
# Clean up
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
# # Merged model load 16 bits model AIME eval
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, False))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
#
|
||||
# # Clean up
|
||||
# del merged_model
|
||||
# del merged_tokenizer
|
||||
# del dataset_ppl
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
#
|
||||
# safe_remove_directory("./unsloth_compiled_cache")
|
||||
#
|
||||
# # Merged model load 8 bits model AIME eval
|
||||
#
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, True))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
|
||||
|
||||
# Merged model load 4 bits AIME eval
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, True, False))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
result_queue = mp.Queue()
|
||||
all_results = []
|
||||
|
||||
|
||||
# run main finetuning and grpo loop
|
||||
p = mp.Process(target=training_run, args=(result_queue,))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
results = result_queue.get()
|
||||
all_results = results
|
||||
|
||||
# evaluate merged model loaded 16bits
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, False))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_16bits = result_queue.get()
|
||||
all_results.append(merged_load_16bits)
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
# Merged model load 8 bits model AIME eval
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_8bits = result_queue.get()
|
||||
all_results.append(merged_load_8bits)
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
# Merged model load 4 bits model AIME eval
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, True, False))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_4bits = result_queue.get()
|
||||
all_results.append(merged_load_4bits)
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
|
||||
# AIME-specific comparison function
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🏆 FINAL TRAINING PIPELINE RESULTS")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Use the AIME-specific comparison
|
||||
compare_aime_results(all_results)
|
||||
289
tests/saving/vision_models/test_index_file_sharded_model.py
Normal file
289
tests/saving/vision_models/test_index_file_sharded_model.py
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
## Import required libraries
|
||||
|
||||
from unsloth import FastVisionModel, is_bf16_supported
|
||||
from unsloth.trainer import UnslothVisionDataCollator
|
||||
|
||||
import torch
|
||||
import os
|
||||
from datasets import load_dataset
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from huggingface_hub import HfFileSystem
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
||||
## Dataset Preparation"""
|
||||
|
||||
print("\n📊 Loading and preparing dataset...")
|
||||
dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split="train")
|
||||
# To select the first 2000 examples
|
||||
train_dataset = dataset.select(range(2000))
|
||||
|
||||
# To select the next 200 examples for evaluation
|
||||
eval_dataset = dataset.select(range(2000, 2200))
|
||||
|
||||
print(f"✅ Dataset loaded successfully!")
|
||||
print(f" 📈 Training samples: {len(train_dataset)}")
|
||||
print(f" 📊 Evaluation samples: {len(eval_dataset)}")
|
||||
# Convert dataset to OAI messages
|
||||
def format_data(sample):
|
||||
return {
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": system_message}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": sample["question"],
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"image": sample["image"],
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": sample["answer"]}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
print("\n🔄 Formatting dataset for vision training...")
|
||||
system_message = "You are an expert french ocr system."
|
||||
# Convert dataset to OAI messages
|
||||
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
|
||||
train_dataset = [format_data(sample) for sample in train_dataset]
|
||||
eval_dataset = [format_data(sample) for sample in eval_dataset]
|
||||
print("✅ Dataset formatting completed!")
|
||||
|
||||
"""## Finetuning Setup and Run"""
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== MODEL LOADING AND SETUP ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# Load Base Model
|
||||
print("🤖 Loading base vision model...")
|
||||
try:
|
||||
model, tokenizer = FastVisionModel.from_pretrained(
|
||||
# model_name = "unsloth/Qwen2-VL-7B-Instruct",
|
||||
model_name="unsloth/Qwen2-VL-7B-Instruct",
|
||||
max_seq_length=2048, # Choose any for long context!
|
||||
load_in_4bit=True, # 4 bit quantization to reduce memory
|
||||
load_in_8bit=False, # [NEW!] A bit more accurate, uses 2x memory
|
||||
full_finetuning=False, # [NEW!] We have full finetuning now!
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to load base model: {e}")
|
||||
raise
|
||||
|
||||
print("\n🔧 Setting up LoRA configuration...")
|
||||
## Lora Finetuning
|
||||
try:
|
||||
model = FastVisionModel.get_peft_model(
|
||||
model,
|
||||
finetune_vision_layers=True, # Turn off for just text!
|
||||
finetune_language_layers=True, # Should leave on!
|
||||
finetune_attention_modules=True, # Attention good for GRPO
|
||||
finetune_mlp_modules=True, # SHould leave on always!
|
||||
r=16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
lora_alpha=32,
|
||||
lora_dropout=0, # Supports any, but = 0 is optimized
|
||||
bias="none", # Supports any, but = "none" is optimized
|
||||
use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
|
||||
random_state=3407,
|
||||
use_rslora=False, # We support rank stabilized LoRA
|
||||
loftq_config=None, # And LoftQ
|
||||
)
|
||||
print("✅ LoRA configuration applied successfully!")
|
||||
print(f" 🎯 LoRA rank (r): 16")
|
||||
print(f" 📊 LoRA alpha: 32")
|
||||
print(f" 🔍 Vision layers: Enabled")
|
||||
print(f" 💬 Language layers: Enabled")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to apply LoRA configuration: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TRAINING SETUP ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("🏋️ Preparing trainer...")
|
||||
FastVisionModel.for_training(model) # Enable for training!
|
||||
|
||||
try:
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
data_collator=UnslothVisionDataCollator(model, tokenizer),
|
||||
train_dataset=train_dataset,
|
||||
args=SFTConfig(
|
||||
# per_device_train_batch_size = 4,
|
||||
# gradient_accumulation_steps = 8,
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
gradient_checkpointing=True,
|
||||
gradient_checkpointing_kwargs={
|
||||
"use_reentrant": False
|
||||
}, # use reentrant checkpointing
|
||||
max_grad_norm=0.3, # max gradient norm based on QLoRA paper
|
||||
warmup_ratio=0.03,
|
||||
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
|
||||
max_steps=10,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bf16_supported(),
|
||||
bf16=is_bf16_supported(),
|
||||
logging_steps=5,
|
||||
save_strategy="epoch",
|
||||
optim="adamw_torch_fused",
|
||||
weight_decay=0.01,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="checkpoints",
|
||||
report_to="none", # For Weights and Biases
|
||||
# You MUST put the below items for vision finetuning:
|
||||
remove_unused_columns=False,
|
||||
dataset_text_field="",
|
||||
dataset_kwargs={"skip_prepare_dataset": True},
|
||||
dataset_num_proc=4,
|
||||
max_seq_length=2048,
|
||||
),
|
||||
)
|
||||
print("✅ Trainer setup completed!")
|
||||
print(f" 📦 Batch size: 2")
|
||||
print(f" 🔄 Gradient accumulation steps: 4")
|
||||
print(f" 📈 Max training steps: 10")
|
||||
print(f" 🎯 Learning rate: 2e-4")
|
||||
print(f" 💾 Precision: {'BF16' if is_bf16_supported() else 'FP16'}")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to setup trainer: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== STARTING TRAINING ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# run training
|
||||
try:
|
||||
print("🚀 Starting training process...")
|
||||
trainer_stats = trainer.train()
|
||||
except Exception as e:
|
||||
print(f"❌ Training failed: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== SAVING MODEL ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
print("💾 Saving adapter model and tokenizer locally...")
|
||||
try:
|
||||
model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
|
||||
tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
print("✅ Model saved locally!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to save model locally: {e}")
|
||||
raise
|
||||
|
||||
|
||||
hf_username = os.environ.get("HF_USER", "")
|
||||
if not hf_username:
|
||||
hf_username = input("Please enter your Hugging Face username: ").strip()
|
||||
os.environ["HF_USER"] = hf_username
|
||||
|
||||
hf_token = os.environ.get("HF_TOKEN", "")
|
||||
if not hf_token:
|
||||
hf_token = input("Please enter your Hugging Face token: ").strip()
|
||||
os.environ["HF_TOKEN"] = hf_token
|
||||
|
||||
repo_name = f"{hf_username}/qwen2-7b-ocr-merged"
|
||||
success = {
|
||||
"upload": False,
|
||||
"safetensors_check": False,
|
||||
"download": False,
|
||||
}
|
||||
# Stage 1: Upload model to Hub
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== UPLOADING MODEL TO HUB ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
print(f"🚀 Uploading to repository: {repo_name}")
|
||||
model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token)
|
||||
success["upload"] = True
|
||||
print("✅ Model uploaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to upload model: {e}")
|
||||
raise Exception("Model upload failed.")
|
||||
|
||||
# Stage 2: Verify safetensors.index.json exists
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VERIFYING REPO CONTENTS ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
fs = HfFileSystem(token=hf_token)
|
||||
file_list = fs.ls(repo_name, detail=True)
|
||||
safetensors_found = any(
|
||||
file["name"].endswith("model.safetensors.index.json") for file in file_list
|
||||
)
|
||||
if safetensors_found:
|
||||
success["safetensors_check"] = True
|
||||
print("✅ model.safetensors.index.json found in repo!")
|
||||
else:
|
||||
raise Exception("model.safetensors.index.json not found in repo.")
|
||||
except Exception as e:
|
||||
print(f"❌ Verification failed: {e}")
|
||||
raise Exception("Repo verification failed.")
|
||||
|
||||
# test downloading model even if cached
|
||||
safe_remove_directory(f"./{hf_username}")
|
||||
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TESTING MODEL DOWNLOAD ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
print("📥 Testing model download...")
|
||||
# Force download even if cached
|
||||
test_model, test_tokenizer = FastVisionModel.from_pretrained(repo_name)
|
||||
success["download"] = True
|
||||
print("✅ Model downloaded successfully!")
|
||||
|
||||
# Clean up test model
|
||||
del test_model, test_tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as e:
|
||||
print(f"❌ Download failed: {e}")
|
||||
raise Exception("Model download failed.")
|
||||
|
||||
# Final report
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VALIDATION REPORT ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
for stage, passed in success.items():
|
||||
status = "✅" if passed else "❌"
|
||||
print(f"{status} {stage.replace('_', ' ').title()}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
if all(success.values()):
|
||||
print("\n🎉 All stages completed successfully!")
|
||||
print(f"🌐 Your model is available at: https://huggingface.co/{repo_name}")
|
||||
else:
|
||||
raise Exception("Validation failed for one or more stages.")
|
||||
|
||||
|
||||
# Final cleanup
|
||||
print("\n🧹 Cleaning up temporary files...")
|
||||
safe_remove_directory("./checkpoints")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
|
||||
print("\n🎯 Pipeline completed successfully!")
|
||||
print("=" * 80)
|
||||
268
tests/saving/vision_models/test_push_to_hub_merged.py
Normal file
268
tests/saving/vision_models/test_push_to_hub_merged.py
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
## Import required libraries
|
||||
|
||||
from unsloth import FastVisionModel, is_bf16_supported
|
||||
from unsloth.trainer import UnslothVisionDataCollator
|
||||
|
||||
import torch
|
||||
import os
|
||||
from datasets import load_dataset
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
||||
## Dataset Preparation"""
|
||||
|
||||
print("\n📊 Loading and preparing dataset...")
|
||||
dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split="train")
|
||||
# To select the first 2000 examples
|
||||
train_dataset = dataset.select(range(2000))
|
||||
|
||||
# To select the next 200 examples for evaluation
|
||||
eval_dataset = dataset.select(range(2000, 2200))
|
||||
|
||||
print(f"✅ Dataset loaded successfully!")
|
||||
print(f" 📈 Training samples: {len(train_dataset)}")
|
||||
print(f" 📊 Evaluation samples: {len(eval_dataset)}")
|
||||
# Convert dataset to OAI messages
|
||||
def format_data(sample):
|
||||
return {
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": system_message}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": sample["question"],
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"image": sample["image"],
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": sample["answer"]}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
print("\n🔄 Formatting dataset for vision training...")
|
||||
system_message = "You are an expert french ocr system."
|
||||
# Convert dataset to OAI messages
|
||||
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
|
||||
train_dataset = [format_data(sample) for sample in train_dataset]
|
||||
eval_dataset = [format_data(sample) for sample in eval_dataset]
|
||||
print("✅ Dataset formatting completed!")
|
||||
|
||||
"""## Finetuning Setup and Run"""
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== MODEL LOADING AND SETUP ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# Load Base Model
|
||||
print("🤖 Loading base vision model...")
|
||||
try:
|
||||
model, tokenizer = FastVisionModel.from_pretrained(
|
||||
# model_name = "unsloth/Qwen2-VL-7B-Instruct",
|
||||
model_name="unsloth/Qwen2-VL-2B-Instruct",
|
||||
max_seq_length=2048, # Choose any for long context!
|
||||
load_in_4bit=True, # 4 bit quantization to reduce memory
|
||||
load_in_8bit=False, # [NEW!] A bit more accurate, uses 2x memory
|
||||
full_finetuning=False, # [NEW!] We have full finetuning now!
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to load base model: {e}")
|
||||
raise
|
||||
|
||||
print("\n🔧 Setting up LoRA configuration...")
|
||||
## Lora Finetuning
|
||||
try:
|
||||
model = FastVisionModel.get_peft_model(
|
||||
model,
|
||||
finetune_vision_layers=True, # Turn off for just text!
|
||||
finetune_language_layers=True, # Should leave on!
|
||||
finetune_attention_modules=True, # Attention good for GRPO
|
||||
finetune_mlp_modules=True, # SHould leave on always!
|
||||
r=16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
lora_alpha=32,
|
||||
lora_dropout=0, # Supports any, but = 0 is optimized
|
||||
bias="none", # Supports any, but = "none" is optimized
|
||||
use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
|
||||
random_state=3407,
|
||||
use_rslora=False, # We support rank stabilized LoRA
|
||||
loftq_config=None, # And LoftQ
|
||||
)
|
||||
print("✅ LoRA configuration applied successfully!")
|
||||
print(f" 🎯 LoRA rank (r): 16")
|
||||
print(f" 📊 LoRA alpha: 32")
|
||||
print(f" 🔍 Vision layers: Enabled")
|
||||
print(f" 💬 Language layers: Enabled")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to apply LoRA configuration: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TRAINING SETUP ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("🏋️ Preparing trainer...")
|
||||
FastVisionModel.for_training(model) # Enable for training!
|
||||
|
||||
try:
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
data_collator=UnslothVisionDataCollator(model, tokenizer),
|
||||
train_dataset=train_dataset,
|
||||
args=SFTConfig(
|
||||
# per_device_train_batch_size = 4,
|
||||
# gradient_accumulation_steps = 8,
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
gradient_checkpointing=True,
|
||||
gradient_checkpointing_kwargs={
|
||||
"use_reentrant": False
|
||||
}, # use reentrant checkpointing
|
||||
max_grad_norm=0.3, # max gradient norm based on QLoRA paper
|
||||
warmup_ratio=0.03,
|
||||
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
|
||||
max_steps=10,
|
||||
learning_rate=2e-4,
|
||||
fp16=not is_bf16_supported(),
|
||||
bf16=is_bf16_supported(),
|
||||
logging_steps=5,
|
||||
save_strategy="epoch",
|
||||
optim="adamw_torch_fused",
|
||||
weight_decay=0.01,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="checkpoints",
|
||||
report_to="none", # For Weights and Biases
|
||||
# You MUST put the below items for vision finetuning:
|
||||
remove_unused_columns=False,
|
||||
dataset_text_field="",
|
||||
dataset_kwargs={"skip_prepare_dataset": True},
|
||||
dataset_num_proc=4,
|
||||
max_seq_length=2048,
|
||||
),
|
||||
)
|
||||
print("✅ Trainer setup completed!")
|
||||
print(f" 📦 Batch size: 2")
|
||||
print(f" 🔄 Gradient accumulation steps: 4")
|
||||
print(f" 📈 Max training steps: 10")
|
||||
print(f" 🎯 Learning rate: 2e-4")
|
||||
print(f" 💾 Precision: {'BF16' if is_bf16_supported() else 'FP16'}")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to setup trainer: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== STARTING TRAINING ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
# run training
|
||||
try:
|
||||
print("🚀 Starting training process...")
|
||||
trainer_stats = trainer.train()
|
||||
except Exception as e:
|
||||
print(f"❌ Training failed: {e}")
|
||||
raise
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("=== SAVING MODEL ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
print("💾 Saving adapter model and tokenizer locally...")
|
||||
try:
|
||||
model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
|
||||
tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
print("✅ Model saved locally!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to save model locally: {e}")
|
||||
raise
|
||||
|
||||
|
||||
hf_username = os.environ.get("HF_USER", "")
|
||||
if not hf_username:
|
||||
hf_username = input("Please enter your Hugging Face username: ").strip()
|
||||
os.environ["HF_USER"] = hf_username
|
||||
|
||||
hf_token = os.environ.get("HF_TOKEN", "")
|
||||
if not hf_token:
|
||||
hf_token = input("Please enter your Hugging Face token: ").strip()
|
||||
os.environ["HF_TOKEN"] = hf_token
|
||||
|
||||
repo_name = f"{hf_username}/qwen2-ocr-merged"
|
||||
success = {
|
||||
"upload": False,
|
||||
"download": False,
|
||||
}
|
||||
# Stage 1: Upload model to Hub
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== UPLOADING MODEL TO HUB ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
print(f"🚀 Uploading to repository: {repo_name}")
|
||||
model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token)
|
||||
success["upload"] = True
|
||||
print("✅ Model uploaded successfully!")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to upload model: {e}")
|
||||
raise Exception("Model upload failed.")
|
||||
|
||||
|
||||
try:
|
||||
print("\n" + "=" * 80)
|
||||
print("=== TESTING MODEL DOWNLOAD ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
print("📥 Testing model download...")
|
||||
# Force download even if cached
|
||||
test_model, test_tokenizer = FastVisionModel.from_pretrained(repo_name)
|
||||
success["download"] = True
|
||||
print("✅ Model downloaded successfully!")
|
||||
|
||||
# Clean up test model
|
||||
del test_model, test_tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as e:
|
||||
print(f"❌ Download failed: {e}")
|
||||
raise Exception("Model download failed.")
|
||||
|
||||
# Final report
|
||||
print("\n" + "=" * 80)
|
||||
print("=== VALIDATION REPORT ===".center(80))
|
||||
print("=" * 80 + "\n")
|
||||
for stage, passed in success.items():
|
||||
status = "✅" if passed else "❌"
|
||||
print(f"{status} {stage.replace('_', ' ').title()}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
if all(success.values()):
|
||||
print("\n🎉 All stages completed successfully!")
|
||||
print(f"🌐 Your model is available at: https://huggingface.co/{repo_name}")
|
||||
else:
|
||||
raise Exception("Validation failed for one or more stages.")
|
||||
|
||||
|
||||
# Final cleanup
|
||||
print("\n🧹 Cleaning up temporary files...")
|
||||
safe_remove_directory("./checkpoints")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
safe_remove_directory(f"./{hf_username}")
|
||||
|
||||
print("\n🎯 Pipeline completed successfully!")
|
||||
print("=" * 80)
|
||||
|
|
@ -0,0 +1,254 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
|
||||
from unsloth import FastVisionModel
|
||||
|
||||
import torch
|
||||
from qwen_vl_utils import process_vision_info
|
||||
import os
|
||||
from datasets import load_dataset
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.ocr_eval import OCRModelEvaluator
|
||||
|
||||
|
||||
## Dataset Preparation
|
||||
from datasets import load_dataset
|
||||
|
||||
dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", 'en', split="train")
|
||||
# To select the first 2000 examples
|
||||
train_dataset = dataset.select(range(2000))
|
||||
|
||||
# To select the next 200 examples for evaluation
|
||||
eval_dataset = dataset.select(range(2000, 2200))
|
||||
|
||||
# Convert dataset to OAI messages
|
||||
def format_data(sample):
|
||||
return {"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": system_message}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": sample["question"],
|
||||
},{
|
||||
"type": "image",
|
||||
"image": sample["image"],
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": sample["answer"]}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
system_message = "You are an expert french ocr system."
|
||||
# Convert dataset to OAI messages
|
||||
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
|
||||
train_dataset = [format_data(sample) for sample in train_dataset]
|
||||
eval_dataset = [format_data(sample) for sample in eval_dataset]
|
||||
|
||||
## Setup OCR main evaluation function and helpers
|
||||
import os
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
from jiwer import wer, cer
|
||||
from qwen_vl_utils import process_vision_info
|
||||
|
||||
#
|
||||
ocr_evaluator = OCRModelEvaluator()
|
||||
model_comparison_results = {}
|
||||
|
||||
## Finetuning Setup and Run
|
||||
# Load Base Model
|
||||
|
||||
model, tokenizer = FastVisionModel.from_pretrained(
|
||||
model_name = "unsloth/Qwen2-VL-7B-Instruct",
|
||||
max_seq_length = 2048, # Choose any for long context!
|
||||
load_in_4bit = True, # 4 bit quantization to reduce memory
|
||||
load_in_8bit = False, # [NEW!] A bit more accurate, uses 2x memory
|
||||
full_finetuning = False, # [NEW!] We have full finetuning now!
|
||||
)
|
||||
|
||||
# benchmark base model performance
|
||||
model_name = "Unsloth Base model"
|
||||
FastVisionModel.for_inference(model)
|
||||
avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_base_model_results")
|
||||
ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
## Lora Finetuning
|
||||
model = FastVisionModel.get_peft_model(
|
||||
model,
|
||||
finetune_vision_layers = True, # Turn off for just text!
|
||||
finetune_language_layers = True, # Should leave on!
|
||||
finetune_attention_modules = True, # Attention good for GRPO
|
||||
finetune_mlp_modules = True, # SHould leave on always!
|
||||
|
||||
r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
#target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
#"gate_proj", "up_proj", "down_proj",],
|
||||
lora_alpha = 32,
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
||||
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
|
||||
random_state = 3407,
|
||||
use_rslora = False, # We support rank stabilized LoRA
|
||||
loftq_config = None, # And LoftQ
|
||||
)
|
||||
|
||||
from unsloth import is_bf16_supported
|
||||
from unsloth.trainer import UnslothVisionDataCollator
|
||||
FastVisionModel.for_training(model) # Enable for training!
|
||||
model.config.use_cache = False
|
||||
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
data_collator = UnslothVisionDataCollator(model, tokenizer),
|
||||
train_dataset = train_dataset,
|
||||
args = SFTConfig(
|
||||
#per_device_train_batch_size = 4,
|
||||
#gradient_accumulation_steps = 8,
|
||||
per_device_train_batch_size = 2,
|
||||
gradient_accumulation_steps = 4,
|
||||
gradient_checkpointing=True,
|
||||
gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
|
||||
max_grad_norm=0.3, # max gradient norm based on QLoRA paper
|
||||
warmup_ratio=0.03,
|
||||
#num_train_epochs = 2, # Set this instead of max_steps for full training runs
|
||||
max_steps=60,
|
||||
learning_rate = 2e-4,
|
||||
fp16 = not is_bf16_supported(),
|
||||
bf16 = is_bf16_supported(),
|
||||
logging_steps = 5,
|
||||
save_strategy="epoch",
|
||||
optim = "adamw_torch_fused",
|
||||
weight_decay = 0.01,
|
||||
lr_scheduler_type = "linear",
|
||||
seed = 3407,
|
||||
output_dir = "unsloth-qwen2-7vl-french-ocr-checkpoints",
|
||||
report_to = "none", # For Weights and Biases
|
||||
|
||||
# You MUST put the below items for vision finetuning:
|
||||
remove_unused_columns = False,
|
||||
dataset_text_field = "",
|
||||
dataset_kwargs = {"skip_prepare_dataset": True},
|
||||
dataset_num_proc = 4,
|
||||
max_seq_length = 2048,
|
||||
),
|
||||
)
|
||||
|
||||
# run training
|
||||
trainer_stats = trainer.train()
|
||||
|
||||
model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
|
||||
tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
|
||||
## Measure Adapter Performance
|
||||
|
||||
# benchmark lora model performance
|
||||
model_name = "Unsloth lora adapter model"
|
||||
FastVisionModel.for_inference(model)
|
||||
avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_lora_model_results")
|
||||
ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
## Merge Model
|
||||
|
||||
def find_lora_base_model(model_to_inspect):
|
||||
current = model_to_inspect
|
||||
if hasattr(current, "base_model"):
|
||||
current = current.base_model
|
||||
if hasattr(current, "model"):
|
||||
current = current.model
|
||||
return current
|
||||
pass
|
||||
|
||||
base = find_lora_base_model(model)
|
||||
|
||||
print((base.__class__.__name__))
|
||||
|
||||
# merge default 16 bits
|
||||
model.save_pretrained_merged(save_directory="qwen2-ocr-merged-finetune-merge-16bit", tokenizer=tokenizer)
|
||||
|
||||
|
||||
## Benchmark merged model performance
|
||||
|
||||
### 16 bits merged model
|
||||
|
||||
model, tokenizer = FastVisionModel.from_pretrained("./qwen2-ocr-merged-finetune-merge-16bit",load_in_4bit=False, load_in_8bit=False)
|
||||
|
||||
# benchmark 4bit loaded, 16bits merged model performance
|
||||
model_name = "Unsloth 16bits-merged model load-16bits"
|
||||
model.config.use_cache = True
|
||||
|
||||
avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_16bits_merged_model_load_16bits_results")
|
||||
ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
# load 16bits-merged model in 4 bits
|
||||
model, tokenizer = FastVisionModel.from_pretrained("./qwen2-ocr-merged-finetune-merge-16bit",load_in_4bit=True, load_in_8bit=False)
|
||||
|
||||
# benchmark 4bit loaded, 16bits merged model performance
|
||||
model_name = "Unsloth 16bits-merged model load-4bits"
|
||||
model.config.use_cache = True
|
||||
|
||||
avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_16bits_merged_model_load_4bits_results")
|
||||
ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
# load model in 8 bits
|
||||
model, tokenizer = FastVisionModel.from_pretrained("./qwen2-ocr-merged-finetune-merge-16bit",load_in_4bit=False, load_in_8bit=True)
|
||||
|
||||
# benchmark 4bit loaded, 16bits merged model performance
|
||||
model_name = "Unsloth 16bits-merged model load-8bits"
|
||||
avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_16bits_merged_model_load_8bits_results")
|
||||
ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
# """### 4 bits merged model"""
|
||||
#
|
||||
# # load 4bits-merged model in 4 bits
|
||||
# model, tokenizer = FastVisionModel.from_pretrained("./qwen2-ocr-merged-finetune-merge-4bit",load_in_4bit=True, load_in_8bit=False)
|
||||
#
|
||||
# # benchmark 4bit loaded, 4bits merged model performance
|
||||
# model_name = "Unsloth 4bits-merged model load-4bits"
|
||||
#
|
||||
# avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_4bits_merged_model_load_4bits_results")
|
||||
# ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
#
|
||||
# # load model in 8 bits
|
||||
# model, tokenizer = FastVisionModel.from_pretrained("./qwen2-ocr-merged-finetune-merge-4bit",load_in_4bit=False, load_in_8bit=True)
|
||||
#
|
||||
# # benchmark 8bit loaded, 4bits merged model performance
|
||||
# model_name = "Unsloth 4bits-merged model load-8bits"
|
||||
#
|
||||
# avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_4bits_merged_model_load_8bits_results")
|
||||
# ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
|
||||
|
||||
# Model comparison report
|
||||
#print model comparison
|
||||
ocr_evaluator.print_model_comparison()
|
||||
|
||||
|
||||
|
||||
# Final cleanup
|
||||
print("\n🧹 Cleaning up temporary files...")
|
||||
safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-adapter")
|
||||
safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-checkpoints")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./qwen2-ocr-merged-finetune-merge-16bit")
|
||||
|
||||
print("\n🎯 Pipeline completed successfully!")
|
||||
print("=" * 80)
|
||||
264
tests/utils/aime_eval.md
Normal file
264
tests/utils/aime_eval.md
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
# AIME Dataset Evaluator
|
||||
|
||||
A Python module for evaluating language models on the AIME (American Invitational Mathematics Examination) dataset. This evaluator automatically downloads and combines multiple AIME test datasets and provides comprehensive mathematical reasoning assessment.
|
||||
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
from aime_utils import evaluate_model_aime
|
||||
|
||||
# Simple AIME evaluation
|
||||
results = evaluate_model_aime(
|
||||
model=your_model,
|
||||
tokenizer=your_tokenizer,
|
||||
model_type="base_model",
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768
|
||||
)
|
||||
|
||||
print(f"AIME Accuracy: {results['accuracy']:.1f}%")
|
||||
print(f"Pass@8: {results['pass_at_k']:.1f}%")
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
```python
|
||||
from aime_utils import evaluate_model_aime, compare_aime_results
|
||||
|
||||
# Evaluate multiple model configurations
|
||||
all_results = []
|
||||
|
||||
# Base model
|
||||
base_results = evaluate_model_aime(
|
||||
model=base_model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="base",
|
||||
temperature=0.3,
|
||||
n_sampling=8
|
||||
)
|
||||
all_results.append(base_results)
|
||||
|
||||
# Fine-tuned model
|
||||
ft_results = evaluate_model_aime(
|
||||
model=finetuned_model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="finetuned",
|
||||
temperature=0.3,
|
||||
n_sampling=8
|
||||
)
|
||||
all_results.append(ft_results)
|
||||
|
||||
# Generate comprehensive comparison
|
||||
compare_aime_results(all_results)
|
||||
```
|
||||
|
||||
## Dataset Format
|
||||
|
||||
The evaluator automatically handles AIME dataset format with problems containing:
|
||||
|
||||
- **Problem**: Mathematical question text
|
||||
- **Answer**: Numerical answer (0-999 range for AIME)
|
||||
- **Solution**: Step-by-step solution (when available)
|
||||
- **Source**: Original dataset identifier (test2024, test2025-I, test2025-II)
|
||||
|
||||
```python
|
||||
# Automatic dataset download and formatting
|
||||
{
|
||||
"global_id": 0,
|
||||
"original_id": "problem_1",
|
||||
"source_dataset": "test2024",
|
||||
"problem": "Find the number of...",
|
||||
"answer": "123",
|
||||
"solution": "Step-by-step solution...",
|
||||
"prompt": [
|
||||
{"role": "system", "content": "You are a mathematical problem solver..."},
|
||||
{"role": "user", "content": "Problem: Find the number of..."}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Configuration Examples
|
||||
|
||||
### Conservative Evaluation
|
||||
```python
|
||||
# Lower temperature for more consistent answers
|
||||
results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="conservative",
|
||||
temperature=0.1,
|
||||
n_sampling=4,
|
||||
top_p=0.9
|
||||
)
|
||||
```
|
||||
|
||||
### High-Sample Evaluation
|
||||
```python
|
||||
# More samples for better Pass@K estimation
|
||||
results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="high_sample",
|
||||
temperature=0.5,
|
||||
n_sampling=16,
|
||||
max_tokens=16384
|
||||
)
|
||||
```
|
||||
|
||||
### Memory-Optimized
|
||||
```python
|
||||
# Reduced parameters for limited resources
|
||||
results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="lite",
|
||||
temperature=0.3,
|
||||
n_sampling=4,
|
||||
max_tokens=8192
|
||||
)
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### Complete Model Pipeline Evaluation
|
||||
```python
|
||||
from aime_utils import evaluate_model_aime, compare_aime_results
|
||||
|
||||
def evaluate_training_pipeline(base_model, finetuned_model, merged_model, tokenizer):
|
||||
"""Evaluate complete training pipeline on AIME"""
|
||||
|
||||
all_results = []
|
||||
|
||||
# Standard evaluation configuration
|
||||
eval_config = {
|
||||
"temperature": 0.3,
|
||||
"n_sampling": 8,
|
||||
"max_tokens": 32768,
|
||||
"top_p": 0.95,
|
||||
"seed": 0
|
||||
}
|
||||
|
||||
# Evaluate base model
|
||||
print("Evaluating base model...")
|
||||
base_results = evaluate_model_aime(
|
||||
model=base_model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="base",
|
||||
**eval_config
|
||||
)
|
||||
all_results.append(base_results)
|
||||
|
||||
# Evaluate fine-tuned model
|
||||
print("Evaluating fine-tuned model...")
|
||||
ft_results = evaluate_model_aime(
|
||||
model=finetuned_model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="finetuned",
|
||||
**eval_config
|
||||
)
|
||||
all_results.append(ft_results)
|
||||
|
||||
# Evaluate merged model
|
||||
print("Evaluating merged model...")
|
||||
merged_results = evaluate_model_aime(
|
||||
model=merged_model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="merged",
|
||||
**eval_config
|
||||
)
|
||||
all_results.append(merged_results)
|
||||
|
||||
# Generate comparison report
|
||||
compare_aime_results(all_results)
|
||||
|
||||
return all_results
|
||||
```
|
||||
|
||||
### Quantization Impact Analysis
|
||||
```python
|
||||
def analyze_quantization_impact(model_paths, tokenizer):
|
||||
"""Analyze impact of different quantization levels"""
|
||||
|
||||
quantization_configs = {
|
||||
"fp16": {"load_in_4bit": False, "load_in_8bit": False},
|
||||
"8bit": {"load_in_4bit": False, "load_in_8bit": True},
|
||||
"4bit": {"load_in_4bit": True, "load_in_8bit": False}
|
||||
}
|
||||
|
||||
all_results = []
|
||||
|
||||
for quant_name, load_config in quantization_configs.items():
|
||||
print(f"Evaluating {quant_name} quantization...")
|
||||
|
||||
# Load model with specific quantization
|
||||
model = load_model_with_config(model_paths["merged"], **load_config)
|
||||
|
||||
results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type=f"merged_{quant_name}",
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768
|
||||
)
|
||||
all_results.append(results)
|
||||
|
||||
# Cleanup
|
||||
del model
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
compare_aime_results(all_results)
|
||||
return all_results
|
||||
```
|
||||
|
||||
## Output Format
|
||||
|
||||
### Individual Evaluation Results
|
||||
```
|
||||
🧮 AIME EVALUATION - BASE MODEL
|
||||
Combined Dataset: test2024 + test2025-I + test2025-II
|
||||
====================================================================
|
||||
|
||||
🎯 Overall Performance:
|
||||
Total problems: 45
|
||||
Correct answers: 12/45 (26.7%)
|
||||
Pass@8: 31.1%
|
||||
|
||||
📈 Performance by Dataset:
|
||||
test2024: 4/15 (26.7%)
|
||||
test2025-I: 5/15 (33.3%)
|
||||
test2025-II: 3/15 (20.0%)
|
||||
|
||||
🎖️ AIME Performance: ✅ EXCELLENT (26.7%)
|
||||
```
|
||||
|
||||
### Comparison Report
|
||||
```
|
||||
COMPREHENSIVE AIME MODEL COMPARISON
|
||||
================================================================================
|
||||
Model Accuracy % Pass@K % Correct Total
|
||||
--------------------------------------------------------------------------------
|
||||
finetuned 31.1 35.6 14 45
|
||||
base 26.7 31.1 12 45
|
||||
merged_4bit 24.4 28.9 11 45
|
||||
|
||||
IMPROVEMENT ANALYSIS
|
||||
==================================================
|
||||
finetuned vs base:
|
||||
Accuracy improvement: +4.4%
|
||||
Pass@K improvement: +4.5%
|
||||
```
|
||||
|
||||
## Performance Tiers
|
||||
|
||||
The evaluator provides performance assessment based on AIME difficulty:
|
||||
|
||||
- **🏆 EXCEPTIONAL**: ≥50% accuracy
|
||||
- **✅ EXCELLENT**: ≥30% accuracy
|
||||
- **🎯 VERY GOOD**: ≥20% accuracy
|
||||
- **⚠️ GOOD**: ≥10% accuracy
|
||||
- **📈 FAIR**: ≥5% accuracy
|
||||
- **❌ NEEDS IMPROVEMENT**: <5% accuracy
|
||||
495
tests/utils/aime_eval.py
Normal file
495
tests/utils/aime_eval.py
Normal file
|
|
@ -0,0 +1,495 @@
|
|||
"""
|
||||
AIME Dataset Evaluation Module
|
||||
|
||||
This module provides functions to evaluate language models on the combined AIME dataset
|
||||
(test2024 + test2025-I + test2025-II).
|
||||
"""
|
||||
|
||||
import json
|
||||
import requests
|
||||
import os
|
||||
import re
|
||||
import logging
|
||||
from typing import List, Dict, Any
|
||||
from tqdm import tqdm
|
||||
from vllm import SamplingParams
|
||||
|
||||
|
||||
def download_and_combine_aime_datasets(data_dir: str = "./data/aime") -> str:
|
||||
"""Download all AIME datasets and combine them into a single file"""
|
||||
|
||||
datasets = {
|
||||
"test2024": "https://raw.githubusercontent.com/GAIR-NLP/AIME-Preview/main/eval/data/aime/test2024.jsonl",
|
||||
"test2025-I": "https://raw.githubusercontent.com/GAIR-NLP/AIME-Preview/main/eval/data/aime/test2025-I.jsonl",
|
||||
"test2025-II": "https://raw.githubusercontent.com/GAIR-NLP/AIME-Preview/main/eval/data/aime/test2025-II.jsonl"
|
||||
}
|
||||
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
combined_filepath = os.path.join(data_dir, "aime.jsonl")
|
||||
|
||||
# Check if combined file already exists
|
||||
if os.path.exists(combined_filepath):
|
||||
print(f"Combined AIME dataset already exists at {combined_filepath}")
|
||||
return combined_filepath
|
||||
|
||||
print("Downloading and combining AIME datasets...")
|
||||
|
||||
all_problems = []
|
||||
global_id = 0
|
||||
|
||||
for dataset_name, url in datasets.items():
|
||||
print(f" Downloading {dataset_name}...")
|
||||
|
||||
try:
|
||||
response = requests.get(url)
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse each line and add source information
|
||||
for line_num, line in enumerate(response.text.strip().split('\n')):
|
||||
if line.strip():
|
||||
try:
|
||||
data = json.loads(line)
|
||||
# Add source dataset information and global ID
|
||||
data['source_dataset'] = dataset_name
|
||||
data['original_id'] = data.get('id', line_num)
|
||||
data['global_id'] = global_id
|
||||
global_id += 1
|
||||
all_problems.append(data)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f" Warning: Error parsing line {line_num + 1} in {dataset_name}: {e}")
|
||||
continue
|
||||
|
||||
except requests.RequestException as e:
|
||||
print(f" Error downloading {dataset_name}: {e}")
|
||||
continue
|
||||
|
||||
# Write combined dataset
|
||||
if all_problems:
|
||||
with open(combined_filepath, 'w', encoding='utf-8') as f:
|
||||
for problem in all_problems:
|
||||
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}")
|
||||
|
||||
# Print summary by dataset
|
||||
for dataset_name in datasets.keys():
|
||||
count = sum(1 for p in all_problems if p['source_dataset'] == dataset_name)
|
||||
print(f" {dataset_name}: {count} problems")
|
||||
|
||||
else:
|
||||
raise RuntimeError("No problems were successfully downloaded")
|
||||
|
||||
return combined_filepath
|
||||
|
||||
|
||||
def load_aime_dataset(data_dir: str = "./data/aime") -> List[Dict[str, Any]]:
|
||||
"""Load combined AIME dataset and format for evaluation"""
|
||||
|
||||
# Download and combine if needed
|
||||
filepath = download_and_combine_aime_datasets(data_dir)
|
||||
|
||||
examples = []
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
for line_num, line in enumerate(f):
|
||||
line = line.strip()
|
||||
if line:
|
||||
try:
|
||||
data = json.loads(line)
|
||||
|
||||
# Format as expected by our evaluation
|
||||
formatted_example = {
|
||||
"global_id": data.get("global_id", line_num),
|
||||
"original_id": data.get("original_id", data.get("id", line_num)),
|
||||
"source_dataset": data.get("source_dataset", "unknown"),
|
||||
"problem": data["problem"],
|
||||
"answer": str(data["answer"]), # Ensure answer is string
|
||||
"solution": data.get("solution", ""),
|
||||
"url": data.get("url", ""),
|
||||
# Format as chat messages for the model
|
||||
"prompt": [
|
||||
{"role": "system", "content": "You are a mathematical problem solver. Solve the given problem step by step and provide your final answer clearly."},
|
||||
{"role": "user", "content": f"Problem: {data['problem']}\n\nSolve this step by step and provide your final numerical answer."}
|
||||
]
|
||||
}
|
||||
examples.append(formatted_example)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Error parsing line {line_num + 1}: {e}")
|
||||
continue
|
||||
|
||||
print(f"Loaded {len(examples)} problems from combined AIME dataset")
|
||||
|
||||
# Print breakdown by source
|
||||
source_counts = {}
|
||||
for example in examples:
|
||||
source = example['source_dataset']
|
||||
source_counts[source] = source_counts.get(source, 0) + 1
|
||||
|
||||
for source, count in source_counts.items():
|
||||
print(f" {source}: {count} problems")
|
||||
|
||||
return examples
|
||||
|
||||
|
||||
def extract_aime_answer(response: str) -> str:
|
||||
"""Extract numerical answer from AIME response"""
|
||||
|
||||
# AIME answers are integers from 0-999
|
||||
# Look for patterns like "The answer is 123" or just standalone numbers
|
||||
patterns = [
|
||||
r"(?:the )?(?:final )?answer is (\d{1,3})",
|
||||
r"(?:therefore|thus|so),?\s*(?:the )?(?:final )?answer is (\d{1,3})",
|
||||
r"\\boxed\{(\d{1,3})\}",
|
||||
r"\$\\boxed\{(\d{1,3})\}\$",
|
||||
r"(?:answer|result):\s*(\d{1,3})",
|
||||
r"(?:^|\n)\s*(\d{1,3})\s*(?:\n|$)", # Standalone number
|
||||
]
|
||||
|
||||
response_lower = response.lower().strip()
|
||||
|
||||
for pattern in patterns:
|
||||
matches = re.findall(pattern, response_lower, re.MULTILINE | re.IGNORECASE)
|
||||
if matches:
|
||||
# Get the last match (most likely to be final answer)
|
||||
answer = matches[-1]
|
||||
try:
|
||||
num = int(answer)
|
||||
if 0 <= num <= 999: # AIME answers are in range 0-999
|
||||
return str(num)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# If no clear pattern found, try to extract any 1-3 digit number
|
||||
numbers = re.findall(r'\b(\d{1,3})\b', response)
|
||||
if numbers:
|
||||
for num_str in reversed(numbers): # Check from end
|
||||
try:
|
||||
num = int(num_str)
|
||||
if 0 <= num <= 999:
|
||||
return str(num)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def get_num_tokens(text, tokenizer_instance):
|
||||
"""Count tokens in text"""
|
||||
if not text:
|
||||
return 0
|
||||
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):
|
||||
"""Evaluate model on combined AIME dataset with official configuration"""
|
||||
|
||||
print(f"\n{'='*70}")
|
||||
print(f"🧮 AIME EVALUATION - {model_type.upper()} MODEL")
|
||||
print(f"Combined Dataset: test2024 + test2025-I + test2025-II")
|
||||
print(f"{'='*70}")
|
||||
|
||||
# Load combined AIME dataset
|
||||
try:
|
||||
eval_dataset = load_aime_dataset()
|
||||
except Exception as e:
|
||||
print(f"Error loading dataset: {e}")
|
||||
return None
|
||||
|
||||
if not eval_dataset:
|
||||
print("No examples found in dataset")
|
||||
return None
|
||||
|
||||
# Initialize tracking variables
|
||||
records = {}
|
||||
input_tokens = []
|
||||
output_tokens = []
|
||||
correct_answers = 0
|
||||
|
||||
# Track performance by source dataset
|
||||
source_stats = {}
|
||||
for example in eval_dataset:
|
||||
source = example['source_dataset']
|
||||
if source not in source_stats:
|
||||
source_stats[source] = {'total': 0, 'correct': 0}
|
||||
source_stats[source]['total'] += 1
|
||||
|
||||
# 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,
|
||||
)
|
||||
|
||||
print(f"\n🔧 Configuration:")
|
||||
print(f" Temperature: {temperature}")
|
||||
print(f" Samples per question: {n_sampling}")
|
||||
print(f" Max tokens: {max_tokens}")
|
||||
print(f" Top-p: {top_p}")
|
||||
print(f" Seed: {seed}")
|
||||
|
||||
# Temporarily suppress verbose logging
|
||||
original_levels = {}
|
||||
loggers_to_suppress = ['vllm', 'vllm.engine', 'vllm.worker', 'vllm.model_executor', 'vllm.executor', 'ray']
|
||||
|
||||
for logger_name in loggers_to_suppress:
|
||||
logger = logging.getLogger(logger_name)
|
||||
original_levels[logger_name] = logger.level
|
||||
logger.setLevel(logging.WARNING)
|
||||
|
||||
try:
|
||||
print(f"\n🚀 Evaluating {len(eval_dataset)} problems...")
|
||||
|
||||
# Main evaluation loop
|
||||
with tqdm(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
|
||||
)
|
||||
|
||||
input_tokens.append(get_num_tokens(prompt_text, tokenizer))
|
||||
|
||||
# Generate multiple responses
|
||||
outputs = model.fast_generate(
|
||||
[prompt_text],
|
||||
sampling_params=sampling_params,
|
||||
lora_request=lora_request,
|
||||
use_tqdm=False,
|
||||
)[0].outputs
|
||||
|
||||
# Process all generated responses
|
||||
responses = [output.text for output in outputs]
|
||||
extracted_answers = [extract_aime_answer(response) for response in responses]
|
||||
|
||||
# Calculate total output tokens
|
||||
total_output_tokens = sum(get_num_tokens(response, tokenizer) for response in responses)
|
||||
output_tokens.append(total_output_tokens)
|
||||
|
||||
# Check if any answer is correct
|
||||
ground_truth = item["answer"]
|
||||
correct_responses = [ans == ground_truth for ans in extracted_answers]
|
||||
is_correct = any(correct_responses)
|
||||
|
||||
if is_correct:
|
||||
correct_answers += 1
|
||||
source_stats[item['source_dataset']]['correct'] += 1
|
||||
|
||||
# Store detailed record
|
||||
records[task_id] = {
|
||||
"global_id": item["global_id"],
|
||||
"original_id": item["original_id"],
|
||||
"source_dataset": item["source_dataset"],
|
||||
"problem": item["problem"],
|
||||
"ground_truth": ground_truth,
|
||||
"responses": responses,
|
||||
"extracted_answers": extracted_answers,
|
||||
"correct_responses": correct_responses,
|
||||
"is_correct": is_correct,
|
||||
"input_tokens": input_tokens[-1],
|
||||
"output_tokens": total_output_tokens,
|
||||
"n_correct": sum(correct_responses),
|
||||
"n_total": len(responses),
|
||||
"solution": item.get("solution", ""),
|
||||
"url": item.get("url", "")
|
||||
}
|
||||
|
||||
# Update progress
|
||||
current_accuracy = correct_answers / (task_id + 1) * 100
|
||||
pbar.set_postfix({
|
||||
'accuracy': f'{current_accuracy:.1f}%',
|
||||
'correct': correct_answers,
|
||||
'total': task_id + 1
|
||||
})
|
||||
pbar.update(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\nError processing problem {task_id}: {str(e)}")
|
||||
records[task_id] = {
|
||||
"global_id": item.get("global_id", task_id),
|
||||
"original_id": item.get("original_id", task_id),
|
||||
"source_dataset": item.get("source_dataset", "unknown"),
|
||||
"problem": item["problem"],
|
||||
"ground_truth": item["answer"],
|
||||
"error": str(e),
|
||||
"is_correct": False
|
||||
}
|
||||
pbar.update(1)
|
||||
continue
|
||||
|
||||
finally:
|
||||
# Restore logging levels
|
||||
for logger_name, level in original_levels.items():
|
||||
logging.getLogger(logger_name).setLevel(level)
|
||||
|
||||
# Calculate metrics
|
||||
total_problems = len(eval_dataset)
|
||||
accuracy = correct_answers / total_problems * 100
|
||||
|
||||
# Calculate Pass@k (probability that at least one of k samples is correct)
|
||||
pass_at_k_scores = []
|
||||
for record in records.values():
|
||||
if "n_correct" in record and "n_total" in record:
|
||||
n_correct = record["n_correct"]
|
||||
n_total = record["n_total"]
|
||||
if n_correct > 0:
|
||||
pass_at_k_scores.append(1.0)
|
||||
else:
|
||||
pass_at_k_scores.append(0.0)
|
||||
|
||||
pass_at_k = sum(pass_at_k_scores) / len(pass_at_k_scores) if pass_at_k_scores else 0
|
||||
|
||||
# Calculate per-source accuracies
|
||||
source_accuracies = {}
|
||||
for source, stats in source_stats.items():
|
||||
source_accuracies[source] = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
|
||||
|
||||
results = {
|
||||
"model_type": model_type,
|
||||
"dataset": "aime_combined",
|
||||
"total_problems": total_problems,
|
||||
"correct_answers": correct_answers,
|
||||
"accuracy": accuracy,
|
||||
"pass_at_k": pass_at_k * 100,
|
||||
"source_stats": source_stats,
|
||||
"source_accuracies": source_accuracies,
|
||||
"temperature": temperature,
|
||||
"n_sampling": n_sampling,
|
||||
"max_tokens": max_tokens,
|
||||
"top_p": top_p,
|
||||
"seed": seed,
|
||||
"avg_input_tokens": sum(input_tokens) / len(input_tokens) if input_tokens else 0,
|
||||
"avg_output_tokens": sum(output_tokens) / len(output_tokens) if output_tokens else 0,
|
||||
"max_input_tokens": max(input_tokens) if input_tokens else 0,
|
||||
"max_output_tokens": max(output_tokens) if output_tokens else 0,
|
||||
}
|
||||
|
||||
# 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)
|
||||
|
||||
# Print comprehensive summary
|
||||
print(f"\n{'='*70}")
|
||||
print(f"📊 AIME EVALUATION RESULTS - {model_type.upper()}")
|
||||
print(f"{'='*70}")
|
||||
|
||||
print(f"\n🎯 Overall Performance:")
|
||||
print(f" Total problems: {total_problems:>6}")
|
||||
print(f" Correct answers: {correct_answers:>6}/{total_problems} ({accuracy:>5.1f}%)")
|
||||
print(f" Pass@{n_sampling}: {pass_at_k:>10.1f}%")
|
||||
|
||||
print(f"\n📈 Performance by Dataset:")
|
||||
for source, stats in source_stats.items():
|
||||
source_acc = source_accuracies[source]
|
||||
print(f" {source:>12}: {stats['correct']:>3}/{stats['total']:>3} ({source_acc:>5.1f}%)")
|
||||
|
||||
print(f"\n🔧 Configuration:")
|
||||
print(f" Temperature: {temperature}")
|
||||
print(f" Samples per problem: {n_sampling}")
|
||||
print(f" Max tokens: {max_tokens}")
|
||||
print(f" Top-p: {top_p}")
|
||||
print(f" Seed: {seed}")
|
||||
|
||||
print(f"\n📝 Token Statistics:")
|
||||
print(f" Avg input tokens: {results['avg_input_tokens']:>10.1f}")
|
||||
print(f" Avg output tokens: {results['avg_output_tokens']:>10.1f}")
|
||||
print(f" Max input tokens: {results['max_input_tokens']:>10}")
|
||||
print(f" Max output tokens: {results['max_output_tokens']:>10}")
|
||||
|
||||
# Performance assessment for AIME
|
||||
if accuracy >= 50:
|
||||
tier = "🏆 EXCEPTIONAL"
|
||||
elif accuracy >= 30:
|
||||
tier = "✅ EXCELLENT"
|
||||
elif accuracy >= 20:
|
||||
tier = "🎯 VERY GOOD"
|
||||
elif accuracy >= 10:
|
||||
tier = "⚠️ GOOD"
|
||||
elif accuracy >= 5:
|
||||
tier = "📈 FAIR"
|
||||
else:
|
||||
tier = "❌ NEEDS IMPROVEMENT"
|
||||
|
||||
print(f"\n🎖️ AIME Performance: {tier} ({accuracy:.1f}%)")
|
||||
print(f"\n💾 Detailed results saved to: {filename}")
|
||||
print(f"\n{'='*70}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# Comparison functions for multiple model results
|
||||
def compare_aime_results(all_results):
|
||||
"""Generate comprehensive comparison for AIME evaluation results"""
|
||||
print(f"\n{'='*80}")
|
||||
print("COMPREHENSIVE AIME MODEL COMPARISON")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Main comparison table
|
||||
print(f"{'Model':<15} {'Accuracy %':<12} {'Pass@K %':<10} {'Correct':<8} {'Total':<8}")
|
||||
print("-" * 80)
|
||||
|
||||
for result in all_results:
|
||||
print(f"{result['model_type']:<15} "
|
||||
f"{result['accuracy']:<12.1f} "
|
||||
f"{result['pass_at_k']:<10.1f} "
|
||||
f"{result['correct_answers']:<8} "
|
||||
f"{result['total_problems']:<8}")
|
||||
|
||||
# Performance improvement analysis
|
||||
if len(all_results) > 1:
|
||||
print(f"\n{'='*50}")
|
||||
print("IMPROVEMENT ANALYSIS")
|
||||
print(f"{'='*50}")
|
||||
|
||||
base_result = all_results[0] # Assume first is base model
|
||||
|
||||
for i, result in enumerate(all_results[1:], 1):
|
||||
print(f"\n{result['model_type']} vs {base_result['model_type']}:")
|
||||
|
||||
accuracy_improvement = result['accuracy'] - base_result['accuracy']
|
||||
pass_k_improvement = result['pass_at_k'] - base_result['pass_at_k']
|
||||
|
||||
print(f" Accuracy improvement: {accuracy_improvement:+.1f}%")
|
||||
print(f" Pass@K improvement: {pass_k_improvement:+.1f}%")
|
||||
|
||||
# Dataset breakdown
|
||||
print(f"\n{'='*50}")
|
||||
print("PERFORMANCE BY DATASET")
|
||||
print(f"{'='*50}")
|
||||
|
||||
# Get all unique datasets from the first result
|
||||
if all_results and 'source_accuracies' in all_results[0]:
|
||||
datasets = list(all_results[0]['source_accuracies'].keys())
|
||||
|
||||
print(f"{'Model':<15}", end="")
|
||||
for dataset in datasets:
|
||||
print(f"{dataset:<15}", end="")
|
||||
print()
|
||||
print("-" * (15 + 15 * len(datasets)))
|
||||
|
||||
for result in all_results:
|
||||
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()
|
||||
|
||||
# Save comparison
|
||||
comparison_data = {
|
||||
"summary": all_results,
|
||||
"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)
|
||||
|
||||
print(f"\nBest performing model: {comparison_data['best_model']['model_type']} "
|
||||
f"({comparison_data['best_model']['accuracy']:.1f}% accuracy)")
|
||||
209
tests/utils/cleanup_utils.py
Normal file
209
tests/utils/cleanup_utils.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
import gc
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import torch
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
def clear_memory(variables_to_clear=None, verbose=False, clear_all_caches=True):
|
||||
"""
|
||||
Comprehensive memory clearing for persistent memory leaks.
|
||||
|
||||
Args:
|
||||
variables_to_clear: List of variable names to clear
|
||||
verbose: Print memory status
|
||||
clear_all_caches: Clear all types of caches (recommended for memory leaks)
|
||||
"""
|
||||
|
||||
# Save current logging levels
|
||||
saved_log_levels = {}
|
||||
for name, logger in logging.Logger.manager.loggerDict.items():
|
||||
if isinstance(logger, logging.Logger):
|
||||
saved_log_levels[name] = logger.level
|
||||
root_level = logging.getLogger().level
|
||||
|
||||
if variables_to_clear is None:
|
||||
variables_to_clear = ["inputs", "model", "base_model", "processor", "tokenizer",
|
||||
"base_processor", "base_tokenizer", "trainer",
|
||||
"peft_model", "bnb_config"]
|
||||
|
||||
# 1. Clear LRU caches FIRST (very important for memory leaks)
|
||||
if clear_all_caches:
|
||||
clear_all_lru_caches(verbose)
|
||||
|
||||
# 2. Delete specified variables
|
||||
g = globals()
|
||||
deleted_vars = []
|
||||
for var in variables_to_clear:
|
||||
if var in g:
|
||||
del g[var]
|
||||
deleted_vars.append(var)
|
||||
|
||||
if verbose and deleted_vars:
|
||||
print(f"Deleted variables: {deleted_vars}")
|
||||
|
||||
# 3. Multiple garbage collection passes (important for circular references)
|
||||
for i in range(3):
|
||||
collected = gc.collect()
|
||||
if verbose and collected > 0:
|
||||
print(f"GC pass {i+1}: collected {collected} objects")
|
||||
|
||||
# 4. CUDA cleanup
|
||||
if torch.cuda.is_available():
|
||||
# Get memory before cleanup
|
||||
if verbose:
|
||||
mem_before = torch.cuda.memory_allocated() / 1024**3
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Additional CUDA cleanup for persistent leaks
|
||||
if clear_all_caches:
|
||||
# Reset memory stats
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.cuda.reset_accumulated_memory_stats()
|
||||
|
||||
# Clear JIT cache
|
||||
if hasattr(torch.jit, '_state') and hasattr(torch.jit._state, '_clear_class_state'):
|
||||
torch.jit._state._clear_class_state()
|
||||
|
||||
# Force another CUDA cache clear
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Final garbage collection
|
||||
gc.collect()
|
||||
|
||||
if verbose:
|
||||
mem_after = torch.cuda.memory_allocated() / 1024**3
|
||||
mem_reserved = torch.cuda.memory_reserved() / 1024**3
|
||||
print(f"GPU memory - Before: {mem_before:.2f} GB, After: {mem_after:.2f} GB")
|
||||
print(f"GPU reserved memory: {mem_reserved:.2f} GB")
|
||||
if mem_before > 0:
|
||||
print(f"Memory freed: {mem_before - mem_after:.2f} GB")
|
||||
|
||||
# restore original logging levels
|
||||
logging.getLogger().setLevel(root_level)
|
||||
for name, level in saved_log_levels.items():
|
||||
if name in logging.Logger.manager.loggerDict:
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
|
||||
def clear_all_lru_caches(verbose=True):
|
||||
"""Clear all LRU caches in loaded modules."""
|
||||
cleared_caches = []
|
||||
|
||||
# Modules to skip to avoid warnings
|
||||
skip_modules = {
|
||||
'torch.distributed',
|
||||
'torchaudio',
|
||||
'torch._C',
|
||||
'torch.distributed.reduce_op',
|
||||
'torchaudio.backend',
|
||||
}
|
||||
|
||||
# Create a static list of modules to avoid RuntimeError
|
||||
modules = list(sys.modules.items())
|
||||
|
||||
# Method 1: Clear caches in all loaded modules
|
||||
for module_name, module in modules:
|
||||
if module is None:
|
||||
continue
|
||||
|
||||
# Skip problematic modules
|
||||
if any(module_name.startswith(skip) for skip in skip_modules):
|
||||
continue
|
||||
|
||||
try:
|
||||
# Look for functions with lru_cache
|
||||
for attr_name in dir(module):
|
||||
try:
|
||||
# Suppress warnings when checking attributes
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", FutureWarning)
|
||||
warnings.simplefilter("ignore", UserWarning)
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
|
||||
attr = getattr(module, attr_name)
|
||||
if hasattr(attr, 'cache_clear'):
|
||||
attr.cache_clear()
|
||||
cleared_caches.append(f"{module_name}.{attr_name}")
|
||||
except Exception:
|
||||
continue # Skip problematic attributes
|
||||
except Exception:
|
||||
continue # Skip problematic modules
|
||||
|
||||
# Method 2: Clear specific known caches
|
||||
known_caches = [
|
||||
'transformers.utils.hub.cached_file',
|
||||
'transformers.tokenization_utils_base.get_tokenizer',
|
||||
'torch._dynamo.utils.counters',
|
||||
]
|
||||
|
||||
for cache_path in known_caches:
|
||||
try:
|
||||
parts = cache_path.split('.')
|
||||
module = sys.modules.get(parts[0])
|
||||
if module:
|
||||
obj = module
|
||||
for part in parts[1:]:
|
||||
obj = getattr(obj, part, None)
|
||||
if obj is None:
|
||||
break
|
||||
if obj and hasattr(obj, 'cache_clear'):
|
||||
obj.cache_clear()
|
||||
cleared_caches.append(cache_path)
|
||||
except Exception:
|
||||
continue # Skip problematic caches
|
||||
|
||||
if verbose and cleared_caches:
|
||||
print(f"Cleared {len(cleared_caches)} LRU caches")
|
||||
|
||||
|
||||
def clear_specific_lru_cache(func):
|
||||
"""Clear cache for a specific function."""
|
||||
if hasattr(func, 'cache_clear'):
|
||||
func.cache_clear()
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# Additional utility for monitoring cache sizes
|
||||
def monitor_cache_sizes():
|
||||
"""Monitor LRU cache sizes across modules."""
|
||||
cache_info = []
|
||||
|
||||
for module_name, module in sys.modules.items():
|
||||
if module is None:
|
||||
continue
|
||||
try:
|
||||
for attr_name in dir(module):
|
||||
try:
|
||||
attr = getattr(module, attr_name)
|
||||
if hasattr(attr, 'cache_info'):
|
||||
info = attr.cache_info()
|
||||
cache_info.append({
|
||||
'function': f"{module_name}.{attr_name}",
|
||||
'size': info.currsize,
|
||||
'hits': info.hits,
|
||||
'misses': info.misses
|
||||
})
|
||||
except:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
||||
return sorted(cache_info, key=lambda x: x['size'], reverse=True)
|
||||
|
||||
|
||||
def safe_remove_directory(path):
|
||||
try:
|
||||
if os.path.exists(path) and os.path.isdir(path):
|
||||
shutil.rmtree(path)
|
||||
return True
|
||||
else:
|
||||
print(f"Path {path} is not a valid directory")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Failed to remove directory {path}: {e}")
|
||||
return False
|
||||
109
tests/utils/ocr_eval.md
Normal file
109
tests/utils/ocr_eval.md
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
|
||||
# OCR Model Evaluator
|
||||
A comprehensive Python module for evaluating Optical Character Recognition (OCR) models using Word Error Rate (WER) and Character Error Rate (CER) metrics. This evaluator supports vision-language models and provides detailed analysis with comparison capabilities across multiple models
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
from ocr_evaluator import evaluate_ocr_model
|
||||
|
||||
# Simple evaluation
|
||||
avg_wer, avg_cer = evaluate_ocr_model(
|
||||
model=your_model,
|
||||
processor=your_processor,
|
||||
dataset=your_dataset,
|
||||
output_dir="evaluation_results"
|
||||
)
|
||||
|
||||
print(f"Average WER: {avg_wer:.4f}")
|
||||
print(f"Average CER: {avg_cer:.4f}")
|
||||
```
|
||||
|
||||
|
||||
### Dataset Format
|
||||
|
||||
The evaluator expects datasets in a chatml conversational format with the following structure:
|
||||
```
|
||||
dataset = [
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "You are an OCR system."}]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Extract text from this image"},
|
||||
{"type": "image", "image": PIL_Image_object}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Ground truth text"}]
|
||||
}
|
||||
]
|
||||
},
|
||||
# ... more samples
|
||||
]
|
||||
```
|
||||
|
||||
|
||||
## Examples
|
||||
|
||||
### Document OCR evaluation
|
||||
|
||||
```python
|
||||
from ocr_evaluator import OCRModelEvaluator
|
||||
from datasets import load_dataset
|
||||
|
||||
# Load document OCR dataset
|
||||
dataset = load_dataset("your-ocr-dataset", split="test")
|
||||
|
||||
# Convert to required format
|
||||
eval_data = [format_document_sample(sample) for sample in dataset]
|
||||
|
||||
# Evaluate models
|
||||
evaluator = OCRModelEvaluator()
|
||||
|
||||
# Compare different model configurations
|
||||
configs = {
|
||||
"Standard Model": {"temperature": 1.0, "max_new_tokens": 512},
|
||||
"Conservative Model": {"temperature": 0.7, "max_new_tokens": 256},
|
||||
"Creative Model": {"temperature": 1.5, "max_new_tokens": 1024}
|
||||
}
|
||||
|
||||
for config_name, params in configs.items():
|
||||
wer, cer = evaluator.evaluate_model(
|
||||
model=base_model,
|
||||
processor=processor,
|
||||
dataset=eval_data,
|
||||
output_dir=f"document_ocr_{config_name.lower().replace(' ', '_')}",
|
||||
**params
|
||||
)
|
||||
evaluator.add_to_comparison(config_name, wer, cer)
|
||||
|
||||
# Generate final report
|
||||
evaluator.print_model_comparison()
|
||||
```
|
||||
|
||||
### Handwritting Recognition
|
||||
```python
|
||||
# Specialized evaluation for handwriting
|
||||
def evaluate_handwriting_models(models, handwriting_dataset):
|
||||
evaluator = OCRModelEvaluator()
|
||||
|
||||
for model_name, (model, processor) in models.items():
|
||||
# Adjust parameters for handwriting recognition
|
||||
wer, cer = evaluator.evaluate_model(
|
||||
model=model,
|
||||
processor=processor,
|
||||
dataset=handwriting_dataset,
|
||||
temperature=1.2, # Slightly higher for handwriting variety
|
||||
max_new_tokens=128, # Usually shorter text
|
||||
output_dir=f"handwriting_{model_name}"
|
||||
)
|
||||
evaluator.add_to_comparison(f"Handwriting - {model_name}", wer, cer)
|
||||
|
||||
return evaluator.print_model_comparison()
|
||||
```
|
||||
352
tests/utils/ocr_eval.py
Normal file
352
tests/utils/ocr_eval.py
Normal file
|
|
@ -0,0 +1,352 @@
|
|||
"""
|
||||
OCR Model Evaluation Module
|
||||
|
||||
This module provides functionality to evaluate OCR models on datasets with
|
||||
word error rate (WER) and character error rate (CER) metrics.
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
from jiwer import wer, cer
|
||||
from qwen_vl_utils import process_vision_info
|
||||
import matplotlib.pyplot as plt
|
||||
from typing import List, Dict, Tuple, Optional, Any
|
||||
import traceback
|
||||
|
||||
|
||||
class OCRModelEvaluator:
|
||||
"""
|
||||
A comprehensive OCR model evaluator that supports multiple models and provides
|
||||
detailed analysis with WER and CER metrics.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the OCR evaluator."""
|
||||
self.model_comparison_results = {}
|
||||
|
||||
def evaluate_model(
|
||||
self,
|
||||
model: Any,
|
||||
processor: Any,
|
||||
dataset: List[Dict],
|
||||
output_dir: str = "ocr_evaluation_results",
|
||||
max_new_tokens: int = 1024,
|
||||
temperature: float = 1.5,
|
||||
min_p: float = 0.1,
|
||||
verbose: bool = True
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
"""
|
||||
Evaluate a model on an OCR dataset.
|
||||
"""
|
||||
# Create output directory if it doesn't exist
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Initialize results storage
|
||||
results = []
|
||||
|
||||
# Process each sample in the dataset
|
||||
for i, sample in enumerate(tqdm(dataset, desc="Evaluating OCR performance", disable=not verbose)):
|
||||
try:
|
||||
# Extract components from sample
|
||||
messages = sample['messages']
|
||||
|
||||
# Get ground truth, image, and question
|
||||
ground_truth, image, question, input_messages = self._extract_sample_components(
|
||||
messages, i, verbose
|
||||
)
|
||||
|
||||
if ground_truth is None or image is None or question is None:
|
||||
continue
|
||||
|
||||
# Generate model response
|
||||
generated_response = self._generate_response(
|
||||
model, processor, input_messages, max_new_tokens, temperature, min_p
|
||||
)
|
||||
|
||||
# Calculate metrics
|
||||
word_error = wer(ground_truth, generated_response)
|
||||
char_error = cer(ground_truth, generated_response)
|
||||
|
||||
# Save individual result
|
||||
self._save_individual_result(
|
||||
output_dir, i, question, generated_response, ground_truth, word_error, char_error
|
||||
)
|
||||
|
||||
# Store results for summary
|
||||
results.append({
|
||||
'sample_id': i,
|
||||
'wer': word_error,
|
||||
'cer': char_error,
|
||||
'model_output': generated_response.strip(),
|
||||
'ground_truth': ground_truth,
|
||||
'question': question
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f"Error processing sample {i}: {str(e)}")
|
||||
traceback.print_exc()
|
||||
|
||||
# Generate summary report
|
||||
return self._generate_summary_report(results, output_dir, verbose)
|
||||
|
||||
def _extract_sample_components(
|
||||
self,
|
||||
messages: List[Dict],
|
||||
sample_idx: int,
|
||||
verbose: bool
|
||||
) -> Tuple[Optional[str], Optional[Any], Optional[str], List[Dict]]:
|
||||
"""Extract ground truth, image, question, and input messages from sample."""
|
||||
|
||||
# Extract system message (if present)
|
||||
system_message = next((msg for msg in messages if msg['role'] == 'system'), None)
|
||||
|
||||
# Extract user message with the image and question
|
||||
user_message = next((msg for msg in messages if msg['role'] == 'user'), None)
|
||||
if not user_message:
|
||||
if verbose:
|
||||
print(f"Skipping sample {sample_idx}: No user message found")
|
||||
return None, None, None, []
|
||||
|
||||
# Extract assistant message with ground truth
|
||||
assistant_message = next((msg for msg in messages if msg['role'] == 'assistant'), None)
|
||||
if not assistant_message:
|
||||
if verbose:
|
||||
print(f"Skipping sample {sample_idx}: No assistant message (ground truth) found")
|
||||
return None, None, None, []
|
||||
|
||||
# Extract ground truth text
|
||||
ground_truth = None
|
||||
for content_item in assistant_message['content']:
|
||||
if content_item['type'] == 'text':
|
||||
ground_truth = content_item['text']
|
||||
break
|
||||
|
||||
if not ground_truth:
|
||||
if verbose:
|
||||
print(f"Skipping sample {sample_idx}: No text found in assistant message")
|
||||
return None, None, None, []
|
||||
|
||||
# Extract image and question from user message
|
||||
image = None
|
||||
question = None
|
||||
|
||||
for content_item in user_message['content']:
|
||||
if content_item['type'] == 'image':
|
||||
image = content_item['image']
|
||||
elif content_item['type'] == 'text':
|
||||
question = content_item['text']
|
||||
|
||||
if not image:
|
||||
if verbose:
|
||||
print(f"Skipping sample {sample_idx}: No image found in user message")
|
||||
return None, None, None, []
|
||||
|
||||
if not question:
|
||||
if verbose:
|
||||
print(f"Skipping sample {sample_idx}: No question found in user message")
|
||||
return None, None, None, []
|
||||
|
||||
# Construct messages for the model input (excluding assistant message)
|
||||
input_messages = []
|
||||
if system_message:
|
||||
input_messages.append(system_message)
|
||||
input_messages.append(user_message)
|
||||
|
||||
return ground_truth, image, question, input_messages
|
||||
|
||||
def _generate_response(
|
||||
self,
|
||||
model: Any,
|
||||
processor: Any,
|
||||
input_messages: List[Dict],
|
||||
max_new_tokens: int,
|
||||
temperature: float,
|
||||
min_p: float
|
||||
) -> str:
|
||||
"""Generate response from the model."""
|
||||
|
||||
# Preparation for inference using Qwen's specific processing
|
||||
text = processor.apply_chat_template(
|
||||
input_messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Process vision info (images/videos) from messages
|
||||
image_inputs, video_inputs = process_vision_info(input_messages)
|
||||
|
||||
# Create model inputs
|
||||
inputs = processor(
|
||||
text=[text],
|
||||
images=image_inputs,
|
||||
videos=video_inputs,
|
||||
padding=True,
|
||||
return_tensors="pt"
|
||||
)
|
||||
inputs = inputs.to(model.device)
|
||||
|
||||
# Generate response
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
temperature=temperature,
|
||||
min_p=min_p,
|
||||
use_cache=True
|
||||
)
|
||||
|
||||
# Extract only the generated part (not the input)
|
||||
generated_ids_trimmed = [
|
||||
out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
||||
]
|
||||
|
||||
# Decode the generated text
|
||||
generated_response = processor.batch_decode(
|
||||
generated_ids_trimmed,
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False
|
||||
)[0]
|
||||
|
||||
return generated_response
|
||||
|
||||
def _save_individual_result(
|
||||
self,
|
||||
output_dir: str,
|
||||
sample_idx: int,
|
||||
question: str,
|
||||
generated_response: str,
|
||||
ground_truth: str,
|
||||
word_error: float,
|
||||
char_error: float
|
||||
):
|
||||
"""Save individual sample result to file."""
|
||||
output_file = os.path.join(output_dir, f"sample_{sample_idx}.txt")
|
||||
with open(output_file, 'w', encoding='utf-8') as f:
|
||||
f.write(f"Sample {sample_idx}\n")
|
||||
f.write(f"Question: {question}\n\n")
|
||||
f.write(f"Model output:\n{generated_response.strip()}\n\n")
|
||||
f.write(f"Ground truth:\n{ground_truth}\n\n")
|
||||
f.write(f"WER: {word_error:.4f}, CER: {char_error:.4f}")
|
||||
|
||||
def _generate_summary_report(
|
||||
self,
|
||||
results: List[Dict],
|
||||
output_dir: str,
|
||||
verbose: bool
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
"""Generate and save summary report."""
|
||||
if not results:
|
||||
if verbose:
|
||||
print("No results to summarize.")
|
||||
return None, None
|
||||
|
||||
df = pd.DataFrame(results)
|
||||
|
||||
# Calculate overall averages
|
||||
avg_wer = df['wer'].mean()
|
||||
avg_cer = df['cer'].mean()
|
||||
|
||||
# Save average metrics
|
||||
with open(os.path.join(output_dir, "avg_metrics.txt"), 'w') as f:
|
||||
f.write(f"Average WER: {avg_wer:.4f}\n")
|
||||
f.write(f"Average CER: {avg_cer:.4f}\n")
|
||||
|
||||
# Save detailed results
|
||||
df.to_csv(os.path.join(output_dir, "detailed_results.csv"), index=False)
|
||||
|
||||
if verbose:
|
||||
print("\nResults Summary:")
|
||||
print(f"Average WER: {avg_wer:.4f}")
|
||||
print(f"Average CER: {avg_cer:.4f}")
|
||||
print(f"\nDetailed results saved to {output_dir}/")
|
||||
|
||||
return avg_wer, avg_cer
|
||||
|
||||
def add_to_comparison(self, model_name: str, wer: float, cer: float):
|
||||
"""Add model results to the comparison tracker."""
|
||||
self.model_comparison_results[model_name] = {
|
||||
"wer": wer,
|
||||
"cer": cer
|
||||
}
|
||||
|
||||
def print_model_comparison(self, save_csv: bool = True, save_plot: bool = True) -> Optional[pd.DataFrame]:
|
||||
"""Print a comparison of all models evaluated so far."""
|
||||
if not self.model_comparison_results:
|
||||
print("No model results available for comparison")
|
||||
return None
|
||||
|
||||
print("\n==== MODEL COMPARISON REPORT ====")
|
||||
|
||||
# Create a comparison dataframe
|
||||
comparison_df = pd.DataFrame({
|
||||
"Model": list(self.model_comparison_results.keys()),
|
||||
"WER": [results["wer"] for results in self.model_comparison_results.values()],
|
||||
"CER": [results["cer"] for results in self.model_comparison_results.values()]
|
||||
})
|
||||
|
||||
# Sort by WER (best performance first)
|
||||
comparison_df = comparison_df.sort_values("WER")
|
||||
|
||||
# Display the comparison table
|
||||
print("\nComparison Table (sorted by WER):")
|
||||
print(comparison_df.to_string(index=False))
|
||||
|
||||
# Save the comparison table
|
||||
if save_csv:
|
||||
comparison_file = "model_comparison_results.csv"
|
||||
comparison_df.to_csv(comparison_file, index=False)
|
||||
print(f"\nComparison table saved to {comparison_file}")
|
||||
|
||||
# Generate a bar chart visualization
|
||||
if save_plot:
|
||||
self._create_comparison_plot(comparison_df)
|
||||
|
||||
return comparison_df
|
||||
|
||||
def _create_comparison_plot(self, comparison_df: pd.DataFrame):
|
||||
"""Create and save comparison plot."""
|
||||
plt.figure(figsize=(12, 6))
|
||||
|
||||
# Plot WER
|
||||
plt.subplot(1, 2, 1)
|
||||
plt.bar(comparison_df["Model"], comparison_df["WER"], color='skyblue')
|
||||
plt.title('Word Error Rate Comparison')
|
||||
plt.ylabel('WER (lower is better)')
|
||||
plt.ylim(bottom=0)
|
||||
plt.xticks(rotation=45, ha='right')
|
||||
|
||||
# Plot CER
|
||||
plt.subplot(1, 2, 2)
|
||||
plt.bar(comparison_df["Model"], comparison_df["CER"], color='lightgreen')
|
||||
plt.title('Character Error Rate Comparison')
|
||||
plt.ylabel('CER (lower is better)')
|
||||
plt.ylim(bottom=0)
|
||||
plt.xticks(rotation=45, ha='right')
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig('ocr_model_comparison.png')
|
||||
plt.show()
|
||||
|
||||
print(f"\nVisualization saved to ocr_model_comparison.png")
|
||||
|
||||
def get_comparison_results(self) -> Dict[str, Dict[str, float]]:
|
||||
"""Get the current comparison results."""
|
||||
return self.model_comparison_results.copy()
|
||||
|
||||
def clear_comparison_results(self):
|
||||
"""Clear all comparison results."""
|
||||
self.model_comparison_results.clear()
|
||||
|
||||
|
||||
def evaluate_ocr_model(model, processor, dataset, output_dir="ocr_evaluation_results", **kwargs):
|
||||
"""
|
||||
Convenience function that maintains backward compatibility with the original function.
|
||||
"""
|
||||
evaluator = OCRModelEvaluator()
|
||||
return evaluator.evaluate_model(model, processor, dataset, output_dir, **kwargs)
|
||||
|
||||
|
||||
def create_evaluator():
|
||||
"""Create a new OCR evaluator instance."""
|
||||
return OCRModelEvaluator()
|
||||
20
tests/utils/perplexity_eval.md
Normal file
20
tests/utils/perplexity_eval.md
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
# Language Model Perplexity Evaluator
|
||||
|
||||
A Python module for evaluating language models using perplexity metrics with sliding window approach for long sequences. This evaluator provides efficient computation of perplexity scores across datasets with model comparison capabilities.
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
from perplexity_evaluator import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Simple perplexity evaluation
|
||||
dataset = {"text": ["Your text samples here...", "Another text sample..."]}
|
||||
perplexity = ppl_model(model, tokenizer, dataset)
|
||||
|
||||
print(f"Model Perplexity: {perplexity:.4f}")
|
||||
|
||||
# Add to comparison tracker
|
||||
add_to_comparison("My Model", perplexity)
|
||||
print_model_comparison()
|
||||
```
|
||||
|
||||
75
tests/utils/perplexity_eval.py
Normal file
75
tests/utils/perplexity_eval.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
from tqdm import tqdm
|
||||
import torch
|
||||
import pandas as pd
|
||||
|
||||
model_comparison_results = {}
|
||||
#return the perplexity of the model on the dataset
|
||||
#The perplexity is computed on each example, individually, with a sliding window for examples longer than 512 tokens.
|
||||
|
||||
|
||||
def ppl_model(model, tokenizer, dataset):
|
||||
nlls = []
|
||||
max_length = 2048
|
||||
stride = 512
|
||||
for s in tqdm(range(len(dataset['text']))):
|
||||
encodings = tokenizer(dataset['text'][s], return_tensors="pt")
|
||||
seq_len = encodings.input_ids.size(1)
|
||||
prev_end_loc = 0
|
||||
for begin_loc in range(0, seq_len, stride):
|
||||
end_loc = min(begin_loc + max_length, seq_len)
|
||||
trg_len = end_loc - prev_end_loc
|
||||
input_ids = encodings.input_ids[:, begin_loc:end_loc].to("cuda")
|
||||
target_ids = input_ids.clone()
|
||||
target_ids[:, :-trg_len] = -100
|
||||
# Create attention mask based on pad token id
|
||||
pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||||
attention_mask = (input_ids != pad_token_id).long()
|
||||
with torch.no_grad():
|
||||
outputs = model(input_ids, labels=target_ids, attention_mask=attention_mask)
|
||||
neg_log_likelihood = outputs.loss
|
||||
nlls.append(neg_log_likelihood)
|
||||
prev_end_loc = end_loc
|
||||
if end_loc == seq_len:
|
||||
break
|
||||
ppl = torch.exp(torch.stack(nlls).mean())
|
||||
return ppl
|
||||
|
||||
|
||||
#--------------------------------------------------------------------
|
||||
|
||||
|
||||
## ----------- Reporting helper function ----------- ##
|
||||
|
||||
# Create a simple function to add results to the comparison
|
||||
def add_to_comparison(model_name, ppl):
|
||||
"""Add model results to the comparison tracker"""
|
||||
model_comparison_results[model_name] = {
|
||||
"ppl": ppl
|
||||
}
|
||||
#return model_comparison_results
|
||||
|
||||
|
||||
|
||||
# Create a function to print the comparison report whenever needed
|
||||
def print_model_comparison():
|
||||
"""Print a comparison of all models evaluated so far"""
|
||||
if not model_comparison_results:
|
||||
print("No model results available for comparison")
|
||||
return
|
||||
|
||||
print("\n==== MODEL COMPARISON REPORT ====")
|
||||
|
||||
# Create a comparison dataframe
|
||||
comparison_df = pd.DataFrame({
|
||||
"Model": list(model_comparison_results.keys()),
|
||||
#"Perplexity": [results["ppl"] for results in model_comparison_results.values()],
|
||||
"Perplexity": [
|
||||
# Convert tensors to CPU and then to float if needed
|
||||
results["ppl"].cpu().item() if torch.is_tensor(results["ppl"]) else results["ppl"]
|
||||
for results in model_comparison_results.values()
|
||||
],
|
||||
})
|
||||
|
||||
# Display the comparison table
|
||||
print("\nComparison Table:")
|
||||
print(comparison_df.to_string(index=False))
|
||||
|
|
@ -1625,9 +1625,9 @@ def create_ollama_modelfile(tokenizer, gguf_location):
|
|||
pass
|
||||
|
||||
def create_ollama_model(
|
||||
username: str,
|
||||
model_name: str,
|
||||
tag: str,
|
||||
username: str,
|
||||
model_name: str,
|
||||
tag: str,
|
||||
modelfile_path: str
|
||||
):
|
||||
try:
|
||||
|
|
@ -1711,7 +1711,7 @@ def push_to_ollama(
|
|||
with open(f"Modelfile_{model_name}", "w") as f:
|
||||
f.write(model_file)
|
||||
f.close()
|
||||
|
||||
|
||||
create_ollama_model(
|
||||
username=username,
|
||||
model_name=model_name,
|
||||
|
|
@ -2320,7 +2320,7 @@ def unsloth_generic_save(
|
|||
)
|
||||
elif save_method == "merged_4bit_forced":
|
||||
save_method = "merged_4bit"
|
||||
|
||||
|
||||
merge_and_overwrite_lora(
|
||||
get_model_name,
|
||||
model = model,
|
||||
|
|
@ -2524,8 +2524,8 @@ def patch_saving_functions(model, vision = False):
|
|||
if not vision:
|
||||
if hasattr(model, "config"):
|
||||
# Counteract tokenizers
|
||||
model.push_to_hub_merged = types.MethodType(unsloth_push_to_hub_merged, model)
|
||||
model.save_pretrained_merged = types.MethodType(unsloth_save_pretrained_merged, model)
|
||||
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
|
||||
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
|
||||
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
|
||||
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
|
||||
model.push_to_hub_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_push_to_hub, model)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue