wip

src/ainu_lm_trainer/services/mt5/mt5_trainer.py

@@ -53,18 +53,23 @@ def train(self) -> None:
             lambda example: len(example["text"]) > 0 and len(example["translation"]) > 0
         )
 
+        # Set task prefix
+        dataset = dataset.map(
+            lambda example: {
+                "text": self.__get_task_prefix(example) + example["translation"],
+                "target": example["text"],
+            },
+            remove_columns=dataset.column_names,
+        )
+
         # https://huggingface.co/docs/transformers/en/tasks/summarization#preprocess
         def preprocess(examples: dict) -> dict:
-            inputs = tokenizer(
-                [
-                    self.__get_task_prefix(example) + example["translation"]
-                    for example in examples
-                ],
-                text_target=examples["text"],
+            return tokenizer(
+                examples["text"],
+                text_target=examples["target"],
                 max_length=self.__context_length,
                 truncation=True,
             )
-            return inputs
 
         dataset = dataset.map(
             preprocess,
@@ -114,10 +119,10 @@ def preprocess(examples: dict) -> dict:
             tokenizer.push_to_hub(self.__model_name)
 
     def __get_task_prefix(self, example: dict) -> str:
-        if example["dialect"] is not None:
-            return f"translate: Japanese to Ainu ({example['dialect']}, {example['pronoun']}): "
+        if example["dialect"] is None:
+            return f"translate: Japanese to Ainu: (沙流, {example['pronoun']}): "
         else:
-            return f"translate: Japanese to Ainu (沙流, {example['pronoun']}): "
+            return f"translate: Japanese to Ainu: ({example['dialect']}, {example['pronoun']}): "
 
     def __compute_metrics(
         self, tokenizer: MT5Tokenizer, eval_preds: EvalPrediction