GPT-2는 오픈 소스이기 때문에 누구나 쉽게 접근할 수 있으며, 특정 데이터셋을 학습하여 원하는 결과를 도출할 수 있습니다.
이번 포스트에서는 GPT-2 모델을 활용하여 위키데이터(WikiText-2)를 학습하고 배포하는 방법을 단계별로 소개하겠습니다.
🔍 AutoModelForCausalLM이란?
AutoModelForCausalLM은 Hugging Face의 transformers 라이브러리에서 제공하는 자동 모델 로더로, Causal Language Modeling(인과적 언어 모델링, 자기회귀 언어 모델)을 수행하는 사전 훈련된 모델을 불러오는 클래스입니다.
이 모델은 이전 단어들을 기반으로 다음 단어를 예측하는 방식으로 학습되며, 대표적인 예시로 GPT-2, GPT-3, GPT-4, LLaMA, Mistral 등이 있습니다.
1. 필요한 라이브러리 설치
우선, Hugging Face의 datasets 라이브러리를 설치합니다. 이 라이브러리는 다양한 공개 데이터셋을 쉽게 로드할 수 있도록 도와줍니다.
!pip install datasets
2. 데이터 로드
이제 WikiText-2 데이터셋을 불러옵니다. Hugging Face의 load_dataset 함수를 이용하면 간편하게 사용할 수 있습니다.
from datasets import load_dataset
dataset = load_dataset("wikitext", name="wikitext-2-raw-v1")
WikiText-2는 위키피디아 기반의 대규모 텍스트 데이터셋으로, 자연어 모델 학습에 적합한 데이터셋입니다.
3. GPT-2 모델 불러오기
사전 훈련된 GPT-2 모델과 토크나이저를 로드합니다.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = 'gpt2'
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
GPT-2에는 기본적으로 pad_token이 없기 때문에 eos_token을 pad_token으로 설정합니다.
tokenizer.pad_token = tokenizer.eos_token
4. 토크나이징 함수 정의
데이터를 모델이 학습할 수 있는 형태로 변환하는 토크나이징(tokenization) 함수를 정의합니다.
GPT-2는 자기 자신을 예측하는 방식(자기회귀 모델)으로 학습되기 때문에, 출력 라벨(labels)도 입력값(input_ids)과 동일하게 설정합니다.
def tokenize_function(examples):
token = tokenizer(examples['text'], padding="max_length", truncation=True, max_length=512)
return {'input_ids': token['input_ids'], 'attention_mask': token['attention_mask'], 'labels': token['input_ids']}
5. 학습 및 테스트 데이터 분리
데이터셋을 학습용(train)과 테스트용(test)으로 나눕니다.
train_dataset = dataset['train']
test_dataset = dataset['test']
이제 토크나이징 함수를 적용하고, 원본 데이터셋에서는 더 쓰지 않을 text 컬럼을 제거합시다.
train_dataset = train_dataset.map(tokenize_function, batched=True, remove_columns=['text'])
test_dataset = test_dataset.map(tokenize_function, batched=True, remove_columns=['text'])
6. 학습 설정 및 Trainer 생성
Hugging Face의 Trainer를 사용하여 모델을 학습할 수 있습니다. 우선, 학습 파라미터를 설정합니다.
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir='./result', # 모델 저장 경로
evaluation_strategy='epoch',
num_train_epochs=2,
weight_decay=0.01, # 과적합 방지
per_device_train_batch_size=8, #한 개의 GPU에서 한 번의 학습 스텝(Iteration) 동안 처리할 샘플 개수
per_device_eval_batch_size=8,
logging_dir='./logs',
save_strategy="epoch",
)
Trainer 객체를 생성하여 학습을 준비합니다.
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset
)
7. 모델 학습 시작
이제 학습을 시작합니다!
trainer.train()
🔍 학습 과정 이해하기
- Training Loss: 학습 데이터에서의 오차
- Validation Loss: 테스트 데이터에서의 오차
- 에포크(epoch)가 늘어날수록 Training Loss는 점점 0에 가까워짐
- Validation Loss는 처음엔 줄어들다가 어느 순간 증가할 수 있음 → 과적합(Overfitting) 발생 가능
8. 학습된 모델로 텍스트 생성하기
이제 학습된 모델을 이용하여 텍스트를 생성해봅시다.
prompt = "The future of AI is"
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
outputs = model.generate(**inputs)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)
9. 모델 서비스화 과정
최적의 체크포인트 선택
training_args에서 results 폴더에 아웃풋 결과가 저장되도록 설정한 상태입니다.
에포크 1번: checkpoint-125 (125 * 8개 씩 학습 = 1000개)
에포크 2번: checkpoint-250 (250 * 8개 씩 학습 = 2000개)
...
에포크가 진행될수록 Training Loss는 0에 가까워지지만,
Validation Loss는 처음에는 감소하다가 일정 지점 이후 증가하는 경향을 보입니다.
이는 모델이 훈련 데이터에 과도하게 최적화되어 새로운 데이터에 대한 일반화 성능이 저하되는 '오버피팅(Overfitting)' 상태에 접어들었음을 의미합니다.
따라서 Validation Loss가 가장 낮았던 시점의 체크포인트(checkpoint)를 선택하여 최적의 모델을 저장하는 것이 중요합니다.
results/checkpoint-XXX 폴더에서 Validation Loss가 최소였던 에포크의 체크포인트를 선택하여 모델을 저장합시다.
Hugging Face에 업로드
from huggingface_hub import HfApi, notebook_login
notebook_login()
my_model_name = 'marurun66/gpt2-wiki' # 모델 이름
api = HfApi()
api.create_repo(repo_id=my_model_name, private=True, exist_ok=True) #덮어쓰기 OK
api.upload_folder(
folder_path='./results/checkpoint-125', # 최적의 체크포인트 선택
repo_id=my_model_name,
repo_type='model',
ignore_patterns=['.git', '.gitattributes']
)
이제 학습된 모델을 Hugging Face에 업로드하여 쉽게 공유하거나 활용할 수 있습니다
10. 정리 및 결론
✔️ GPT-2 모델을 활용하여 WikiText-2 데이터셋을 학습하는 방법을 실습해보았습니다.
✔️ Hugging Face의 Trainer를 이용하면 쉽게 모델을 훈련하고 평가할 수 있습니다.
✔️ 학습한 모델을 이용하여 자연어 생성 테스트를 진행할 수 있습니다.
이제 같은 방법으로 다른 데이터셋을 활용하여 GPT-2를 학습시키고, 원하는 NLP 애플리케이션을 구축해보세요! 🚀