32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
#!/usr/bin/env python3
|
|
"""WEVAL Brain Fine-Tune — Run on Kaggle T4 GPU (30h/week free)"""
|
|
# pip install trl transformers datasets peft bitsandbytes
|
|
from trl import SFTTrainer, SFTConfig
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from datasets import load_dataset
|
|
import torch
|
|
|
|
MODEL = "Qwen/Qwen2.5-3B-Instruct" # Small enough for T4 16GB
|
|
DATA = "weval-finetune-data.jsonl"
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL)
|
|
mdl = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float16, device_map="auto")
|
|
ds = load_dataset("json", data_files=DATA, split="train")
|
|
|
|
cfg = SFTConfig(
|
|
output_dir="./weval-brain-v4",
|
|
num_train_epochs=3,
|
|
per_device_train_batch_size=2,
|
|
gradient_accumulation_steps=4,
|
|
learning_rate=2e-5,
|
|
max_seq_length=1024,
|
|
logging_steps=10,
|
|
save_steps=50,
|
|
fp16=True,
|
|
)
|
|
|
|
trainer = SFTTrainer(model=mdl, args=cfg, train_dataset=ds, tokenizer=tok)
|
|
trainer.train()
|
|
trainer.save_model("./weval-brain-v4-final")
|
|
print("DONE — upload to HF: huggingface-cli upload yace222/weval-brain-v4 ./weval-brain-v4-final")
|