a6dc10dad2* Reduce and tighten comments and docstrings in tests Shorten verbose comments and docstrings across the test suite without changing any test logic. Remove narration that restates the next line, collapse long module and test docstrings to a single line, and drop banner separators. Keep regression context (issue and PR references, run ids), skip reasons, mocking and timing rationale, license headers, lint and type directives, and commented-out code. Comments and docstrings only: an AST signature check confirms no code, assertions, or string literals changed, and the suite byte-compiles cleanly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
95 lines
2.5 KiB
Python
95 lines
2.5 KiB
Python
from unsloth import FastLanguageModel
|
|
from trl import SFTTrainer, SFTConfig
|
|
from datasets import load_dataset
|
|
import torch
|
|
import gc
|
|
import os
|
|
import shutil
|
|
|
|
|
|
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
|
|
|
|
|
|
# Used by formatting_prompts_func below.
|
|
tokenizer = None
|
|
|
|
|
|
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}
|
|
|
|
|
|
# --- Load 4-bit Model and Train ---
|
|
print("Loading 4-bit Mxfp4 gpt-oss model for training...")
|
|
max_seq_length = 1024
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
"unsloth/gpt-oss-20b", max_seq_length = max_seq_length, load_in_4bit = True
|
|
)
|
|
|
|
dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split = "train[:50]").map(
|
|
formatting_prompts_func, batched = True
|
|
)
|
|
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r = 8,
|
|
target_modules = [
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"o_proj",
|
|
"gate_proj",
|
|
"up_proj",
|
|
"down_proj",
|
|
],
|
|
lora_alpha = 16,
|
|
use_gradient_checkpointing = "unsloth",
|
|
random_state = 3407,
|
|
)
|
|
|
|
trainer = SFTTrainer(
|
|
model = model,
|
|
tokenizer = tokenizer,
|
|
train_dataset = dataset,
|
|
args = SFTConfig(
|
|
per_device_train_batch_size = 1,
|
|
gradient_accumulation_steps = 4,
|
|
max_steps = 10,
|
|
learning_rate = 2e-4,
|
|
output_dir = "outputs",
|
|
report_to = "none",
|
|
),
|
|
)
|
|
|
|
print("Starting fine-tuning...")
|
|
trainer.train()
|
|
print("Fine-tuning complete.")
|
|
|
|
# --- Merge and Save ---
|
|
print("\n💾 Merging and saving the 16-bit model to './gpt-oss-finetuned-merged'...")
|
|
model.save_pretrained_merged(save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer)
|
|
print("✅ Model merged and saved.")
|
|
|
|
# --- Cleanup ---
|
|
print("\n🧹 Cleaning up training artifacts...")
|
|
del model, trainer, tokenizer, dataset
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
safe_remove_directory("./outputs")
|
|
safe_remove_directory("./unsloth_compiled_cache") # cache created by this process
|
|
print("✅ Cleanup complete. Exiting training script.")
|