chore: add dataset, setup tutorial info

Author: ryanlinjui

train.ipynb

@@ -26,7 +26,16 @@
     "- After finishing, push to HuggingFace Datasets.\n",
     "- For labeling:\n",
     "    - [Google AI Studio](https://aistudio.google.com) or [OpenAI ChatGPT](https://chatgpt.com).\n",
-    "    - Use function calling by API. Start the gradio app locally or visit [here](https://huggingface.co/spaces/ryanlinjui/menu-text-detection)."
+    "    - Use function calling by API. Start the gradio app locally or visit [here](https://huggingface.co/spaces/ryanlinjui/menu-text-detection).\n",
+    "\n",
+    "### Menu Type\n",
+    "- **h**: horizontal menu\n",
+    "- **v**: vertical menu\n",
+    "- **d**: document-style menu\n",
+    "- **s**: in-scene menu (non-document style)\n",
+    "- **i**: irregular menu (menu with irregular text layout)\n",
+    "\n",
+    "> Please see the [examples](https://github.com/ryanlinjui/menu-text-detection/tree/main/examples) for more details."
    ]
   },
   {
@@ -37,8 +46,8 @@
    "source": [
     "from datasets import load_dataset\n",
     "\n",
-    "dataset = load_dataset(path=\"datasets/menu-zh-TW\")\n",
-    "dataset.push_to_hub(repo_id=\"ryanlinjui/menu-zh-TW\")"
+    "dataset = load_dataset(path=\"datasets/menu-zh-TW\")      # load dataset from the local directory including the metadata.jsonl, images files.\n",
+    "dataset.push_to_hub(repo_id=\"ryanlinjui/menu-zh-TW\")    # push to the huggingface dataset hub"
    ]
   },
   {
@@ -56,27 +65,31 @@
    "source": [
     "from datasets import load_dataset\n",
     "from transformers import DonutProcessor, VisionEncoderDecoderModel, VisionEncoderDecoderConfig\n",
+    "\n",
     "from menu.donut import DonutDatasets\n",
     "\n",
-    "DATASETS_REPO_ID = \"ryanlinjui/menu-zh-TW\"\n",
-    "PRETRAINED_MODEL_REPO_ID = \"naver-clova-ix/donut-base\"\n",
-    "TASK_PROMPT_NAME = \"<s_menu>\"\n",
-    "MAX_LENGTH = 768\n",
-    "IMAGE_SIZE = [1280, 960]\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",
+    "IMAGE_SIZE = [1280, 960]                                # set your image size for training\n",
     "\n",
     "raw_datasets = load_dataset(DATASETS_REPO_ID)\n",
     "\n",
-    "# Config: 預訓練模型載入 Encoder–Decoder 的設定\n",
+    "# Config: set the model config\n",
     "config = VisionEncoderDecoderConfig.from_pretrained(PRETRAINED_MODEL_REPO_ID)\n",
     "config.encoder.image_size = IMAGE_SIZE\n",
     "config.decoder.max_length = MAX_LENGTH\n",
     "\n",
-    "# Processor: 影像前處理與文字後處理\n",
+    "# Processor: use the processor to process the dataset. \n",
+    "# Convert the image to the tensor and the text to the token ids.\n",
     "processor = DonutProcessor.from_pretrained(PRETRAINED_MODEL_REPO_ID)\n",
     "processor.feature_extractor.size = IMAGE_SIZE[::-1]\n",
     "processor.feature_extractor.do_align_long_axis = False\n",
     "\n",
-    "# Donut Datasets: \n",
+    "# DonutDatasets: use the DonutDatasets to process the dataset.\n",
+    "# For model inpit, the image must be converted to the tensor and the json text must be converted to the token with the task prompt string.\n",
+    "# This example sets the column name by \"image\" and \"menu\". So that image file is included in the \"image\" column and the json text is included in the \"menu\" column.\n",
     "datasets = DonutDatasets(\n",
     "    datasets=raw_datasets,\n",
     "    processor=processor,\n",
@@ -91,7 +104,7 @@
     "    seed=42\n",
     ")\n",
     "\n",
-    "# Model: 載入預訓練模型\n",
+    "# Model: load the pretrained model and set the config.\n",
     "model = VisionEncoderDecoderModel.from_pretrained(PRETRAINED_MODEL_REPO_ID, config=config)\n",
     "model.decoder.resize_token_embeddings(len(processor.tokenizer))\n",
     "model.config.pad_token_id = processor.tokenizer.pad_token_id\n",
@@ -114,15 +127,17 @@
     "import torch\n",
     "from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments\n",
     "\n",
-    "HUGGINGFACE_MODEL_ID = \"ryanlinjui/donut-base-finetuned-menu\"\n",
-    "EPOCHS = 100\n",
-    "TRAIN_BATCH_SIZE = 4\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",
     "\n",
-    "if torch.cuda.is_available():\n",
-    "    print(\"Using GPU\")\n",
-    "    model.to(\"cuda\")\n",
-    "else:\n",
-    "    print(\"Using default device\")\n",
+    "device = (\n",
+    "    \"cuda\"\n",
+    "    if torch.cuda.is_available()\n",
+    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
+    ")\n",
+    "print(f\"Using {device} device\")\n",
+    "model.to(device)\n",
     "\n",
     "training_args = Seq2SeqTrainingArguments(\n",
     "    num_train_epochs=EPOCHS,\n",
@@ -189,33 +204,22 @@
     "outputs = ocr_pipeline(image)\n",
     "\n",
     "# 5. 印出辨識文字\n",
-    "print(outputs[0][\"generated_text\"])\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Plot the results"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "print(outputs[0][\"generated_text\"])\n",
+    "\n",
+    "'''\n",
     "# test model\n",
     "import re\n",
     "\n",
     "from transformers import VisionEncoderDecoderModel\n",
+    "from transformers import DonutProcessor\n",
+    "import torch\n",
     "from PIL import Image\n",
     "\n",
     "image = Image.open(\"./examples/menu-hd.jpg\").convert(\"RGB\")\n",
     "\n",
     "processor = DonutProcessor.from_pretrained(\"ryanlinjui/donut-base-finetuned-menu\")\n",
     "model = VisionEncoderDecoderModel.from_pretrained(\"ryanlinjui/donut-base-finetuned-menu\")\n",
-    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"mps\"\n",
     "\n",
     "model.eval()\n",
     "model.to(device)\n",
@@ -241,9 +245,28 @@
     "\n",
     "seq = processor.batch_decode(outputs.sequences)[0]\n",
     "seq = seq.replace(processor.tokenizer.eos_token, \"\").replace(processor.tokenizer.pad_token, \"\")\n",
-    "seq = re.sub(r\"<.*?>\", \"\", seq, count=1).strip()  # remove first task start token\n",
+    "# seq = re.sub(r\"<.*?>\", \"\", seq, count=1).strip()  # remove first task start token\n",
     "seq = processor.token2json(seq)\n",
-    "print(seq)"
+    "print(seq)\n",
+    "'''\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Plot the results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training Loss\n",
+    "# Validation Normal ED per each epoch 1~0, 1 -> 0.22\n",
+    "# Test Accuracy TED Accuracy, F1 Score Accuracy 0.687058, 0.51119 "
    ]
   }
  ],
@@ -263,7 +286,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.10"
+   "version": "3.11.12"
   }
  },
  "nbformat": 4,