68 lines
3.2 KiB
Python
68 lines
3.2 KiB
Python
|
"""C17d: 모든 풀이 + 길이 필터 (1500자 이하만) + NaN 방지"""
|
||
|
|
import json, re, random, torch, numpy as np, os
|
||
|
|
from collections import defaultdict
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
from trl import SFTTrainer, SFTConfig
|
||
|
|
from transformers import EarlyStoppingCallback
|
||
|
|
from datasets import Dataset
|
||
|
|
|
||
|
|
SEED = 42
|
||
|
|
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
torch.cuda.manual_seed_all(SEED)
|
||
|
|
if torch.cuda.get_device_capability()[0] >= 8: torch.set_float32_matmul_precision('high')
|
||
|
|
|
||
|
|
SP = "주어진 수학 문제를 단계별로 풀고 답변을 작성하세요.\n반드시 최종 답변을 \\boxed{정수} 형식으로 마지막 줄에 출력하세요.\n예시: \\boxed{42}"
|
||
|
|
|
||
|
|
print("=== C17d: All solutions, length-filtered (≤1500 chars) ===")
|
||
|
|
|
||
|
|
with open("data/GSM8K_full_qwen3_30b.json") as f:
|
||
|
|
data = json.load(f)
|
||
|
|
|
||
|
|
# 길이 필터: 1500자 이하만
|
||
|
|
filtered = [d for d in data if len(d['answer']) <= 1500]
|
||
|
|
print(f"원본: {len(data)}개 → 필터 후: {len(filtered)}개 (제거: {len(data)-len(filtered)})")
|
||
|
|
|
||
|
|
random.shuffle(filtered)
|
||
|
|
uq = len(set(d["question"] for d in filtered))
|
||
|
|
print(f"Unique: {uq}, avg {len(filtered)/uq:.1f}/q")
|
||
|
|
|
||
|
|
split = int(len(filtered) * 0.95)
|
||
|
|
train, test = filtered[:split], filtered[split:]
|
||
|
|
def to_sft(ex):
|
||
|
|
return {"prompt": [{"role":"user","content":SP+"\n\n"+ex["question"]}],
|
||
|
|
"completion": [{"role":"assistant","content":ex["answer"]}]}
|
||
|
|
|
||
|
|
cols = [c for c in Dataset.from_list(train[:1]).column_names if c not in ["prompt","completion"]]
|
||
|
|
train_ds = Dataset.from_list(train).map(to_sft, remove_columns=cols)
|
||
|
|
test_ds = Dataset.from_list(test).map(to_sft, remove_columns=cols)
|
||
|
|
print(f"학습: {len(train_ds)} / 검증: {len(test_ds)}")
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("outputs/models/gemma-3-1b-it")
|
||
|
|
model = AutoModelForCausalLM.from_pretrained("outputs/models/gemma-3-1b-it", dtype=torch.bfloat16, device_map="auto", attn_implementation='flash_attention_2')
|
||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
||
|
|
model.gradient_checkpointing_enable(); model.config.use_cache = False
|
||
|
|
|
||
|
|
cfg = SFTConfig(
|
||
|
|
report_to='none', seed=SEED, eval_strategy="steps", eval_steps=200,
|
||
|
|
save_total_limit=2, load_best_model_at_end=True, metric_for_best_model="eval_loss",
|
||
|
|
save_steps=200, num_train_epochs=3, warmup_ratio=0.05, weight_decay=0.01, max_grad_norm=1.0,
|
||
|
|
neftune_noise_alpha=5, per_device_train_batch_size=8, gradient_accumulation_steps=4,
|
||
|
|
per_device_eval_batch_size=2, max_length=2048, lr_scheduler_type='cosine',
|
||
|
|
learning_rate=2e-5, bf16=True, optim="paged_adamw_8bit",
|
||
|
|
output_dir="outputs/c17d_checkpoints", logging_steps=50, save_strategy="steps",
|
||
|
|
)
|
||
|
|
|
||
|
|
trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=train_ds, eval_dataset=test_ds, args=cfg,
|
||
|
|
callbacks=[EarlyStoppingCallback(early_stopping_patience=3)])
|
||
|
|
print("학습 시작 (3 epochs, 모든 풀이, ≤1500자)")
|
||
|
|
r = trainer.train()
|
||
|
|
print(f"완료! Loss: {r.training_loss:.4f}")
|
||
|
|
|
||
|
|
SAVE = "outputs/models/c17d-gemma-3-1b-it-Math"
|
||
|
|
os.makedirs(SAVE, exist_ok=True)
|
||
|
|
model.eval(); model.save_pretrained(SAVE, safe_serialization=False); tokenizer.save_pretrained(SAVE)
|
||
|
|
print(f"저장: {SAVE}")
|
||
|
|
del model, trainer; torch.cuda.empty_cache()
|
||
|
|
print("GPU 해제")
|