ELECTRA 학습 구조 소개 및 Domain Adaptation 수행하기

 
from transformers import ElectraForPreTraining, ElectraTokenizer, ElectraForMaskedLM
 
tokenizer = ElectraTokenizer.from_pretrained("monologg/koelectra-base-v3-generator")
 
generator = ElectraForMaskedLM.from_pretrained('monologg/koelectra-base-v3-generator')
 
discriminator = ElectraForPreTraining.from_pretrained("monologg/koelectra-base-v3-discriminator")
 

 
# Huggingface Transformers 내부 ElectraForMaskedLM 코드
 
class ElectraForMaskedLM(ElectraPreTrainedModel):
    _keys_to_ignore_on_load_missing = ["generator_lm_head.weight"]
 
    def __init__(self, config):
        super().__init__(config)
 
        # ElectraForMaskedLM 내부에서 ElectraModel 모듈을 불러와 사용함.
        self.electra = ElectraModel(config)
 
        # [Mask] 토큰에 적합한 토큰을 예측하는 generator predictor가 포함되어 있음을 확인할 수 있음.
        self.generator_predictions = ElectraGeneratorPredictions(config)
 
        self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
 
        # Initialize weights and apply final processing
        self.post_init()
 
    def forward(...) :
 
            # electra를 통해 마지막 encoder의 output(=last_hidden_state)를 받음.
            generator_hidden_states = self.electra(
                input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                head_mask=head_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
            # last_hidden_states를 prediction scores에 input 데이터로 활용함.
            generator_sequence_output = generator_hidden_states[0]
 
            # last_hidden_states를 MLM prediction layer의 input 데이터로 활용함.
            prediction_scores = self.generator_predictions(generator_sequence_output)
            prediction_scores = self.generator_lm_head(prediction_scores)
 
    return MaskedLMOutput(
                loss=loss,
                # Prediction_score을 리턴
                # logits의 shape은 (batch_size, src_token_len, vocab_size)
                logits=prediction_scores,
                hidden_states=generator_hidden_states.hidden_states,
                attentions=generator_hidden_states.attentions,
            )
 

 
# Huggingface Transformers 내부 ElectraForPreTraining 코드
 
class ElectraForPreTraining(ElectraPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
 
        # ElectraForPreTraining 내부에서 ElectraModel 모듈을 불러와 사용함.
        self.electra = ElectraModel(config)
 
        # Token의 진위여부를 판별하는 Precdiction 모델을 불러옴.
        self.discriminator_predictions = ElectraDiscriminatorPredictions(config)
 
    ....
 
 
    def forward(...)
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
        # electra를 통해 마지막 encoder의 output(=last_hidden_state)를 받음.
        discriminator_hidden_states = self.electra(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
 
        # last_hidden_states를 classification layer의 input 데이터로 활용함.
        discriminator_sequence_output = discriminator_hidden_states[0]
 
        logits = self.discriminator_predictions(discriminator_sequence_output)
 
        return ElectraForPreTrainingOutput(
            loss=loss,
            # logits을 리턴
            # logits의 shape은 (batch_size, src_token_len)
            logits=logits,
            hidden_states=discriminator_hidden_states.hidden_states,
            attentions=discriminator_hidden_states.attentions,
        )

discriminator = ElectraForPreTraining.from_pretrained('...')
 
# Electra Model 추출하기
trained_electra = discriminator.electra

import math
from functools import reduce
from collections import namedtuple
 
import torch
from torch import nn
import torch.nn.functional as F
 
# constants
 
Results = namedtuple(
    "Results",
    [
        "loss",
        "mlm_loss",
        "disc_loss",
        "gen_acc",
        "disc_acc",
        "disc_labels",
        "disc_predictions",
    ],
)
 
# 모델 내부에서 활용되는 함수 정의
 
def log(t, eps=1e-9):
    return torch.log(t + eps)
 
def gumbel_noise(t):
    noise = torch.zeros_like(t).uniform_(0, 1)
    return -log(-log(noise))
 
def gumbel_sample(t, temperature=1.0):
    return ((t / temperature) + gumbel_noise(t)).argmax(dim=-1)
 
def prob_mask_like(t, prob):
    return torch.zeros_like(t).float().uniform_(0, 1) < prob
 
def mask_with_tokens(t, token_ids):
    init_no_mask = torch.full_like(t, False, dtype=torch.bool)
    mask = reduce(lambda acc, el: acc | (t == el), token_ids, init_no_mask)
    return mask
 
def get_mask_subset_with_prob(mask, prob):
    batch, seq_len, device = *mask.shape, mask.device
    max_masked = math.ceil(prob * seq_len)
 
    num_tokens = mask.sum(dim=-1, keepdim=True)
    mask_excess = mask.cumsum(dim=-1) > (num_tokens * prob).ceil()
    mask_excess = mask_excess[:, :max_masked]
 
    rand = torch.rand((batch, seq_len), device=device).masked_fill(~mask, -1e9)
    _, sampled_indices = rand.topk(max_masked, dim=-1)
    sampled_indices = (sampled_indices + 1).masked_fill_(mask_excess, 0)
 
    new_mask = torch.zeros((batch, seq_len + 1), device=device)
    new_mask.scatter_(-1, sampled_indices, 1)
    return new_mask[:, 1:].bool()
 
# main electra class
 
class Electra(nn.Module):
    def __init__(
        self,
        generator,
        discriminator,
        tokenizer,
        *,
        num_tokens=35000,
        mask_prob=0.15,
        replace_prob=0.85,
        mask_token_id=4,
        pad_token_id=0,
        mask_ignore_token_ids=[2, 3],
        disc_weight=50.0,
        gen_weight=1.0,
        temperature=1.0,
    ):
        super().__init__()
 
        """
        num_tokens: 모델 vocab_size
        mask_prob: 토큰 중 [MASK] 토큰으로 대체되는 비율
        replace_prop:  토큰 중 [MASK] 토큰으로 대체되는 비율
        mask_token_i: [MASK] Token id
        pad_token_i: [PAD] Token id
        mask_ignore_token_id: [CLS],[SEP] Token id
        disc_weigh: discriminator loss의 Weight 조정을 위한 값
        gen_weigh: generator loss의 Weight 조정을 위한 값
        temperature: gumbel_distribution에 활용되는 arg, 값이 높을수록 모집단 분포와 유사한 sampling 수행
        """
 
        # Generator, Discriminator, Tokenizer
        self.generator = generator
        self.discriminator = discriminator
        self.tokenizer = tokenizer
 
        # mlm related probabilities
        self.mask_prob = mask_prob # 0.15
        self.replace_prob = replace_prob # 0.85
 
        self.num_tokens = num_tokens # 35000
 
        # token ids
        self.pad_token_id = pad_token_id
        self.mask_token_id = mask_token_id
        self.mask_ignore_token_ids = set([*mask_ignore_token_ids, pad_token_id])
 
        # sampling temperature
        self.temperature = temperature
 
        # loss weights
        # Discriminor Weight이 50인 이유는 오차에 대한 Loss가 Generator 보다 작기 때문
        # 따라서 적절한 학습을 수행하기 위해 Weight을 통해 Discriminator와 Generator의 차이를 줄여줌
        self.disc_weight = disc_weight # 50.0
        self.gen_weight = gen_weight # 1.0
 
    def forward(self, input_ids, **kwargs):
 
        input = input_ids["input_ids"]
 
        # ------ 1단계 Input Data Masking --------#
 
        """
        - Generator는 Bert와 구조도 동일하고 학습하는 방법도 동일함.
 
        - Generator 학습을 위해선 [Masked] 토큰이 필요하므로 input data를 Masking하는 과정이 필요함.
 
        """
 
        replace_prob = prob_mask_like(input, self.replace_prob)
 
        # do not mask [pad] tokens, or any other tokens in the tokens designated to be excluded ([cls], [sep])
        # also do not include these special tokens in the tokens chosen at random
        no_mask = mask_with_tokens(input, self.mask_ignore_token_ids)
        mask = get_mask_subset_with_prob(~no_mask, self.mask_prob)
 
        # get mask indices
        mask_indices = torch.nonzero(mask, as_tuple=True)
 
        # mask input with mask tokens with probability of `replace_prob` (keep tokens the same with probability 1 - replace_prob)
        masked_input = input.clone().detach()
 
        # set inverse of mask to padding tokens for labels
        gen_labels = input.masked_fill(~mask, self.pad_token_id)
 
        # clone the mask, for potential modification if random tokens are involved
        # not to be mistakened for the mask above, which is for all tokens, whether not replaced nor replaced with random tokens
        masking_mask = mask.clone()
 
        # [mask] input
        masked_input = masked_input.masked_fill(
            masking_mask * replace_prob, self.mask_token_id
        )
 
        # ------ 2단계 Masking 된 문장을 Generator가 학습하고 가짜 Token을 생성 --------#
 
        """
        - Generator를 학습하여 MLM_loss 계산(combined_loss 계산에 활용)
        - Generator에서 예측한 문장을 Discriminator 학습에 활용
        - ex) 원본 문장 : 특히 안드로이드 플랫폼 기반의 (웹)앱과 (하이)브드리앱에 초점을 맞추고 있다
              가짜 문장 : 특히 안드로이드 플랫폼 기반의 (마이크로)앱과 (이)브드리앱에 초점을 맞추고 있다
        """
 
        # get generator output and get mlm loss(수정)
        logits = self.generator(masked_input, **kwargs).logits
 
        mlm_loss = F.cross_entropy(
            logits.transpose(1, 2), gen_labels, ignore_index=self.pad_token_id
        )
 
        # use mask from before to select logits that need sampling
        sample_logits = logits[mask_indices]
 
        # sample
        sampled = gumbel_sample(sample_logits, temperature=self.temperature)
 
        # scatter the sampled values back to the input
        disc_input = input.clone()
        disc_input[mask_indices] = sampled.detach()
 
        # generate discriminator labels, with replaced as True and original as False
        disc_labels = (input != disc_input).float().detach()
 
        # ------ 3단계 Token의 진위여부를 Discriminator가 예측하고 이를 학습 --------#
 
        """
        - 가짜 문장을 학습해 개별 토큰에 대해 진위여부를 판단
        - 진짜 token이라 판단하면 0, 가짜 토큰이라 판단하면 1을 부여
        - 정답과 비교해 disc_loss를 계산(combined_loss 계산에 활용)
        - combined_loss : 학습의 최종 loss임. 모델은 combined_loss의 최솟값을 얻기 위한 방식으로 학습 진행
        """
 
        # get discriminator predictions of replaced / original
        non_padded_indices = torch.nonzero(input != self.pad_token_id, as_tuple=True)
 
        # get discriminator output and binary cross entropy loss
        disc_logits = self.discriminator(disc_input, **kwargs).logits
        disc_logits_reshape = disc_logits.reshape_as(disc_labels)
 
        disc_loss = F.binary_cross_entropy_with_logits(
            disc_logits_reshape[non_padded_indices], disc_labels[non_padded_indices]
        )
 
        combined_loss = (self.gen_weight * mlm_loss + self.disc_weight * disc_loss,)
 
        # ------ 모델 성능 및 학습 과정을 추적하기 위한 지표(Metrics) 설계(선택사항) --------#
 
        with torch.no_grad():
            # gen mask 예측
            gen_predictions = torch.argmax(logits, dim=-1)
 
            # fake token 진위 예측
            disc_predictions = torch.round(
                (torch.sign(disc_logits_reshape) + 1.0) * 0.5
            )
            # generator_accuracy
            gen_acc = (gen_labels[mask] == gen_predictions[mask]).float().mean()
            # discriminator_accuracy
            disc_acc = (
                0.5 * (disc_labels[mask] == disc_predictions[mask]).float().mean()
                + 0.5 * (disc_labels[~mask] == disc_predictions[~mask]).float().mean()
            )
 
 
        return Results(
            combined_loss,
            mlm_loss,
            disc_loss,
            gen_acc,
            disc_acc,
            disc_labels,
            disc_predictions,
        )

from datasets import load_dataset
 
# local file을 불러오고 싶을땐 '확장자명', '경로' 적기
train = load_dataset('csv',data_files='data/book_train_128.csv')
validation = load_dataset('csv',data_files='data/book_validation_128.csv')
 
# tokenizing 방법 정의
def tokenize_function(examples):
    return tokenizer(examples['sen'], max_length=128, padding=True, truncation=True)
 
# datasets의 map 매서드를 활용해 모든 문장에 대한 토크나이징 수행
train_data_set = train['train'].map(tokenize_function,batch_size=True)
validation_data_set = validation['train'].map(tokenize_function,batch_size=True)

train = load_dataset('csv',data_files='data/book_train_128.csv')
 
print(train)
 
>>> DatasetDict({
    train: Dataset({
        features: ['Unnamed: 0', 'sen'],
        num_rows: 175900
    })
})
 
#----------
 
# column 제거
train = train.remove_columns('Unnamed: 0')
 
print(train)
 
>>> DatasetDict({
    train: Dataset({
        features: ['sen'],
        num_rows: 175900
    })
})
 
#----------
 
# train 데이터셋으로 이동
train_data_set = train['train']
 
print(train_data_set)
 
>>> Dataset({
    features: ['Unnamed: 0', 'sen'],
    num_rows: 175900
})
 
#----------
 
# 데이터 불러오기
print(train_data_set[0])
 
>> {'sen': '이 책의 특징ㆍ코딩의 기초 기초수학 논리의 ... 기초수학'}
 
#----------
 
# 데이터 추출
print(train_data_set[0]['sen'])
 
>>> '이 책의 특징ㆍ코딩의 기초 기초수학 논리의 ... 기초수학'
 
#----------
 
# type 확인
print(train_data_set.feature)
 
>>> {'sen': Value(dtype='string', id=None)}
 
#----------
 
# 저장
print(train_data_set.to_csv(''))
 

from transformers import TrainingArguments
 
training_args = TrainingArguments(
    output_dir="test_trainer",
    per_device_eval_batch_size=8,
    per_device_train_batch_size=8,
    logging_steps=2,
    num_train_epochs=2,
    eval_steps = 100
    evaluation_strategy='steps'
)

from transformers import TrainerCallback
 
# custom callback 만들기, 이때 TrainerCallback을 상속 받아야함.
class myCallback(TrainerCallback):
 
  # step이 시작할때 아래 코드 실행
  def on_step_begin(self, args, state, control, logs=None, **kwargs):
 
      if state.global_step % args.logging_steps == 0:
        # state는 현재 step, epoch 등 진행 상태에 대한 값을 불러옴
        # arg는 훈련 옵션으로 설정한 값을 불러옴.
          print("")
          print(
              f"{int(state.epoch)}번째 epoch 진행 중 --- {state.global_step}번째 step 결과"
          )

0번째 epoch 진행 중 ------- 0번째 step 결과
input 문장 : 장 수학 기호 수식에 많이 쓰이는 그리스 알 [MASK] [MASK] [MASK] 읽고 쓰는 법을 배웁니다
output 문장 : 장 수학 기호 수식에 많이 쓰이는 그리스 알 [##파] [##벳] [##을]을 쓰는 법을 배웁니다
 
0번째 epoch 진행 중 ------- 20번째 step 결과
input 문장 : [MASK]이 출간된지 꽤 됬다고 생각하는데 실습하는데 전혀 [MASK]없습니다
output 문장 : [책]이 출간된지 꽤 됬다고 생각하는데 실습하는데 전혀 [문제]없습니다
 

class customtrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
 
    # compute_loss 함수 덮어쓰기
    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer.
 
        By default, all models return the loss in the first element.
 
        Subclass and override for custom behavior.
        """
        if self.label_smoother is not None and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        prds,outputs = model(inputs)
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]
 
        if labels is not None:
            loss = self.label_smoother(outputs, labels)
        else:
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
 
        ############## 모델 학습 과정 확인을 위한 코드 추가
 
        if self.state.global_step % self.args.logging_steps == 0:
            # self.state.global_step = 현 step 파악
            # args.logging_steps = argument에서 지정한 logging_step
 
            # batch 중 0 번째 위치한 문장 선택
            num = 1
            input_id = inputs.input_ids[num].reshape(-1).data.tolist()
            output_id = outputs.logits[num].argmax(dim=-1).reshape(-1).data.tolist()
            attention_mask = inputs.attention_mask[num]
 
            # mask가 위치한 idx 추출하기
            mask_idx = (inputs.input_ids[num] == 4).nonzero().data.reshape(-1).tolist()
 
            # padding 제거
            input_id_without_pad = [
                input_id[i] for i in range(len(input_id)) if attention_mask[i]
            ]
            output_id_without_pad = [
                output_id[i] for i in range(len(output_id)) if attention_mask[i]
            ]
 
            # id to token
            # [1:-1] [CLS,SEP] 제거
            inputs_tokens = self.tokenizer.convert_ids_to_tokens(input_id_without_pad)[
                1:-1
            ]
            outputs_tokens = self.tokenizer.convert_ids_to_tokens(
                output_id_without_pad
            )[1:-1]
 
            # output mask 부분 표시하기
            for i in mask_idx:
                # [CLS,SEP 위치 조정]
                outputs_tokens[i - 1] = "[" + outputs_tokens[i - 1] + "]"
 
            inputs_sen = self.tokenizer.convert_tokens_to_string(inputs_tokens)
            outputs_sen = self.tokenizer.convert_tokens_to_string(outputs_tokens)
 
            print(f"input 문장 : {''.join(inputs_sen)}")
            print(f"output 문장 : {''.join(outputs_sen)}")
 
            # input 문장 : 수학 기호 수식에 많이 쓰이는 그리스 알 [MASK] [MASK] [MASK] 읽고 쓰는 법을 배웁니다
            # output 문장 : 수학 기호 수식에 많이 쓰이는 그리스 알 [##파] [##벳] [##을]을 쓰는 법을 배웁니다
 
        return (loss, outputs) if return_outputs else loss

trainer = customtrainer(
    model=model.to(device),
    train_dataset=train_data_set,
    eval_dataset=validation_data_set,
    args=training_args,
    tokenizer=tokenizer,
    callbacks=[myCallback],
)
 
trainer.train()