lightning-transformers를 활용하여, Causal Language Modeling을 구현하는 방법은 LANGUAGE MODELING에 간략히 소개되어있고, 본 포스트에서는 위의 링크에서 간략히 소개된 내용들이 lightning-transformers 라이브러리 내부에서 어떤 코드에 해당하는 지를 분석하였습니다.
실제로 huggingface datasets 라이브러리를 lightning-transformers에서 사용하고있기때문에, https://huggingface.co/datasets 에 등록되어있는 데이터라면 매우 쉽게 사용할 수 있습니다. 예를 들어 아래의 command로 lightning-transformers 라이브러리를 사용한다고 합시다.
python train.py task=nlp/language_modeling dataset=nlp/language_modeling/wikitext
위의 command에서 task=nlp/langague_modeling은 conf directory 하위에 존재하는 conf/task/nlp/language_modeling.yaml과 conf/dataset/nlp/language_modeling/wikitext.yaml를 train.py에 argument로 전달하는 것을 의미합니다. (전체 lightning-transformers 코드 구성에 관한 설명은 Masked Language Modeling with lightning-transformers 를 참고해주세요!)
├── conf
│ ├── __init__.py
│ ├── backbone
│ │ └── nlp
│ │ ├── default.yaml
│ │ └── seq2seq.yaml
│ ├── config.yaml
│ ├── dataset
│ │ ├── default.yaml
│ │ └── nlp
│ │ ├── default.yaml
│ │ ├── language_modeling
│ │ │ ├── default.yaml
│ │ │ └── wikitext.yaml
...
│ ├── task
│ │ ├── default.yaml
│ │ └── nlp
│ │ ├── default.yaml
│ │ ├── language_modeling.yaml
│ │ ├── multiple_choice.yaml
│ │ ├── question_answering.yaml
│ │ ├── summarization.yaml
│ │ ├── text_classification.yaml
│ │ ├── token_classification.yaml
│ │ └── translation.yaml
...
│ ├── tokenizer
│ │ └── autotokenizer.yaml
...
각각의 yaml 파일을 분석해보겠습니다. 먼저 conf/dataset/nlp/language_modeling/wikitext.yaml 아래와 같으며, lighting-transformers가 사용하고있는 hydra에서 지원하는 문법이 적용되어 있음을 확인할 수 있습니다. 특히 아래의 yaml 파일에서 __target__에 python class 명을 적으면 hydra가 알아서 instantiate를 해주는 기능입니다. (자세한 내용은 Instantiating objects with Hydra 를 참고해주세요!)
conf/dataset/nlp/language_modeling/wikitext.yaml
# @package dataset
defaults:
- nlp/default
_target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingDataModule
cfg:
dataset_name: wikitext
dataset_config_name: wikitext-2-raw-v1
block_size: 512
실제로 위의 yaml 파일에 정의된대로 cfg 하위에 정의된 dataset_name, dataset_config_name, block_size key에 대응되는 value로 LanguageModelingDataModule class를 instantiate하게 됩니다.
lightning_transformers.task.nlp.language_modeling.LanguageModelingDataModule
class LanguageModelingDataModule(HFDataModule):
"""
Defines ``LightningDataModule`` for Language Modeling Datasets.
Args:
*args: ``HFDataModule`` specific arguments.
cfg: Contains data specific parameters when processing/loading the dataset
(Default ``LanguageModelingDataConfig``)
**kwargs: ``HFDataModule`` specific arguments.
"""
cfg: LanguageModelingDataConfig
def __init__(self, *args, cfg: LanguageModelingDataConfig = LanguageModelingDataConfig(), **kwargs) -> None:
super().__init__(*args, cfg=cfg, **kwargs)
def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
column_names = dataset["train" if stage == "fit" else "validation"].column_names
text_column_name = "text" if "text" in column_names else column_names[0]
tokenize_function = partial(self.tokenize_function, tokenizer=self.tokenizer, text_column_name=text_column_name)
dataset = dataset.map(
tokenize_function,
batched=True,
num_proc=self.cfg.preprocessing_num_workers,
remove_columns=column_names,
load_from_cache_file=self.cfg.load_from_cache_file,
)
convert_to_features = partial(self.convert_to_features, block_size=self.effective_block_size)
dataset = dataset.map(
convert_to_features,
batched=True,
num_proc=self.cfg.preprocessing_num_workers,
load_from_cache_file=self.cfg.load_from_cache_file,
)
return dataset
...
LanguageModelingDataModule이 상속하는 class인 HFDataModule의 load_dataset method에서 cfg 하위에 정의된 value들을 사용합니다.
lightning_transformers.core.nlp.data.HFDataModule
class HFDataModule(TokenizerDataModule):
"""
Base ``LightningDataModule`` for HuggingFace Datasets. Provides helper functions and boilerplate logic
to load/process datasets.
Args:
tokenizer: ``PreTrainedTokenizerBase`` for tokenizing data.
cfg: Contains data specific parameters when processing/loading the dataset (Default ``HFTransformerDataConfig``)
"""
cfg: HFTransformerDataConfig
tokenizer: PreTrainedTokenizerBase
...
def load_dataset(self) -> Dataset:
# Allow custom data files when loading the dataset
data_files = {}
if self.cfg.train_file is not None:
data_files["train"] = self.cfg.train_file
if self.cfg.validation_file is not None:
data_files["validation"] = self.cfg.validation_file
if self.cfg.test_file is not None:
data_files["test_file"] = self.cfg.test_file
data_files = data_files if data_files else None
if self.cfg.dataset_name is not None:
# Download and load the Huggingface dataset.
return load_dataset(
path=self.cfg.dataset_name,
name=self.cfg.dataset_config_name,
cache_dir=self.cfg.cache_dir,
data_files=data_files
)
...
또 다른 yaml 파일인 conf/task/nlp/language_modeling.yaml 아래와 같으며 마찬가지로 __target__에 해당하는 lightning_transformers.task.nlp.language_modeling.LanguageModelingTransformer를 downstream_model_type key의 value인 transformers.AutoModelForCausalLM 전달하여 instantiate 하게됩니다.
conf/task/nlp/language_modeling.yaml
# @package task
defaults:
- nlp/default
- override /dataset@_group_: nlp/language_modeling/default
_target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingTransformer
downstream_model_type: transformers.AutoModelForCausalLM
lightning_transformers.task.nlp.language_modeling.LanguageModelingTransformer
class LanguageModelingTransformer(HFTransformer):
"""
Defines ``LightningModule`` for the Language Modeling Task.
Args:
*args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
downstream_model_type: Downstream HuggingFace AutoModel to load. (default ``transformers.AutoModelForCausalLM``)
**kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
"""
def __init__(self, *args, downstream_model_type: str = 'transformers.AutoModelForCausalLM', **kwargs) -> None:
super().__init__(downstream_model_type, *args, **kwargs)
def on_fit_start(self):
tokenizer_length = len(self.tokenizer)
self.model.resize_token_embeddings(tokenizer_length)
def _step(self, batch, batch_idx):
outputs = self.model(**batch)
loss = outputs[0]
return loss
def training_step(self, batch, batch_idx):
loss = self._step(batch, batch_idx)
self.log("train_loss", loss)
return loss
def validation_step(self, batch, batch_idx, dataloader_idx=0):
loss = self._step(batch, batch_idx)
self.log("val_loss", loss, sync_dist=True)
def test_step(self, batch, batch_idx, dataloader_idx=0):
loss = self._step(batch, batch_idx)
self.log("test_loss", loss, sync_dist=True)
@property
def hf_pipeline_task(self) -> str:
return "text-generation"
예시로든 command에 아래와 같이 backbone.pretrained_model_name_or_path=gpt2를 argument로 전달하면 conf/backbone/nlp/default.yaml에서 pretrained_model_name_or_path key의 value를 gpt2로 설정하여, train.py에 argument로 전달하는 것과 같습니다.
python train.py task=nlp/language_modeling dataset=nlp/language_modeling/wikitext backbone.pretrained_model_name_or_path=gpt2