실습 2 — 작은 번역기를 만들고 처음부터 학습하기

# MarianConfig = Marian 번역 모델의 설정표, MarianMTModel = Marian machine translation model(Encoder-Decoder 번역 모델)
from transformers import MarianConfig, MarianMTModel  # 작은 Encoder-Decoder Transformer를 만듭니다.

translation_config = MarianConfig(  # 모델 크기와 특수 토큰을 정합니다.
    # len(tokenizer) = 61,790. 단어 표의 행 수이자 출력 Linear 의 칸 수
    vocab_size=len(translation_tokenizer),  # 토크나이저의 전체 토큰 수입니다.
    # d_model = 토큰 하나를 나타내는 숫자 개수(열 수). 원조 Transformer 는 512
    d_model=256,  # 토큰 벡터의 크기입니다.
    # 원조 Transformer 는 6층씩. 여기서는 빨리 돌도록 2층씩
    encoder_layers=2, decoder_layers=2,  # Encoder·Decoder를 두 층씩 쌓습니다.
    # 헤드 4개 → 헤드 하나가 256 ÷ 4 = 64칸씩 맡는다(√64 = 8 로 나눈다)
    encoder_attention_heads=4, decoder_attention_heads=4,  # Attention을 네 관점으로 나눕니다.
    # FFN: 256 → 512 → 256 으로 넓혔다가 줄인다
    encoder_ffn_dim=512, decoder_ffn_dim=512,  # FFN의 중간 벡터 크기입니다.
    max_position_embeddings=MAX_LENGTH,  # 모델이 처리할 입력과 번역문의 최대 길이를 토큰 수로 지정합니다.
    pad_token_id=pad_id,  # 문장 길이를 맞추기 위해 빈자리에 넣을 토큰의 ID입니다.
    eos_token_id=eos_id,  # 번역이 끝났음을 알리는 토큰의 ID입니다.
    # decoder_start_token_id: 정답을 한 칸 밀 때 맨 앞에 넣는 <시작> 토큰. 여기서는 [CLS](2)
    decoder_start_token_id=bos_id,  # Decoder가 번역을 시작할 때 맨 앞에 넣는 토큰의 ID입니다.
)
# 설정표만 주고 만들면 가중치가 무작위 숫자로 채워진다(아무것도 모르는 모델)
translation_model = MarianMTModel(translation_config)  # 사전학습 가중치 없이 설정값만으로 초기화합니다.
# p.numel() = number of elements(행렬 하나에 든 숫자 개수). 모두 더하면 전체 파라미터 수
# :, = 세 자리마다 쉼표. 결과: 파라미터 수: 18,470,400
print(f"파라미터 수: {sum(p.numel() for p in translation_model.parameters()):,}")

# DataCollatorForSeq2Seq = 묶음(batch) 정리 담당, Seq2SeqTrainingArguments = 학습 조건표, Seq2SeqTrainer = 학습 반복 담당
from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer  # 문장 길이를 맞추고 번역 모델을 학습할 도구를 불러옵니다.

# model= 을 주면 모델의 shift 규칙으로 labels 를 한 칸 밀어 decoder_input_ids 를 만든다(Teacher Forcing)
# labels 의 빈칸은 -100(채점에서 빼는 칸)으로, input_ids 의 빈칸은 [PAD] 로 채운다
data_collator = DataCollatorForSeq2Seq(translation_tokenizer, model=translation_model)  # 문장 길이를 맞추고, 정답 문장을 한 칸 밀어 Decoder 입력을 만듭니다.

translation_args = Seq2SeqTrainingArguments(  # 반복 횟수, 배치 크기, 학습률과 검증 주기 등을 정합니다.
    output_dir="outputs/translation-scratch",  # 학습 결과가 저장될 폴더를 지정합니다.
    # epoch = 문제집 한 바퀴. 10바퀴 × 238단계 = 2,380단계
    num_train_epochs=10,  # 전체 학습 데이터를 열 번 반복합니다.
    per_device_train_batch_size=32,  # 한 번의 학습 단계에서 문장쌍 32개를 처리합니다.
    per_device_eval_batch_size=32,  # 한 번의 검증 단계에서 문장쌍 32개를 처리합니다.
    # 5e-4 = 0.0005. 한 걸음의 보폭. 처음부터 배우는 작은 모델이라 파인튜닝보다 크게 잡았다
    learning_rate=5e-4,  # 학습할 때 가중치를 얼마나 크게 바꿀지 정하는 학습률입니다.
    # "epoch": 한 바퀴 돌 때마다 검증 데이터(949쌍)로 모의고사를 본다
    eval_strategy="epoch",  # epoch마다 검증 손실을 확인합니다.
    logging_steps=50,  # 학습 50단계마다 손실 등을 기록합니다.
    report_to="none",  # 학습 기록을 외부 서비스로 보내지 않습니다.
    save_strategy="no",  # 학습 중간에는 모델을 저장하지 않습니다.
)

translation_trainer = Seq2SeqTrainer(  # 번역 모델, 데이터, 학습 조건을 Trainer에 연결합니다.
    model=translation_model,  # 위에서 만든 번역 모델의 가중치를 학습합니다.
    args=translation_args,  # 앞에서 정한 배치 크기와 학습률 등의 조건을 적용합니다.
    train_dataset=tokenized["train"],  # 이 데이터로 모델의 가중치를 업데이트합니다.
    eval_dataset=tokenized["validation"],  # 이 데이터로 학습 중 번역 성능을 확인합니다.
    data_collator=data_collator,  # 길이가 다른 문장을 패딩하고 Decoder 입력을 준비합니다.
    # processing_class: 입력을 처리하는 도구(여기서는 토크나이저)를 Trainer 에 알려 준다
    processing_class=translation_tokenizer,  # 영어·한국어 문장을 토큰 ID로 바꾸는 토크나이저를 전달합니다.
)
# train() = 학습 시작. 결과: epoch 별 학습 손실·검증 손실 표와 TrainOutput(걸린 시간 등)
# 이 PC(RTX 3080)에서 약 52초. 마지막 검증 손실 약 3.19
translation_trainer.train()