Experiment T5

src/ainu_lm_pipeline/components/get_mt_training_job_spec.py

@@ -15,8 +15,8 @@ def get_mt_training_job_spec(
                 "image_uri": train_image_uri,
                 "args": [
                     "train",
-                    "mt",
-                    "--base-model=google/mt5-small",
+                    "t5",
+                    "--base-tokenizer=aynumosir/sentencepiece-ainu",
                     f"--dataset-name={dataset_name}",
                     f"--dataset-revision={dataset_revision}",
                     "--num-train-epochs=20",
@@ -26,7 +26,6 @@ def get_mt_training_job_spec(
                     "--learning-rate=5e-4",
                     "--warmup-ratio=0.06",
                     "--weight-decay=0.01",
-                    "--experiment-task-prefix=all",
                     f"--hub-model-id={hub_model_id}",
                     f"--push-to-hub={'yes' if push_to_hub else 'no'}",
                 ],

src/ainu_lm_pipeline/pipelines/ainu_mt_pipeline.py

@@ -85,7 +85,7 @@ def ainu_mt_pipeline(
         train_image_uri=train_image_uri,
         dataset_name=hf_dataset_repo,
         dataset_revision=get_dataset_revision_op.output,
-        hub_model_id="aynumosir/mt5-base-ainu",
+        hub_model_id="aynumosir/t5-base-ainu",
         push_to_hub=True,
     ).set_display_name("MTジョブの仕様を取得")
 

src/ainu_lm_trainer/app/commands/train.py

@@ -68,6 +68,9 @@ def add_parser(parser: ArgumentParser) -> None:
     gpt2_parser = subparsers.add_parser("gpt2", parents=common)
     gpt2_parser.add_argument("--base-tokenizer", type=str)
 
+    t5_parser = subparsers.add_parser("t5", parents=common)
+    t5_parser.add_argument("--base-tokenizer", type=str)
+
     subparsers.add_parser("byte-level-bpe", parents=common)
     subparsers.add_parser("sentencepiece", parents=common)
 
@@ -151,6 +154,14 @@ def main(args: Namespace) -> None:
             config_workspace=config_workspace,
         )
 
+    if args.task == "t5":
+        pretraining.t5.train(
+            tokenizer_name=args.base_tokenizer,
+            config_dataset=config_dataset,
+            config_training=config_training,
+            config_workspace=config_workspace,
+        )
+
     if args.task == "byte-level-bpe":
         bpe_trainer = pretraining.ByteLevelBpeTokenizerTrainer(
             config_dataset=config_dataset,

src/ainu_lm_trainer/trainers/fine_tuning/mt/mt.py

@@ -19,8 +19,7 @@
     TrainingConfig,
     WorkspaceConfig,
 )
-from ....utils import HyperparameterTuningCallback
-from . import task_prefix
+from ....utils import HyperparameterTuningCallback, task_prefix
 
 sacrebleu = evaluate.load("sacrebleu")
 

src/ainu_lm_trainer/trainers/pretraining/__init__.py

@@ -1,3 +1,3 @@
-from . import gpt2, roberta
+from . import gpt2, roberta, t5
 from .byte_level_bpe import ByteLevelBpeTokenizerTrainer
 from .sentencepiece import SentencepieceTokenizerTrainer

src/ainu_lm_trainer/trainers/pretraining/sentencepiece/sentencepiece.py

@@ -26,6 +26,11 @@ def __batch_iterator(self, batch_size: int = 1_000) -> Iterator[Iterator[str]]:
             yield self.__dataset[i : i + batch_size]["text"]
             yield self.__dataset[i : i + batch_size]["translation"]
 
+            # 方言名がちゃんとトークナイズされること
+            for dialect in self.__dataset[i : i + batch_size]["dialect"]:
+                if dialect:
+                    yield dialect
+
     def __prepare(self) -> None:
         self.__config_workspace.model_dir.mkdir(parents=True, exist_ok=True)
 
@@ -36,6 +41,7 @@ def train(self) -> SentencePieceUnigramTokenizer:
 
         tokenizer.train_from_iterator(
             iterator=self.__batch_iterator(),
+            unk_token="<unk>",
             special_tokens=[
                 "<unk>",
                 "<pad>",

src/ainu_lm_trainer/trainers/pretraining/t5/__init__.py

@@ -0,0 +1 @@
+from .mt import train

src/ainu_lm_trainer/trainers/pretraining/t5/mt.py

@@ -0,0 +1,129 @@
+import evaluate
+import numpy as np
+import torch
+from datasets import DatasetDict, interleave_datasets
+from transformers import (
+    DataCollatorForSeq2Seq,
+    EarlyStoppingCallback,
+    EvalPrediction,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    T5Config,
+    T5ForConditionalGeneration,
+    T5TokenizerFast,
+)
+
+from ....config import DatasetConfig, TrainingConfig, WorkspaceConfig
+from ....utils import task_prefix
+
+sacrebleu = evaluate.load("sacrebleu")
+
+
+def compute_metrics(tokenizer: T5TokenizerFast, eval_preds: EvalPrediction) -> dict:
+    predictions, labels = eval_preds
+
+    predictions = np.where(predictions != -100, predictions, tokenizer.pad_token_id)
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    bleu = sacrebleu.compute(predictions=decoded_preds, references=decoded_labels)
+    return {"bleu": bleu["score"]}
+
+
+def train(
+    config_dataset: DatasetConfig,
+    config_training: TrainingConfig,
+    config_workspace: WorkspaceConfig,
+    tokenizer_name: str,
+    context_length: int = 128,
+) -> None:
+    tokenizer = T5TokenizerFast.from_pretrained(tokenizer_name)
+
+    config = T5Config.from_pretrained("google-t5/t5-base")
+
+    model = T5ForConditionalGeneration(config)
+    model = model.to("cuda") if torch.cuda.is_available() else model
+
+    dataset_dict = config_dataset.load()
+    dataset_dict = dataset_dict.filter(
+        lambda example: len(example["text"]) > 0 and len(example["translation"]) > 0
+    )
+
+    # --------------------------------------
+    # Prepare Japanese to Ainu
+    # --------------------------------------
+    dataset_dict_ja2ain = dataset_dict.map(
+        lambda example: {
+            "text": task_prefix.ja2ain(example) + example["translation"],
+            "text_target": example["text"],
+        },
+        remove_columns=dataset_dict.column_names["train"],
+    )
+
+    # --------------------------------------
+    # Prepare Ainu to Japanese
+    # --------------------------------------
+    dataset_dict_ain2ja = dataset_dict.map(
+        lambda example: {
+            "text": task_prefix.ain2ja(example) + example["text"],
+            "text_target": example["translation"],
+        },
+        remove_columns=dataset_dict.column_names["train"],
+    )
+
+    dataset_dict = DatasetDict(
+        {
+            "train": interleave_datasets(
+                [dataset_dict_ain2ja["train"], dataset_dict_ja2ain["train"]],
+                stopping_strategy="all_exhausted",
+            ),
+            "test": interleave_datasets(
+                [dataset_dict_ain2ja["test"], dataset_dict_ja2ain["test"]],
+                stopping_strategy="all_exhausted",
+            ),
+        }
+    )
+
+    dataset_dict = dataset_dict.map(
+        lambda examples: tokenizer(
+            examples["text"],
+            text_target=examples["text_target"],
+            max_length=context_length,
+            truncation=True,
+        ),
+        batched=True,
+    )
+
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    training_args = Seq2SeqTrainingArguments(
+        save_strategy="epoch",
+        evaluation_strategy="epoch",
+        output_dir=str(config_workspace.checkpoint_dir),
+        generation_max_length=context_length,
+        predict_with_generate=True,
+        load_best_model_at_end=True,
+        metric_for_best_model="bleu",
+        greater_is_better=True,
+        logging_dir=str(config_workspace.logging_dir),
+        report_to=["tensorboard"],
+    )
+    training_args = config_training.extend(training_args)
+
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=dataset_dict["train"],
+        eval_dataset=dataset_dict["test"],
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
+        compute_metrics=lambda eval_preds: compute_metrics(tokenizer, eval_preds),
+    )
+
+    trainer.train()
+    trainer.save_model(str(config_workspace.model_dir))
+
+    if config_training.push_to_hub:
+        model.push_to_hub(training_args.hub_model_id)
+        tokenizer.push_to_hub(training_args.hub_model_id)

src/ainu_lm_trainer/utils/__init__.py

@@ -1,2 +1,3 @@
+from . import task_prefix
 from .get_path_from_uri import get_path_from_uri, get_path_str_from_uri
 from .hyperparameter_tuning_callback import HyperparameterTuningCallback

src/ainu_lm_trainer/trainers/fine_tuning/mt/task_prefix.py → src/ainu_lm_trainer/utils/task_prefix.py

@@ -1,4 +1,4 @@
-from ....config import TaskPrefixType
+from ..config import TaskPrefixType
 
 
 def __make_ainu_language_identifier(

src/ainu_lm_trainer/trainers/fine_tuning/mt/task_prefix_test.py → src/ainu_lm_trainer/utils/task_prefix_test.py

@@ -1,4 +1,4 @@
-from ....config import TaskPrefixType
+from ..config import TaskPrefixType
 from .task_prefix import ain2ja
 
 sentence = {