1. 데이터 로드
1.1 토크나이저 로드
from transformers import BertTokenizer
def get_tokenizer():
tokenizer_name = "bert-base-multilingual-cased"
tokenizer = BertTokenizer.form_pretrained(tokenizer_name)
return tokenizer
1.2 문자열 데이터를 토큰화 하는 함수와 토큰을 자연어로 바꾸는 두개의 함수 만들기
import torch
def collate_fn(batch):
'''
자연어 데이터 토큰화 적용 함수
'''
tokenizer = get_tokenizer()
# 각 문장 쌍을 토크나이징하고 토큰 ID로 변환
en_batch = [tokenizer.encode(example['en'], max_length=128, truncation=True, padding='max_length', return_tensors="pt") for example in batch]
ko_batch = [tokenizer.encode(example['ko'], max_length=128, truncation=True, padding='max_length', return_tensors="pt") for example in batch]
# 토크나이즈된 리스트를 텐서로 변환
en_batch = torch.cat(en_batch, dim=0)
ko_batch = torch.cat(ko_batch, dim=0)
return en_batch, ko_batch
def decode_sentences(tokenizer, tensor):
decoded_sentences = [tokenizer.decode(t, skip_special_tokens=True) for t in tensor]
return decoded_sentences
1.3 데이터로더 만드는 함수 정의
from datasets import load_from_disk
from torch.utils.data import DataLoader
def load_and_create_dataloaders(batch_size=32, use_subset=True, subset_size=100):
# 데이터셋 로드
load_path = '/usr/src/app/daconsrc/dataset_dict'
dataset = load_from_disk(load_path)
# 훈련 세트와 테스트 세트 분리
train_dataset = dataset['train']
test_dataset = dataset['test']
# 데이터셋의 일부만 사용하는 경우
if use_subset:
train_dataset = train_dataset.select(range(subset_size))
test_dataset = test_dataset.select(range(subset_size))
# DataLoader 설정
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=False)
return train_dataloader, test_dataloader
2. 모델 정의
2.1 인코더 클래스
- 임베딩층(드랍아웃 적용)
- GRU 층(드랍아웃 적용)
import math
import random
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
# Encoder 정의
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, hidden_dim, n_layers, dropout):
super().__init__()
self.embedding = nn.Embedding(input_dim, emb_dim)
self.dropout = nn.Dropout(dropout)
self.gru = nn.GRU(emb_dim, hidden_dim, n_layers, batch_first=True, dropout=dropout)
def forward(self, src):
embedded = self.dropout(self.embedding(src))
outputs, hidden = self.gru(embedded)
return outputs, hidden
2.2 디코더 클래스
- 임베딩층(드랍아웃 적용)
- GRU 층(드랍아웃 적용)
- 완전 연결 층 (Fully Connected Layer)
- SoftMax
# Decoder 정의
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, hidden_dim, n_layers, dropout):
super().__init__()
self.output_dim = output_dim
self.embedding = nn.Embedding(output_dim, emb_dim)
self.gru = nn.GRU(emb_dim, hidden_dim, n_layers, batch_first=True, dropout=dropout if n_layers > 1 else 0)
self.fc_out = nn.Linear(hidden_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, trg, hidden):
trg = trg.unsqueeze(1)
embedded = self.dropout(self.embedding(trg))
output, hidden = self.gru(embedded, hidden)
prediction = self.fc_out(output.squeeze(1))
prediction = F.softmax(prediction, dim=1)
return prediction, hidden
2.3 Seq2Seq 클래스
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, device, trg_vocab_size, tokens):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
self.sos_token = tokens['sos_token']
self.eos_token = tokens['eos_token']
self.trg_vocab_size = trg_vocab_size
def forward(self, src, trg=None, teacher_force_ratio=1.0, max_len=50):
# 인코더 실행, 컨텍스트 벡터 준비
_, hidden = self.encoder(src)
# 첫번째 토큰은 sos로 세팅
current_token = torch.tensor([self.sos_token] * src.size(0)).to(self.device)
# 최종 출력값 설정
batch_size = src.shape[0]
outputs = torch.zeros(batch_size, max_len, self.trg_vocab_size).to(self.device)
all_tokens = torch.zeros(batch_size, max_len).to(self.device)
for t in range(1, max_len):
# 디코더 실행
output, hidden = self.decoder(current_token, hidden)
# 출력값 저장
outputs[:, t, :] = output # outputs 갱신
top1 = output.argmax(1) # output에서 토큰 추출
all_tokens[:, t] = top1 # 토큰 저장
# 티처 포싱
if teacher_force_ratio > 0 and trg is not None and t < trg.size(1):
current_token = trg[:, t] if random.random() < teacher_force_ratio else top1
else:
current_token = top1
# eos 토큰 발견 시 종료
if (current_token == self.eos_token).all():
break
actual_lengths = (all_tokens != self.eos_token).sum(dim=1)
return outputs, actual_lengths
3. 모델 학습과 평가
3.1 모델 학습 함수
def train_model(model, train_dataloader, optimizer, criterion, clip, device):
model.train()
epoch_loss = 0
progress_bar = tqdm(train_dataloader, desc='Training', leave=False)
for src, trg in progress_bar:
src, trg = src.to(device), trg.to(device)
optimizer.zero_grad()
output, actual_lengths = model(src, trg)
output, trg = remove_special_tokens(output, trg, model.eos_token)
loss = criterion(output, trg)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(train_dataloader)
3.2 모델 평가 함수
def evaluate_model(model, test_dataloader, criterion, device):
model.eval()
epoch_loss = 0
progress_bar = tqdm(test_dataloader, desc='Evaluating', leave=False)
with torch.no_grad():
for src, trg in progress_bar:
src, trg = src.to(device), trg.to(device)
output, actual_lengths = model(src, trg, teacher_force_ratio=0)
output, trg = remove_special_tokens(output, trg, model.eos_token)
loss = criterion(output, trg)
epoch_loss += loss.item()
return epoch_loss / len(test_dataloader)
3.3 스페셜 토큰 제거 전처리 함수
def remove_special_tokens(output, trg, eos_token):
# output, trg의 sos 토큰 제거
output_dim = output.shape[-1]
max_len = trg.shape[1]
output = output[:, 1:max_len].reshape(-1, output_dim)
trg = trg[:, 1:max_len].reshape(-1)
# eos토큰 제거
trg_valid_idx = (trg != eos_token).nonzero(as_tuple=True)[0]
output_valid_idx = (output.argmax(1) != eos_token).nonzero(as_tuple=True)[0]
trg_valid_idx = trg_valid_idx[trg_valid_idx < output.size(0)] # 길이가 너무길면 제거
output_valid_idx = output_valid_idx[output_valid_idx < output.size(0)]
trg = trg[trg_valid_idx]
output = output[output_valid_idx]
# 출력
return output, trg
3.4 예측 함수
def predict(model, src, max_len, device):
# 문자열을 토큰화 한다.
tokenizer = get_tokenizer()
en_tokens = tokenizer.encode(src, max_length=128, truncation=True, padding='max_length', return_tensors="pt").to(device)
# 모델로 결과 예측
model.eval()
with torch.no_grad():
outputs, _ = model(en_tokens, trg=None, teacher_force_ratio=0, max_len=max_len)
predictions = outputs.argmax(2)
# 모델이 예측한 토큰을 문자로 만들어 준다.
decoded_sentences = [tokenizer.decode(t, skip_special_tokens=True) for t in predictions]
return decoded_sentences
4. 코드 실행
4.1 모델 학습
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
tokenizer = get_tokenizer()
# 모델의 하이퍼파라미터 정의
INPUT_DIM = len(tokenizer.vocab)
OUTPUT_DIM = len(tokenizer.vocab)
ENC_EMB_DIM = 32 # 인코더 임베딩 차원
DEC_EMB_DIM = 32 # 디코더 임베딩 차원
HIDDEN_DIM = 64 # GRU 히든 상태 차원
N_LAYERS = 2 # GRU 레이어 수
ENC_DROPOUT = 0.5 # 인코더 드롭아웃 비율
DEC_DROPOUT = 0.5 # 디코더 드롭아웃 비율
N_EPOCHS = 2
tokens = {
'sos_token': 0, # <sos> 토큰
'eos_token': 1 # <eos> 토큰
}
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_dataloader, test_dataloader = load_and_create_dataloaders(batch_size=4, use_subset=True, subset_size=64) # 50개의 샘플만 사용
endocer = Encoder(INPUT_DIM, ENC_EMB_DIM, HIDDEN_DIM, N_LAYERS, ENC_DROPOUT)
decoder = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HIDDEN_DIM, N_LAYERS, DEC_DROPOUT)
model = Seq2Seq(endocer, decoder, device, OUTPUT_DIM, tokens).to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
best_test_loss = float('inf')
for epoch in range(N_EPOCHS):
# 모델 학습
train_loss = train_model(model, train_dataloader, optimizer, criterion, 1.0, device)
# 평가
test_loss = evaluate_model(model, test_dataloader, criterion, device)
print(f'Epoch: {epoch+1:02}')
print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
print(f'\t Val. Loss: {test_loss:.3f} | Val. PPL: {math.exp(test_loss):7.3f}')
4.2 예측
src = '나는 학생입니다.'
max_len = 128
result = predict(model, src, max_len, device)
print(result)