chore: add compute metrix

Author: ryanlinjui

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "gradio>=5.29.0",
     "huggingface-hub>=0.31.1",
     "matplotlib>=3.10.1",
+    "nltk>=3.9.1",
     "notebook>=7.4.2",
     "openai>=1.77.0",
     "pillow>=11.2.1",

train.ipynb

@@ -70,8 +70,8 @@
     "\n",
     "DATASETS_REPO_ID = \"ryanlinjui/menu-zh-TW\"              # set your dataset repo id for training\n",
     "PRETRAINED_MODEL_REPO_ID = \"naver-clova-ix/donut-base\"  # set your pretrained model repo id for fine-tuning\n",
-    "TASK_PROMPT_NAME = \"<s_menu>\"                           # set your task prompt name for training\n",
-    "MAX_LENGTH = 768                                        # set your max length for maximum output length\n",
+    "TASK_PROMPT_NAME = \"<s_menu-text-detection>\"            # set your task prompt name for training\n",
+    "MAX_LENGTH = 1024                                       # set your max length for maximum output length, max to 1536 for donut-base\n",
     "IMAGE_SIZE = [1280, 960]                                # set your image size for training\n",
     "\n",
     "raw_datasets = load_dataset(DATASETS_REPO_ID)\n",
@@ -97,11 +97,13 @@
     "    annotation_column=\"menu\",\n",
     "    task_start_token=TASK_PROMPT_NAME,\n",
     "    prompt_end_token=TASK_PROMPT_NAME,\n",
+    "    max_length=MAX_LENGTH,\n",
     "    train_split=0.8,\n",
     "    validation_split=0.1,\n",
     "    test_split=0.1,\n",
-    "    sort_json_key=True,\n",
-    "    seed=42\n",
+    "    sort_json_key=False,\n",
+    "    seed=42,\n",
+    "    shuffle=False\n",
     ")\n",
     "\n",
     "# Model: load the pretrained model and set the config.\n",
@@ -124,12 +126,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from functools import reduce\n",
+    "\n",
     "import torch\n",
+    "import numpy as np\n",
+    "from nltk.metrics import edit_distance\n",
+    "from transformers.trainer_utils import EvalPrediction\n",
     "from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments\n",
     "\n",
     "HUGGINGFACE_MODEL_ID = \"ryanlinjui/donut-base-finetuned-menu\" # set your huggingface model repo id for saving / pushing to the hub\n",
     "EPOCHS = 100            # set your training epochs\n",
-    "TRAIN_BATCH_SIZE = 4    # set your training batch size\n",
+    "TRAIN_BATCH_SIZE = 3    # set your training batch size\n",
     "\n",
     "device = (\n",
     "    \"cuda\"\n",
@@ -139,16 +146,46 @@
     "print(f\"Using {device} device\")\n",
     "model.to(device)\n",
     "\n",
+    "train_datasets = datasets[\"train\"]\n",
+    "validation_datasets = datasets[\"validation\"]\n",
+    "filtered_tokens = [\n",
+    "    processor.tokenizer.bos_token,\n",
+    "    processor.tokenizer.eos_token,\n",
+    "    processor.tokenizer.pad_token,\n",
+    "    processor.tokenizer.unk_token,\n",
+    "]\n",
+    "def compute_metrics(eval_pred: EvalPrediction) -> dict:\n",
+    "    decoded_preds = processor.tokenizer.batch_decode(eval_pred.predictions, skip_special_tokens=False)\n",
+    "\n",
+    "    normed_eds = []\n",
+    "    for idx, pred in enumerate(decoded_preds):\n",
+    "        prediction_sequence = reduce(lambda s, t: s.replace(t, \"\"), filtered_tokens, pred)\n",
+    "        target_sequence = reduce(lambda s, t: s.replace(t, \"\"), filtered_tokens, validation_datasets[idx][\"target_sequence\"])\n",
+    "        ed = edit_distance(prediction_sequence, target_sequence) / max(len(prediction_sequence), len(target_sequence))\n",
+    "        normed_eds.append(ed)\n",
+    "\n",
+    "        print(f\"[Sample {idx}]\")\n",
+    "        print(f\"  Prediction: {prediction_sequence}\")\n",
+    "        print(f\"  Target: {target_sequence}\")\n",
+    "        print(f\"  Normalized Edit Distance: {ed:.4f}\")\n",
+    "        print(\"-\" * 40)\n",
+    "\n",
+    "    return {\"normed_edit_distance\": float(np.mean(normed_eds))}\n",
+    "\n",
     "training_args = Seq2SeqTrainingArguments(\n",
     "    num_train_epochs=EPOCHS,\n",
     "    per_device_train_batch_size=TRAIN_BATCH_SIZE,\n",
     "    learning_rate=3e-5,\n",
     "    per_device_eval_batch_size=1,\n",
     "    output_dir=\"./.checkpoints\",\n",
     "    seed=2022,\n",
-    "    warmup_steps=30,\n",
+    "    warmup_steps=300,\n",
     "    eval_strategy=\"steps\",\n",
-    "    eval_steps=100,\n",
+    "    eval_steps=200,\n",
+    "    fp16=(device == \"cuda\"),\n",
+    "    predict_with_generate=True,\n",
+    "    generation_max_length=MAX_LENGTH,\n",
+    "    generation_num_beams=1,\n",
     "    logging_strategy=\"steps\",\n",
     "    logging_steps=50,\n",
     "    save_strategy=\"steps\",\n",
@@ -157,17 +194,19 @@
     "    hub_model_id=HUGGINGFACE_MODEL_ID,\n",
     "    hub_strategy=\"every_save\",\n",
     "    report_to=\"tensorboard\",\n",
-    "    logging_dir=\"./.checkpoints/logs\",\n",
+    "    logging_dir=\"./.checkpoints/logs\"\n",
     ")\n",
     "trainer = Seq2SeqTrainer(\n",
     "    model=model,\n",
     "    args=training_args,\n",
-    "    train_dataset=datasets[\"train\"],\n",
-    "    eval_dataset=datasets[\"test\"],\n",
-    "    tokenizer=processor\n",
+    "    train_dataset=train_datasets,\n",
+    "    eval_dataset=validation_datasets,\n",
+    "    tokenizer=processor,\n",
+    "    compute_metrics=compute_metrics\n",
     ")\n",
     "\n",
-    "trainer.train()"
+    "trainer.train()\n",
+    "trainer.push_to_hub()"
    ]
   },
   {