{
 "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 = \"<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: 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",
    "    train_split=0.8,\n",
    "    validation_split=0.1,\n",
    "    test_split=0.1,\n",
    "    sort_json_key=True,\n",
    "    seed=42\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": [
    "import torch\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",
    "\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",
    "    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",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=100,\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",
    "    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",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import (\n",
    "    VisionEncoderDecoderModel,\n",
    "    DonutProcessor,\n",
    "    pipeline\n",
    ")\n",
    "from PIL import Image\n",
    "\n",
    "model_id = \"ryanlinjui/donut-base-finetuned-menu\"\n",
    "\n",
    "# 1. 下載並載入 model + processor\n",
    "processor = DonutProcessor.from_pretrained(model_id)\n",
    "model     = VisionEncoderDecoderModel.from_pretrained(model_id)\n",
    "\n",
    "# 2. 建立一個 image-to-text pipeline\n",
    "ocr_pipeline = pipeline(\n",
    "    \"image-to-text\",             # 使用 image-to-text 任務\n",
    "    model=model,                 # 傳入已載入的 model\n",
    "    tokenizer=processor.tokenizer,\n",
    "    feature_extractor=processor.feature_extractor,\n",
    ")\n",
    "\n",
    "# 3. 載入一張測試圖片\n",
    "image = Image.open(\"./examples/menu-hd.jpg\")\n",
    "\n",
    "# 4. 呼叫 pipeline，取得結果\n",
    "outputs = ocr_pipeline(image)\n",
    "\n",
    "# 5. 印出辨識文字\n",
    "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 \"mps\"\n",
    "\n",
    "model.eval()\n",
    "model.to(device)\n",
    "\n",
    "pixel_values = processor(image, return_tensors=\"pt\").pixel_values\n",
    "pixel_values = pixel_values.to(device)\n",
    "\n",
    "task_prompt = \"<s_menu>\"\n",
    "decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors=\"pt\").input_ids\n",
    "decoder_input_ids = decoder_input_ids.to(device)\n",
    "outputs = model.generate(\n",
    "    pixel_values,\n",
    "    decoder_input_ids=decoder_input_ids,\n",
    "    max_length=model.decoder.config.max_position_embeddings,\n",
    "    early_stopping=True,\n",
    "    pad_token_id=processor.tokenizer.pad_token_id,\n",
    "    eos_token_id=processor.tokenizer.eos_token_id,\n",
    "    use_cache=True,\n",
    "    num_beams=1,\n",
    "    bad_words_ids=[[processor.tokenizer.unk_token_id]],\n",
    "    return_dict_in_generate=True,\n",
    ")\n",
    "\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 = processor.token2json(seq)\n",
    "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 "
   ]
  }
 ],
 "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
}