{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Login to HuggingFace (just login once)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from huggingface_hub import interpreter_login\n", "interpreter_login()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Collect Menu Image Datasets\n", "- Use `metadata.jsonl` to label the images's ground truth. You can visit [here](https://github.com/ryanlinjui/menu-text-detection/tree/main/examples) to see the examples.\n", "- 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).\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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from datasets import load_dataset\n", "\n", "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" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Setup for Fine-tuning" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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\" # 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 = \"\" # 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", "\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: 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", "# 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", " image_column=\"image\",\n", " 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=False,\n", " seed=42,\n", " shuffle=False\n", ")\n", "\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", "model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([TASK_PROMPT_NAME])[0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Start Fine-tuning" ] }, { "cell_type": "code", "execution_count": null, "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 = 1 # set your training batch size\n", "LEARNING_RATE = 3e-5 # set your learning rate\n", "WEIGHT_DECAY = 0.1 # set your weight decay\n", "\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", "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=LEARNING_RATE,\n", " weight_decay=WEIGHT_DECAY,\n", " per_device_eval_batch_size=1,\n", " output_dir=\"./.checkpoints\",\n", " seed=42,\n", " warmup_steps=30,\n", " eval_strategy=\"steps\",\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", " save_steps=200,\n", " push_to_hub=True if HUGGINGFACE_MODEL_ID else False,\n", " hub_model_id=HUGGINGFACE_MODEL_ID,\n", " hub_strategy=\"every_save\",\n", " report_to=\"tensorboard\"\n", ")\n", "trainer = Seq2SeqTrainer(\n", " model=model,\n", " args=training_args,\n", " train_dataset=train_datasets,\n", " eval_dataset=validation_datasets,\n", " processing_class=processor,\n", " compute_metrics=compute_metrics\n", ")\n", "\n", "trainer.train()\n", "trainer.push_to_hub()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from PIL import Image\n", "from transformers import pipeline\n", "from transformers import DonutProcessor\n", "\n", "MODEL_REPO_ID = \"ryanlinjui/donut-base-finetuned-menu\"\n", "TASK_PROMPT_NAME = \"\"\n", "MAX_LENGTH = 1024\n", "IMAGE_SIZE = [1280, 960]\n", "\n", "processor = DonutProcessor.from_pretrained(MODEL_REPO_ID)\n", "pipe = pipeline(\"image-text-to-text\", model=MODEL_REPO_ID, processor=processor)\n", "image = Image.open(\"./examples/menu-hd.jpg\")\n", "\n", "outputs = pipe(text=TASK_PROMPT_NAME, images=image)[0][\"generated_text\"]\n", "\n", "print(outputs)\n", "print(processor.token2json(outputs))" ] }, { "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 " ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.12" } }, "nbformat": 4, "nbformat_minor": 2 }