실습 5 — 작은 GPT 사전학습하기

import math
import torch
from transformers import PreTrainedTokenizerFast
from datasets import Dataset

# 실습 4와 똑같이 데이터를 만든다
tok = PreTrainedTokenizerFast(tokenizer_file="out_tok/tokenizer.json", pad_token="<pad>", eos_token="<eos>")
with open("shop_corpus.txt", encoding="utf-8") as f:
    lines = [line for line in f.read().splitlines() if line]
ds = Dataset.from_dict({"text": lines}).train_test_split(test_size=0.05, seed=42)
BLOCK = 128

def tokenize(batch):
    return tok([t + "<eos>" for t in batch["text"]])

def group(batch):
    ids = sum(batch["input_ids"], [])
    n = len(ids) // BLOCK * BLOCK
    return {"input_ids": [ids[i:i + BLOCK] for i in range(0, n, BLOCK)]}

lm_data = ds.map(tokenize, batched=True, remove_columns=["text"]).map(group, batched=True, remove_columns=["attention_mask"])
print("학습용 묶음:", len(lm_data["train"]), "/ 검증용 묶음:", len(lm_data["test"]))
print("GPU 사용:", torch.cuda.is_available())

from transformers import GPT2Config, GPT2LMHeadModel, set_seed

set_seed(42)     # 파이썬, NumPy, PyTorch 의 씨앗을 한꺼번에 고정한다
config = GPT2Config(
    vocab_size=len(tok),          # 우리 토크나이저의 어휘 크기
    n_positions=BLOCK,            # 한 번에 보는 최대 토큰 수
    n_embd=256,                   # 좌표 수
    n_layer=4,                    # 블록 수
    n_head=4,                     # 헤드 수
    bos_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id,
)
model = GPT2LMHeadModel(config)   # 설정대로 빈 모델을 만든다(가중치는 난수)
print(f"파라미터 수: {sum(p.numel() for p in model.parameters()):,}")

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

# mlm=False: 가리고 맞히기가 아니라 다음 토큰 맞히기. labels 를 input_ids 로 채워 준다
collator = DataCollatorForLanguageModeling(tok, mlm=False)
args = TrainingArguments(
    "out_pretrain_ckpt",
    per_device_train_batch_size=32,
    num_train_epochs=8,
    learning_rate=1e-3,
    warmup_steps=50,                          # 처음 50걸음은 학습률을 천천히 올린다
    eval_strategy="epoch",                    # 에포크마다 검증 손실을 잰다
    logging_steps=50,
    save_strategy="no",                       # 중간 저장은 하지 않는다
    report_to="none",
    seed=42,
    fp16=torch.cuda.is_available(),           # GPU 가 있을 때만 혼합 정밀도
    disable_tqdm=True,                        # 진행 막대를 끈다
)
trainer = Trainer(model=model, args=args, train_dataset=lm_data["train"],
                  eval_dataset=lm_data["test"], data_collator=collator)

before = trainer.evaluate()["eval_loss"]      # 학습 전 검증 손실
trainer.train()
after = trainer.evaluate()["eval_loss"]       # 학습 뒤 검증 손실
print(f"학습 전: 손실 {before:.3f}, 퍼플렉서티 {math.exp(before):.1f}")
print(f"학습 뒤: 손실 {after:.3f}, 퍼플렉서티 {math.exp(after):.2f}")

model = trainer.model.eval()
for start in ["[안내] 교환은", "[상품] 텀블러는", "[공지]"]:
    enc = tok(start, return_tensors="pt").to(model.device)     # 모델이 있는 장치(GPU/CPU)로
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=30, do_sample=False,
                             eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
    print(tok.decode(out[0], skip_special_tokens=True))

# 모델과 토크나이저를 한 폴더에 저장한다. 실습 6과 프로젝트에서 from_pretrained 로 불러온다
model.save_pretrained("out_pretrained")
tok.save_pretrained("out_pretrained")
print("저장: out_pretrained")