load_model.py

import logging
import os
import math
import pandas as pd
import numpy as np
import argparse
import json
import gc
import re
import copy
import random
from tqdm import tqdm


from typing import Dict, List, Optional, Tuple

from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from transformers import PreTrainedTokenizerFast, LEDForConditionalGeneration, AutoModel
from transformers import BartForConditionalGeneration, BartConfig
from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
from transformers.models.longformer.modeling_longformer import LongformerSelfAttention
from transformers import get_linear_schedule_with_warmup, AdamW, TrainingArguments

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import CrossEntropyLoss, MSELoss
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset


# Kobart의 attention layer를 대체
class LongformerSelfAttentionForBart(nn.Module):
    def __init__(self, config : dict , layer_id : int):
        super().__init__()
        self.embed_dim = config.d_model
        self.longformer_self_attn = LongformerSelfAttention(config, layer_id=layer_id)
        self.output = nn.Linear(self.embed_dim, self.embed_dim)

    # kobart의 기존 layer와 동일한 형태의 입력을 받고, 동일한 형태의 출력을 할 수 있도록 해줘야함.
    def forward(self,
        hidden_states: torch.Tensor,
        key_value_states: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

        is_cross_attention = key_value_states is not None
        bsz, tgt_len, embed_dim = hidden_states.size()

        # bs x seq_len x seq_len -> bs x seq_len 으로 변경
        attention_mask = attention_mask.squeeze(dim=1)
        attention_mask = attention_mask[:,0]

        is_index_masked = attention_mask < 0
        is_index_global_attn = attention_mask > 0
        is_global_attn = is_index_global_attn.flatten().any().item()

        outputs = self.longformer_self_attn(
            hidden_states,
            attention_mask=attention_mask,
            layer_head_mask=None,
            is_index_masked=is_index_masked,
            is_index_global_attn=is_index_global_attn,
            is_global_attn=is_global_attn,
            output_attentions=output_attentions,
        )

        attn_output = self.output(outputs[0])

        return (attn_output,) + outputs[1:] if len(outputs) == 2 else (attn_output, None, None)

class LongformerBartForConditionalGeneration(BartForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        if config.attention_mode == 'n2':
            pass  # do nothing, use BertSelfAttention instead
        else:
            self.model.encoder.embed_positions = BartLearnedPositionalEmbedding(
                config.max_encoder_position_embeddings, 
                config.d_model, 
                config.pad_token_id)

            self.model.decoder.embed_positions = BartLearnedPositionalEmbedding(
                config.max_decoder_position_embeddings, 
                config.d_model, 
                config.pad_token_id)
            
            for i, layer in enumerate(self.model.encoder.layers):
                layer.self_attn = LongformerSelfAttentionForBart(config, layer_id=i)

#longformer bart모델의 config 생성 class
class LongformerBartConfig(BartConfig):
    def __init__(self, attention_window: List[int] = [512], attention_dilation: List[int] = [1],
                 autoregressive: bool = False, attention_mode: str = 'sliding_chunks',
                 gradient_checkpointing: bool = False, max_seq_len: int = 4096, max_pos: int = 4104,  **kwargs):
        """
        Args:
            attention_window: list of attention window sizes of length = number of layers.
                window size = number of attention locations on each side.
                For an affective window size of 512, use `attention_window=[256]*num_layers`
                which is 256 on each side.
            attention_dilation: list of attention dilation of length = number of layers.
                attention dilation of `1` means no dilation.
            autoregressive: do autoregressive attention or have attention of both sides
            attention_mode: 'n2' for regular n^2 self-attention, 'tvm' for TVM implemenation of Longformer
                selfattention, 'sliding_chunks' for another implementation of Longformer selfattention
        """

        super().__init__(**kwargs)
        self.attention_window = attention_window
        self.attention_dilation = attention_dilation
        self.autoregressive = autoregressive
        self.attention_mode = attention_mode
        self.gradient_checkpointing = gradient_checkpointing
        assert self.attention_mode in ['tvm', 'sliding_chunks', 'n2']
        
if __name__ == '__main__':
    # Longformer weight 만드는 코드
	max_pos = 4104
	max_seq_len = 4096
	attention_window = 512
	save_path = '../LED_KoBART/model'

	# 기존 pretrained 된 kobart tokenizer & model load
	tokenizer = PreTrainedTokenizerFast.from_pretrained('ainize/kobart-news', model_max_length=max_pos)
	kobart_longformer = BartForConditionalGeneration.from_pretrained('ainize/kobart-news')
	config  = LongformerBartConfig.from_pretrained('ainize/kobart-news')

	kobart_longformer.config = config

	config.attention_probs_dropout_prob = config.attention_dropout
	config.architectures = ['LongformerEncoderDecoderForConditionalGeneration', ]

	# Tokenizer의 max_positional_embedding_size 확장
	# extend position embeddings
	tokenizer.model_max_length = max_pos
	tokenizer.init_kwargs['model_max_length'] = max_pos
	current_max_pos, embed_size = kobart_longformer.model.encoder.embed_positions.weight.shape
	assert current_max_pos == config.max_position_embeddings + 2

	config.max_encoder_position_embeddings = max_pos
	config.max_decoder_position_embeddings = config.max_position_embeddings
	del config.max_position_embeddings
	max_pos += 2  # NOTE: BART has positions 0,1 reserved, so embedding size is max position + 2
	assert max_pos >= current_max_pos

	new_encoder_pos_embed = kobart_longformer.model.encoder.embed_positions.weight.new_empty(max_pos, embed_size)

	# Positional Embedding 확장
	k = 2
	step = 1028 - 2
	while k < max_pos - 1:
		new_encoder_pos_embed[k:(k + step)] = kobart_longformer.model.encoder.embed_positions.weight[2:]
		k += step
	kobart_longformer.model.encoder.embed_positions.weight.data = new_encoder_pos_embed

	config.attention_window = [attention_window] * config.num_hidden_layers
	config.attention_dilation = [1] * config.num_hidden_layers

	# Kobart Self attention > Longformer Self Attention
	for i, layer in enumerate(kobart_longformer.model.encoder.layers):
		longformer_self_attn_for_bart = LongformerSelfAttentionForBart(kobart_longformer.config, layer_id=i)

		longformer_self_attn_for_bart.longformer_self_attn.query = layer.self_attn.q_proj
		longformer_self_attn_for_bart.longformer_self_attn.key = layer.self_attn.k_proj
		longformer_self_attn_for_bart.longformer_self_attn.value = layer.self_attn.v_proj

		longformer_self_attn_for_bart.longformer_self_attn.query_global = copy.deepcopy(layer.self_attn.q_proj)
		longformer_self_attn_for_bart.longformer_self_attn.key_global = copy.deepcopy(layer.self_attn.k_proj)
		longformer_self_attn_for_bart.longformer_self_attn.value_global = copy.deepcopy(layer.self_attn.v_proj)

		longformer_self_attn_for_bart.output = layer.self_attn.out_proj

		layer.self_attn = longformer_self_attn_for_bart

	kobart_longformer.save_pretrained(save_path)
	tokenizer.save_pretrained(save_path, None)