간단한 SLM 학습 및 테스트 예제

Hugging Face의 Transformers 라이브러리를 사용하여 DistilBERT를 간단한 텍스트 분류 작업에 학습시키는 과정
필요한 라이브러리 설치

pip install transformers datasets torch

학습용 데이터 준비

Hugging Face의 datasets 라이브러리를 사용해 간단한 감정 분석 데이터셋(SST-2)을 불러옵니다.

from datasets import load_dataset
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
from transformers import Trainer, TrainingArguments
import torch

# 1. 데이터셋 로드
dataset = load_dataset("glue", "sst2")

# 2. 토크나이저 로드
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

# 3. 데이터 전처리
def tokenize_function(examples):
    return tokenizer(examples["sentence"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 4. PyTorch 데이터셋으로 변환
tokenized_datasets = tokenized_datasets.remove_columns(["sentence", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")

# 5. 학습 및 검증 데이터 분리
train_dataset = tokenized_datasets["train"]
valid_dataset = tokenized_datasets["validation"]

# 6. 모델 로드
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)

# 7. 학습 설정
training_args = TrainingArguments(
    output_dir="./results",          # 모델 저장 디렉토리
    evaluation_strategy="epoch",    # 평가 주기
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",            # 로그 저장 디렉토리
    logging_steps=10,
    save_steps=500,
)

# 8. Trainer 설정
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
)

# 9. 모델 학습
trainer.train()

# 10. 모델 저장
model.save_pretrained("./distilbert-sst2")
tokenizer.save_pretrained("./distilbert-sst2")

* 코드 설명

1. 데이터 로드 및 전처리: • SST-2는 감정 분석용 데이터셋입니다(긍정/부정 분류). • 문장을 토큰화하고 패딩/트렁케이션 처리하여 모델 입력 형식에 맞게 변환합니다. 2. 모델 설정: • DistilBERT 모델을 로드하고 출력 레이블 개수를 2(긍정/부정)로 설정합니다. 3. Trainer 사용: • Trainer는 학습 과정을 간단히 관리할 수 있는 Hugging Face의 유틸리티입니다. • 학습 파라미터(학습률, 배치 크기, 에포크 수 등)를 설정합니다. 4. 학습 실행: • trainer.train()으로 학습을 시작합니다. • 학습이 완료되면 모델과 토크나이저를 저장합니다.

* 결과

• 학습된 모델은 SST-2 데이터셋에서 긍정/부정 감정을 예측하는 데 사용할 수 있습니다. • 저장된 모델(./distilbert-sst2)을 로드하여 새로운 문장의 감정을 분류할 수 있습니다.

추가 활용

• 다른 데이터셋으로 모델을 Fine-tuning하거나, 텍스트 분류 외의 작업(질문 응답, 요약 등)에도 적용할 수 있습니다. • SLM(Small Language Model) 학습은 LLM에 비해 자원 소모가 적고, 로컬 컴퓨터에서도 실행이 가능합니다.

학습용 데이터 준비

Comments