Create train.py

train.py

@@ -0,0 +1,78 @@
+from datasets import DatasetDict, load_dataset
+from evaluate import load as load_metric
+from transformers import *
+
+
+def train(batch_size: int, model_name: str="t5-small", max_steps: int=10_000) -> None:
+    total_batch_size_per_step = 512
+    grad_acc_steps = total_batch_size_per_step // batch_size
+    assert grad_acc_steps * batch_size == total_batch_size_per_step
+
+    model_name_for_path = model_name.split("/")[-1]
+    output_dir = f"wmt19-ende-{model_name_for_path}"
+    args = Seq2SeqTrainingArguments(
+        output_dir=output_dir,
+        learning_rate=1e-4,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size * 2,
+        gradient_accumulation_steps=grad_acc_steps,
+        max_steps=max_steps,
+        weight_decay=1e-2,
+        optim="adamw_torch_fused",
+        lr_scheduler_type="constant",
+        evaluation_strategy="steps",
+        eval_steps=100,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=1,
+        save_safetensors=True,
+        metric_for_best_model="bleu",
+        push_to_hub=True,
+        bf16=True,
+        bf16_full_eval=True,
+        seed=42,
+        predict_with_generate=True,
+        log_level="error",
+        logging_steps=1,
+        logging_dir=output_dir,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    bleu = load_metric("bleu")
+    def compute_metrics(eval_preds: EvalPrediction):
+        logits, label_ids = eval_preds
+        label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+        references = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+        predictions = tokenizer.batch_decode(logits, skip_special_tokens=True)
+
+        bleu_outputs = bleu.compute(predictions=predictions, references=references)
+        return {
+            "bleu": 100 * bleu_outputs["bleu"],
+            "brevity_penalty": bleu_outputs["brevity_penalty"],
+        }
+
+    def map_fn(inputs):
+        map_fn = lambda s: tokenizer([d[s] for d in inputs["translation"]], return_attention_mask=False, max_length=64, truncation=True).input_ids
+        return {
+            "input_ids": map_fn("de"),
+            "labels": map_fn("en"),
+        }
+
+    get_dataset_split = lambda s: load_dataset("wmt19", "de-en", split=s, streaming=True).map(map_fn, batched=True)
+    apply_length_filter = lambda d: d.filter(lambda e: len(e["input_ids"]) >= 8 and len(e["labels"]) >= 8)
+
+    trainer = Seq2SeqTrainer(
+        model=AutoModelForSeq2SeqLM.from_pretrained(model_name),
+        args=args,
+        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
+        train_dataset=apply_length_filter(get_dataset_split("train")),
+        eval_dataset=get_dataset_split("validation"),
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+    trainer.remove_callback(PrinterCallback)
+
+    trainer.train()
+    trainer.push_to_hub()