{"metadata":{"kernelspec":{"display_name":"Python 3","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.7.12"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":43873,"databundleVersionId":5585780}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# End-to-End Image-to-Text Model: Donut (Document Understanding Transformer) 🍩\n\nDonut is a combination of a swin encoder and a bart decoder. This means that it can generate text based on an image without needing any OCR. In this notebook, I show how you can use it to train a model for this competition. This is only the starting point, so I'm sure there is much room for improvement. Still, it is so cool that this one model can do everything necessary for this competition!\n\n![donut diagram](https://raw.githubusercontent.com/clovaai/donut/master/misc/overview.png)","metadata":{}},{"cell_type":"code","source":"'''\n让电脑能看懂图表，并读出里面的数据，\n比如看到一张折线图，能说出这是折线图，x轴：1月，2月，y轴：10，20\n\n\n为什么用Donut：\n传统方法需要好几步：\n（1）先识别图表类型（分类）\n（2）再用OCR识别文字\n（3）再用目标检测找数据点\n而Donut是端到端的：【图表图片】-->【Donut模型】-->【直接输出文本】\n\n\nDount模型是什么：\nDount=Document understanding Transformer\n是一个不需要OCR的文档理解模型，输入图片，直接输出文字\n它有Swin Transformer编码器和BART解码器组成\n\nSwin Transformer编码器：\n作用：把图片转成特征向量\n怎么做的：把图片切成一个个小块(patch),提取视觉特征\n\nBART解码器：\n作用：根据视觉特征，一个字一个字地生成文本\n怎么做的：像ChatGPT一样，自回归生成（一个词一个词往外蹦）\n\n\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# You will have to restart the notebook after running this\n#!pip install -U datasets transformers pyarrow polyleven -q\n!pip install -U datasets\n%reset","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:27.869264Z","iopub.status.busy":"2026-03-10T08:35:27.868533Z","iopub.status.idle":"2026-03-10T08:35:39.264145Z","shell.execute_reply":"2026-03-10T08:35:39.263284Z","shell.execute_reply.started":"2026-03-10T08:35:27.869230Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -U pyarrow>=8.0.0","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:39.266157Z","iopub.status.busy":"2026-03-10T08:35:39.265763Z","iopub.status.idle":"2026-03-10T08:35:48.032291Z","shell.execute_reply":"2026-03-10T08:35:48.031222Z","shell.execute_reply.started":"2026-03-10T08:35:39.266130Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//导入工具包\n#用了huggingFace的Transformers库加载预训练模型,用datasets库处理数据\n\nimport os\nimport re\nimport json\nfrom collections import Counter\nfrom itertools import chain\nfrom pathlib import Path\nfrom typing import List, Dict, Union, Tuple, Any\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom PIL import Image\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    DonutProcessor,\n    VisionEncoderDecoderConfig,\n    VisionEncoderDecoderModel,\n    get_scheduler\n)\nfrom datasets import Dataset\nfrom datasets import Image as ds_img\n#from polyleven import levenshtein # a faster version of levenshtein\ndef levenshtein(s1, s2):\n    if len(s1) < len(s2):\n        return levenshtein(s2, s1)\n    if len(s2) == 0:\n        return len(s1)\n    previous_row = range(len(s2) + 1)\n    for i, c1 in enumerate(s1):\n        current_row = [i + 1]\n        for j, c2 in enumerate(s2):\n            insertions = previous_row[j + 1] + 1\n            deletions = current_row[j] + 1\n            substitutions = previous_row[j] + (c1 != c2)\n            current_row.append(min(insertions, deletions, substitutions))\n        previous_row = current_row\n    return previous_row[-1]","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:48.034094Z","iopub.status.busy":"2026-03-10T08:35:48.033828Z","iopub.status.idle":"2026-03-10T08:35:52.731152Z","shell.execute_reply":"2026-03-10T08:35:52.730036Z","shell.execute_reply.started":"2026-03-10T08:35:48.034068Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#数据路径配置\ndata_dir = Path(\"/kaggle/input/benetech-making-graphs-accessible/train\")\n\nimages_path = data_dir / \"images\"\ntrain_json_files = list((data_dir / \"annotations\").glob(\"*.json\"))\n\n\n#//配置类\nclass CFG:\n\n    # General\n    debug = False\n    num_proc = 2\n    num_workers = 2\n    gpus = 2\n\n    # Data  数据参数\n    max_length = 1024  #//模型最多生成1024个字符\n    image_height = 560  #/图片缩放大小,保持比例\n    image_width = 560\n\n    # Training\n    epochs = 2 #//训练2轮,每轮约3小时\n    val_check_interval = 1.0  # how many times we want to validate during an epoch\n    check_val_every_n_epoch = 1\n    gradient_clip_val = 1.0\n    lr = 3e-5  #//学习率0.00003,微调时的常用值\n    lr_scheduler_type = \"cosine\"\n    num_warmup_steps = 100\n    seed = 42\n    warmup_steps = 300  \n    output_path = \"output\"\n    log_steps = 200\n    batch_size = 2  #//一次处理2张图\n    use_wandb = False","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:52.734698Z","iopub.status.busy":"2026-03-10T08:35:52.733433Z","iopub.status.idle":"2026-03-10T08:35:53.214471Z","shell.execute_reply":"2026-03-10T08:35:53.213325Z","shell.execute_reply.started":"2026-03-10T08:35:52.734663Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# First step is to transform the given annotations into a format the model can work with\n\nSince the predictions need to be in the form x1;x2;x3 and y1;y2;y3, let's make that the format we will generate. We also need to add special tokens to predict the chart type and separators to distinguish what is what.\n\nRemember, Donut takes an image as input and will generate text as output.","metadata":{}},{"cell_type":"code","source":"#数据预处理模块\n\n\n\n\n\n'''\n把原始的JSON标注文件，转换成Donut模型能看懂的文本格式\n\n为什么需要这一步：\nDonut是生成式模型，它需要学会输出这样的文本：\n<line><x_start>0;2;4<x_end><y_start>10;20;30<y_end>\n但原始数据长这样：\n{\n  \"chart-type\": \"line\",\n  \"data-series\": [\n    {\"x\": 0, \"y\": 10},\n    {\"x\": 2, \"y\": 20},\n    {\"x\": 4, \"y\": 30}\n  ]\n}\n所以需要一个翻译官，把JSON翻译成带特殊token的字符串，这个翻译官就是get_gt_string_and_xy函数\n'''\n\n\n\n'''\n数据预处理这部分你做了什么：\n将原始JSON标注转换成Donut模型能学习的文本格式，主要工作包括：\n1.设计特殊token：定义了<x_start>,<y_start>等标记，让模型学会结构化输出\n2.处理数值精度：实现了自适应小数精度处理\n3.数据清洗：过滤掉NaN等无效值，确保训练数据质量\n4.最终输出：生成类似<line><x_start>0;2;4<x_end><y_start>10;20;30<y_end>的目标文本，供模型学习\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Let's add chart types as special tokens and a special BOS token\n\n#//定义特殊token\n'''\n为什么需要这些token：\n模型输出是一串长文本，比如：<line><x_start>0;2;4<x_end><y_start>10;20;30<y_end>\n没有这些标记，你分不清哪里是x轴，哪里是y轴\n训练时，模型要学会在正确的地方插入这些标记\n'''\nPROMPT_TOKEN = \"<|PROMPT|>\"  #//提示符,告诉模型要开始生成了\nX_START = \"<x_start>\"        #//x轴数据开始标记\nX_END = \"<x_end>\"            #//x轴数据结束标记\nY_START = \"<y_start>\"        #//y轴数据开始标记\nY_END = \"<y_end>\"            #//y轴数据结束标记\n\n#//图表类型token\n#//这5种token对应5种图表类型,模型输出时,第一个词就是图表类型\nSEPARATOR_TOKENS = [\n    PROMPT_TOKEN,\n    X_START,\n    X_END,\n    Y_START,\n    Y_END,\n]\n\nLINE_TOKEN =  \"<line>\" \nVERTICAL_BAR_TOKEN = \"<vertical_bar>\"\nHORIZONTAL_BAR_TOKEN = \"<horizontal_bar>\"\nSCATTER_TOKEN = \"<scatter>\"\nDOT_TOKEN = \"<dot>\"\n\nCHART_TYPE_TOKENS = [\n    LINE_TOKEN,\n    VERTICAL_BAR_TOKEN,\n    HORIZONTAL_BAR_TOKEN,\n    SCATTER_TOKEN,\n    DOT_TOKEN,\n]\n\nnew_tokens = SEPARATOR_TOKENS + CHART_TYPE_TOKENS\n\n\n#//处理小数精度\n'''\n整数大的数字（如100.12345），保留1位小数就够了（100.1）\n整数小的数字（如0.12345），需要保留更多精度（0.1234）\n设计了自适应小数精度处理，根据整数部分大小动态调整小数位数，既保证精度又避免冗余\n'''\ndef round_float(value: Union[int, float, str]) -> Union[str, float]:\n    \"\"\"\n    Convert a float value to a string with the specified number of decimal places. \n    If there is more than 1 digit in the integer, then we will truncate to 1 decimal.\n    Otherwise, will truncate to 4 decimals.\n\n    Args:\n        value (int, float, str): The float value to convert\n\n    Returns:\n        str: The rounded float value as a string\n    \"\"\"\n    if isinstance(value, float):\n        value = str(value)\n\n        if \".\" in value:\n            integer, decimal = value.split(\".\")\n            if abs(float(integer)) > 1:\n                decimal = decimal[:1]\n            else:\n                decimal = decimal[:4]\n\n            value = integer + \".\" + decimal\n    return value\n\n\n\n#//过滤无效值\ndef is_nan(value: Union[int, float, str]) -> bool:\n    \"\"\"\n    Check if a value is NaN (not a number).\n\n    Args:\n        value (int, float, str): The value to check\n\n    Returns:\n        bool: True if the value is NaN, False otherwise\n    \"\"\"\n    return isinstance(value, float) and str(value) == \"nan\"\n\n\n\n\n\n#//这是整个预处理模块的核心\ndef get_gt_string_and_xy(filepath: Union[str, os.PathLike]) -> Dict[str, str]:\n    \"\"\"\n    Get the ground truth string and x-y data from the given JSON file.\n\n    Args:\n        filepath (str): The path to the JSON file\n\n    Returns:\n        dict: A dictionary containing the ground truth string, x-y data, chart type, id, and source\n    \"\"\"\n    filepath = Path(filepath)\n\n    #//1.读取JSON\n    with open(filepath) as fp:\n        data = json.load(fp)\n\n    #//2.提取数据序列\n    data_series = data[\"data-series\"] #//拿到所以数据点\n\n    all_x, all_y = [], []\n\n    #//3.处理每个数据点\n    for d in data_series:\n        x = d[\"x\"]\n        y = d[\"y\"]\n\n        x = round_float(x)\n        y = round_float(y)\n\n        # Ignore nan values\n        if is_nan(x) or is_nan(y):\n            continue\n\n        all_x.append(x)\n        all_y.append(y)\n        \n\n    #//4.组装成带token的字符串\n    chart_type = f\"<{data['chart-type']}>\"\n    x_str = X_START + \";\".join(list(map(str, all_x))) + X_END\n    y_str = Y_START + \";\".join(list(map(str, all_y))) + Y_END\n    \n    gt_string = PROMPT_TOKEN + chart_type + x_str + y_str\n\n    #//5.返回结果\n    return {\n        \"ground_truth\": gt_string,\n        \"x\": json.dumps(all_x),\n        \"y\": json.dumps(all_y),\n        \"chart-type\": data[\"chart-type\"],\n        \"id\": filepath.stem,\n        \"source\": data[\"source\"],\n    }","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:53.216823Z","iopub.status.busy":"2026-03-10T08:35:53.215974Z","iopub.status.idle":"2026-03-10T08:35:53.230246Z","shell.execute_reply":"2026-03-10T08:35:53.229106Z","shell.execute_reply.started":"2026-03-10T08:35:53.216781Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"get_gt_string_and_xy(data_dir / \"annotations\" / \"000d269c8e26.json\")","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:53.231854Z","iopub.status.busy":"2026-03-10T08:35:53.231530Z","iopub.status.idle":"2026-03-10T08:35:53.248760Z","shell.execute_reply":"2026-03-10T08:35:53.247563Z","shell.execute_reply.started":"2026-03-10T08:35:53.231820Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Creating the dataset","metadata":{}},{"cell_type":"code","source":"'''\n预处理模块        输入：单个JSON         输出：带token的字典     作用：翻译成模型能学的格式\n\n数据集构建模块     输入：所有JSON列表     输出：Dataset对象       作用：打包成训练集\n  \n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//数据集建构模块\n\n\n\n'''\n是数据预处理后的下一步-----把处理好的数据打包成PyTorch能用的格式\n把上一步生成的字典数据，转换成HuggingFace的Dataset对象，供训练时使用\n\n为什么需要这一步：\n上一步的get_gt_string_and_xy只是处理单个文件，现在有几万个JSON文件，需要：\n1.批量处理所有文件\n2.把数据组织成统一的格式\n3.让PyTorch能高效读取\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# This generator function will be used to create the Dataset object\n\n\n#//1.生成器函数\n'''\n问：你如何处理大规模训练数据：\n答：我用了Python的生成器，训练集有几万张图，如果一次性加载到内存会爆，\n生成器是惰性加载的-----每次只处理一个样本，交给模型训练，用完就释放，\n这样内存里永远只存一个样本，可以处理任意大的数据集\n'''\ndef gen_data(files: List[Union[str, os.PathLike]]) -> Dict[str, str]:\n    \"\"\"\n    This function takes a list of json files and returns a generator that yields a\n    dictionary with the ground truth string and the path to the image.\n\n    Args:\n        files (list): A list of json files\n\n    Returns:\n        generator: A generator that yields a dictionary with the ground truth string and\n            the path to the corresponding image.\n    \"\"\"\n\n    for f in files:\n\n        yield {\n            **get_gt_string_and_xy(f),\n            \"image_path\": str(images_path / f\"{f.stem}.jpg\"),\n        }\n\n\n#//2.创建Dataset\n#//从生成器创建Dataset对象\nds = Dataset.from_generator(\n    gen_data, gen_kwargs={\"files\": train_json_files}, num_proc=CFG.num_proc\n)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:53.250232Z","iopub.status.busy":"2026-03-10T08:35:53.249965Z","iopub.status.idle":"2026-03-10T08:35:56.827703Z","shell.execute_reply":"2026-03-10T08:35:56.826573Z","shell.execute_reply.started":"2026-03-10T08:35:53.250208Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Checking image sizes","metadata":{}},{"cell_type":"code","source":"#//数据探索和可视化模块\n\n\n\n'''\n检查训练集中所有图片的尺寸，确保它们大小差不多，方便后续统一处理\n\n为什么需要这一步：\nDonut模型要求输入图片固定尺寸（代码里设的560*560）\n如果原始图片大小差异太大，resize时会变形或丢失信息\n所以要先看看数据分布，确认560*560是否合理\n'''\n\n\n\n'''\n问：你如何确保数据质量：\n答：在训练前，我做了数据探索性分析，用PIL读取所有图片的尺寸，画散点图观察分布，\n发现图片宽度集中在400-600，高度集中在100-600，图片宽高比差异大，直接resize会变形，因此预处理时应random_padding补边\n\n问：为什么需要统一图片尺寸：\n答：Donut模型要求固定尺寸输入，通过分析数据分布，我确认560*560能覆盖大部分图片，且不会因过度压缩丢失关键信息\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def add_image_sizes(examples: Dict[str, Union[str, os.PathLike]]) -> Dict[str, List[int]]:\n    \"\"\"\n    This function takes a dictionary of examples and adds the width and height of the\n    image to the dictionary. This is to be used with the `Dataset.map` function.\n\n    Args:\n        examples (dict): A dictionary of examples (from `map` function)\n\n    Returns:\n        dict: The dictionary with the width and height of the image added\n    \"\"\"\n\n    sizes = [Image.open(x).size for x in examples[\"image_path\"]]\n\n    width, height = list(zip(*sizes))\n\n    return {\n        \"width\": list(width),\n        \"height\": list(height),\n    }\n\n\nds = ds.map(add_image_sizes, batched=True, num_proc=CFG.num_proc)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:56.829361Z","iopub.status.busy":"2026-03-10T08:35:56.828930Z","iopub.status.idle":"2026-03-10T08:35:56.851040Z","shell.execute_reply":"2026-03-10T08:35:56.849914Z","shell.execute_reply.started":"2026-03-10T08:35:56.829323Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Most fall within 560 x 560\nplt.figure(figsize=(10, 10))\n\nplt.scatter(x=ds[\"width\"], y=ds[\"height\"], marker=\"o\", alpha=0.5)\n\n# Add axis labels\nplt.xlabel(\"Width\")\nplt.ylabel(\"Height\")\n\n# Add a title\nplt.title(\"Image Dimensions\")\n\n# Add gridlines\nplt.grid(True)\n\n# Set aspect ratio to be equal\nplt.gca().set_aspect(\"equal\", adjustable=\"box\")\n\n# Show the plot\nplt.show()","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:56.852837Z","iopub.status.busy":"2026-03-10T08:35:56.852469Z","iopub.status.idle":"2026-03-10T08:35:57.624574Z","shell.execute_reply":"2026-03-10T08:35:57.623248Z","shell.execute_reply.started":"2026-03-10T08:35:56.852801Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\nx轴：图片宽度\ny轴：图片高度\n每个点代表一张图片\n\n图片尺寸分布在一个范围内，说明：\n1.原始图片大小不一\n2.有的宽高比不同\n3.训练时需要统一resize到560*560\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model configuration\n\nNeed to set:  \n\n- image height\n- image width\n- max sequence length to generate","metadata":{}},{"cell_type":"code","source":"'''\n数据探索模块--->模型配置模块--->加载模型模块--->训练模块\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//模型配置模块\n\n\n'''\n是在加载预训练模型之前，先设置好模型的参数\n告诉Dount模型：输入图片要resize成多大，输出文本最长能有多长\n\n为什么需要这一步：\n预训练的Donut模型有自己的默认配置（比如图片尺寸256*256）\n你的任务需要560*560的图片\n所以要覆盖默认配置，让模型适配你的数据\n'''\n\n\n'''\n问：你怎么配置模型的：\n我从HuggingFace加载Donut的默认配置后，修改了2个关键参数：\n1.图片尺寸：设为560*560\n2.最大生成长度：设为1024,经过统计，每个图表的token数大约200-300，1024留了足够余量，防止模型生成到一半被截断\n\n问：为什么不用默认配置：\nDonut的默认配置可能针对通用文档（比如论文，发票），尺寸较小，而我的图表数据需要更高分辨率才能看清轴标签和小数点\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//1.加载默认配置\nconfig = VisionEncoderDecoderConfig.from_pretrained(\"naver-clova-ix/donut-base\")\n#//2.设置图片尺寸\nconfig.encoder.image_size = (CFG.image_height, CFG.image_width)\n#//3.设置最大生成长度\n#//修改解码器BART的最大输出长度,告诉模型:生成的文本最多1024个token\nconfig.decoder.max_length = CFG.max_length\n\nprint(CFG.image_height, CFG.image_width, CFG.max_length)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:57.629708Z","iopub.status.busy":"2026-03-10T08:35:57.629041Z","iopub.status.idle":"2026-03-10T08:35:57.810112Z","shell.execute_reply":"2026-03-10T08:35:57.809133Z","shell.execute_reply.started":"2026-03-10T08:35:57.629677Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Checking tokenizer for unknowns\n","metadata":{}},{"cell_type":"code","source":"#//词表对齐模块\n\n\n\n'''\n检查你自定义的特殊token模型认不认识\n\n检查Donut的tokenizer是否认识你新加的token(如<line>,<x_start>),如果不认识就加进去\n\n为什么需要这一步：\nDonut预训练时用的词表里没有你自定义的token\n比如<line>,<x_start>这些，模型压根没见过\n如果直接训练，模型会把它们当成未知词，学不会\n'''\n\n\n\n'''\n问：你如何确保自定义token能被模型识别：\n答：我做了词表对齐检查，用 check_for_unk 函数遍历整个数据集，\n发现 3.7 万条数据都存在未知 token，主要是数字 \"1\" 和特殊 token 没被识别。\n\n通过分析，发现 Donut 的 BPE 词表里没有单独的 \"1\"，\n所以我添加了 <one> token，同时把 <line>、<x_start> 等 11 个特殊 token 都加到词表里。\n\n这一步保证了模型在训练时能正确理解我要它学习的文本格式，避免因未知 token 导致的学习偏差。\"\n\n\n问：为什么不用默认的 tokenizer：\n答：\"默认 tokenizer 只包含预训练时见过的词。\n我新增的图表类型 token 和数据分隔 token 它都没见过，\n如果不加进去，模型会把它们当成未知词，永远学不会。所以必须做词表扩展。\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//1.加载tokenizer\nprocessor = DonutProcessor.from_pretrained(\"naver-clova-ix/donut-base\")\nprocessor.image_processor.size = {\n    \"height\": CFG.image_height,\n    \"width\": CFG.image_width,\n}","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:57.811923Z","iopub.status.busy":"2026-03-10T08:35:57.811278Z","iopub.status.idle":"2026-03-10T08:35:58.358890Z","shell.execute_reply":"2026-03-10T08:35:58.357950Z","shell.execute_reply.started":"2026-03-10T08:35:57.811893Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//2.检查位置token\ndef check_for_unk(examples: Dict[str, str]) -> Dict[str, List[str]]:\n    \"\"\"\n    Check for unknown tokens in the given examples.\n\n    This function takes a dictionary containing a list of ground truth texts and\n    tokenizes them using the processor's tokenizer. It then checks for any unknown\n    tokens in the tokenized text and returns a dictionary containing a list of the\n    unknown tokens for each example.\n\n    Args:\n        examples (dict): A dictionary containing a list of ground truth texts. \n            Example: {\"ground_truth\": [\"text1\", \"text2\", ...]}\n\n    Returns:\n        dict: A dictionary containing a list of unknown tokens for each example. \n            Example: {\"unk_tokens\": [[\"unk1\", \"unk2\"], [], [\"unk3\"], ...]}\n    \"\"\"\n\n    texts = examples[\"ground_truth\"]\n\n    ids = processor.tokenizer(texts).input_ids\n    tokens = [processor.tokenizer.tokenize(x, add_special_tokens=True) for x in texts]\n\n    unk_tokens = []\n    for example_ids, example_tokens in zip(ids, tokens):\n        example_unk_tokens = []\n        for i in range(len(example_ids)):\n            if example_ids[i] == processor.tokenizer.unk_token_id:\n                example_unk_tokens.append(example_tokens[i])\n\n        unk_tokens.append(example_unk_tokens)\n\n    return {\"unk_tokens\": unk_tokens}\n\n\n#//3.应用到整个数据集\nunk = ds.map(check_for_unk, batched=True, num_proc=CFG.num_proc)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.360655Z","iopub.status.busy":"2026-03-10T08:35:58.360168Z","iopub.status.idle":"2026-03-10T08:35:58.389948Z","shell.execute_reply":"2026-03-10T08:35:58.389201Z","shell.execute_reply.started":"2026-03-10T08:35:58.360625Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Let's look at only the examples with unknown tokens\n#//4.筛选有未知token的样本\nunk = unk.filter(lambda x: len(x[\"unk_tokens\"]) > 0, num_proc=CFG.num_proc)\n\nprint(len(unk))\n\nunk[0]","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.391300Z","iopub.status.busy":"2026-03-10T08:35:58.391040Z","iopub.status.idle":"2026-03-10T08:35:58.412125Z","shell.execute_reply":"2026-03-10T08:35:58.411029Z","shell.execute_reply.started":"2026-03-10T08:35:58.391275Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Let's count which tokens show up most often as unknowns\nall_unk_tokens = [x for y in unk[\"unk_tokens\"] for x in y]\n\nCounter(all_unk_tokens)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.413862Z","iopub.status.busy":"2026-03-10T08:35:58.413346Z","iopub.status.idle":"2026-03-10T08:35:58.680696Z","shell.execute_reply":"2026-03-10T08:35:58.679544Z","shell.execute_reply.started":"2026-03-10T08:35:58.413836Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"It seems that the donut tokenizer doesn't contain \"1\" which will result in many unknown tokens in the ground truth string. There is \"_1\" (space in front of 1) and \"01\", \"11\", \"21\", \"31\", \"41\" etc. But if there is a \"0.1\" or \">1\" then the 1 will turn into an unknown token.\n\nI'll handle this in the `preprocess` function.","metadata":{}},{"cell_type":"code","source":"example_str = \"0.1 1 1990\"\n\ntemp_ids = processor.tokenizer(example_str).input_ids\nprint(\"ids:\", temp_ids)\nprint(\"tokenized:\", processor.tokenizer.tokenize(example_str))\nprint(\"decoded:\", processor.tokenizer.decode(temp_ids))\nprint(\"unk id:\", processor.tokenizer.unk_token_id)\n\n# Adding these tokens should mean that there should be very few unknown tokens\nnum_added = processor.tokenizer.add_tokens([\"<one>\"] + new_tokens)\nprint(num_added, \"tokens added\")\n\nconfig.pad_token_id = processor.tokenizer.pad_token_id\nconfig.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([PROMPT_TOKEN])[0]","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.682452Z","iopub.status.busy":"2026-03-10T08:35:58.682019Z","iopub.status.idle":"2026-03-10T08:35:58.691758Z","shell.execute_reply":"2026-03-10T08:35:58.690588Z","shell.execute_reply.started":"2026-03-10T08:35:58.682408Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Preprocessing function to get pixel values and input_ids\n\nWhen ever the `_getitem__` function is called for the dataset, it will run this function.","metadata":{}},{"cell_type":"code","source":"#//预处理模块\n#//是整个训练流程中最核心的数据准备环节---把原始图片和文本转换成模型能直接吃的tensor\n\n\n\n'''\n每次从数据集中取一个batch时，实时把图片转成tensor，把文本转成token ids，并修复之前发现的“1”问题\n为什么需要这一步：\n1.图片不能直接喂给模型，要转成tensor（数值矩阵）\n2.文本不能直接喂给模型，要转成token ids（数字序列）\n3.之前发现“1”会被当成未知token，要在这一步恢复\n'''\n\n\n\n\n\n'''\n问：你做了哪些数据预处理工作？\n答：我做了两层预处理：\n1.第一层是离线预处理，一次性把所有 JSON 文件转换成带特殊 token 的文本格式，\n比如把原始数据变成 <line><x_start>0;2;4<x_end><y_start>10;20;30<y_end>。这一步只做一次，结果存到 Dataset 里。\n\n2.第二层是在线预处理，用 set_transform 挂载到数据集上。\n每次训练取 batch 时，实时把图片转成 tensor，把文本转成 token ids\n，同时修复之前发现的数字 \"1\" 被识别为未知 token 的问题。这样既保证数据格式正确，又节省内存。\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from functools import partial\n\none_token_id = processor.tokenizer(\"<one>\", add_special_tokens=False).input_ids[0]\nunk_token_id = processor.tokenizer.unk_token_id\n\n\n#//1.修复函数\ndef replace_unk_tokens_with_one(example_ids: List[int], example_tokens: List[str], one_token_id:int, unk_token_id:int) -> List[int]:\n    \"\"\"\n    Replace unknown tokens that represent \"1\" with the correct token id.\n\n    Args:\n        example_ids (list): List of token ids for a given example\n        example_tokens (list): List of tokens for the same given example\n        one_token_id (int): Token id for the \"<one>\" token\n        unk_token_id (int): Token id for the unknown token\n\n    Returns:\n        list: The updated list of token ids with the correct token id for \"1\"\n    \"\"\"\n    \n    temp_ids = []\n    for id_, token in zip(example_ids, example_tokens):\n        if id_ == unk_token_id and token == \"1\":\n            id_ = one_token_id\n        temp_ids.append(id_)\n    return temp_ids\n\n\n#//2.主预处理函数\ndef preprocess(examples: Dict[str, str], processor: DonutProcessor, CFG: CFG) -> Dict[str, Union[torch.Tensor, List[int], List[str]]]:\n    \"\"\"\n    Preprocess the given examples.\n\n    This function processes the input examples by tokenizing the texts, replacing\n    any unknown tokens that represent \"1\" with the correct token id, and loading\n    the images.\n\n    Args:\n        examples (dict): A dictionary containing ground truth texts, image paths, and ids\n        processor: An object responsible for tokenizing texts and processing images\n        CFG: A configuration object containing settings and hyperparameters\n\n    Returns:\n        dict: A dictionary containing preprocessed images, token ids, and ids\n    \"\"\"\n    \n    pixel_values = []\n\n    texts = examples[\"ground_truth\"]\n\n    ids = processor.tokenizer(\n        texts,\n        add_special_tokens=False,\n        max_length=CFG.max_length,\n        padding=True,\n        truncation=True,\n    ).input_ids\n\n    if isinstance(texts, str):\n        texts = [texts]\n\n    tokens = [processor.tokenizer.tokenize(text, add_special_tokens=False) for text in texts]\n    \n    one_token_id = processor.tokenizer(\"<one>\", add_special_tokens=False).input_ids[0]\n    unk_token_id = processor.tokenizer.unk_token_id\n    \n    final_ids = [\n        replace_unk_tokens_with_one(example_ids, example_tokens, one_token_id, unk_token_id)\n        for example_ids, example_tokens in zip(ids, tokens)\n    ]\n\n    for sample in examples[\"image_path\"]:\n        pixel_values.append(processor(sample, random_padding=True).pixel_values)\n\n    return {\n        \"pixel_values\": torch.tensor(np.vstack(pixel_values)),\n        \"input_ids\": final_ids,\n        \"id\": examples[\"id\"],\n    }\n\n\n#//3.应用到数据集\nimage_ds = ds.cast_column(\"image_path\", ds_img())\nimage_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.694020Z","iopub.status.busy":"2026-03-10T08:35:58.693188Z","iopub.status.idle":"2026-03-10T08:35:58.744371Z","shell.execute_reply":"2026-03-10T08:35:58.743296Z","shell.execute_reply.started":"2026-03-10T08:35:58.693979Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample = image_ds[[0, 1, 2]]\n\nprint(sample[\"pixel_values\"].shape)\nprint(processor.decode(sample[\"input_ids\"][2]))\nprint(len(sample[\"input_ids\"][2]))\nprint(processor.tokenizer.convert_ids_to_tokens(sample[\"input_ids\"][2]))","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.746282Z","iopub.status.busy":"2026-03-10T08:35:58.745651Z","iopub.status.idle":"2026-03-10T08:35:58.859268Z","shell.execute_reply":"2026-03-10T08:35:58.858006Z","shell.execute_reply.started":"2026-03-10T08:35:58.746245Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Creating splits\n\n\nBasic split: break extracted into folds, all generated and some extracted into train. Only extracted in validation.","metadata":{}},{"cell_type":"code","source":"#//数据集划分模块\n\n\n\n'''\n把数据分成训练集和验证集，确保训练时模型将各种图表类型\n\n为什么需要这一步：\n数据来源有2种：extracted（真实扫描图）和generated（合成图）\n验证集必须用真实数据（extracted），才能反映模型在真实场景的表现\n'''\n\n\n\n'''\n问：你怎么划分训练集和验证集的？\n答：我做了分层划分，考虑了两个维度：\n1. 数据来源：真实数据（extracted）和合成数据（generated）。\n验证集只用真实数据，因为最终要在真实场景用；合成数据全部进训练集，帮助模型学习。\n2. 图表类型：各类图表数量不均衡，条形图多，水平条形图少。\n我用 StratifiedKFold 做4折分层抽样，确保每一折里各类图表比例和总体一致。\n最后训练集包含3折真实数据+全部合成数据，验证集用1折真实数据。这样既保证了训练数据量，又保证了验证的公平性。\n\n\n\n问：为什么不用随机划分？\n答：\"随机划分可能导致某一类图表全进验证集或全进训练集，比如水平条形图本来就少，万一全进验证集，模型就学不到这类。\n分层划分保证了每一类都在训练集和验证集中有合理分布。\n\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# separate by extracted and generated\n\n#//1.按来源分离数据\nextracted_ds = ds.filter(lambda x: x[\"source\"] == \"extracted\", num_proc=CFG.num_proc)\ngenerated_ds = ds.filter(lambda x: x[\"source\"] == \"generated\", num_proc=CFG.num_proc)\n\n#//2.查看图表类型分布\nchart_types = extracted_ds[\"chart-type\"]\n\nprint(Counter(chart_types))\n\n\n#//3.分层k折划分\nskf = StratifiedKFold(n_splits=4)\n\nfold_idxs = []\n\nfor _, val_idxs in skf.split(chart_types, y=chart_types):\n    fold_idxs.append(val_idxs)\n\n#//4.查看每一折的分布\nfor n, idxs in enumerate(fold_idxs):\n    print(Counter([chart_types[i] for i in idxs]))\n    if n > 3:\n        break","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.860902Z","iopub.status.busy":"2026-03-10T08:35:58.860531Z","iopub.status.idle":"2026-03-10T08:35:58.897080Z","shell.execute_reply":"2026-03-10T08:35:58.896259Z","shell.execute_reply.started":"2026-03-10T08:35:58.860862Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\n垂直条形图最多（457）\n折线图次之（423）\n散点图（165）\n水平条形图最少（73）\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Combine generated and some of the extracted examples into the train set\n\nfrom datasets import concatenate_datasets\n\n#//5.构建训练集\nfold = 0\n\ntrain_extracted = extracted_ds.select(\n    list(chain(*[x for i, x in enumerate(fold_idxs) if i != fold]))\n)\ntrain_ds = concatenate_datasets([train_extracted, generated_ds])\ntrain_ds = train_ds.cast_column(\"image_path\", ds_img())\ntrain_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))\n\n# Creat validation set from only extracted examples\n\n#//6.构建验证集\nval_gt_ds = extracted_ds.select(fold_idxs[fold])\nval_ds = val_gt_ds.cast_column(\"image_path\", ds_img())\nval_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))\n\ngt_chart_type = val_gt_ds[\"chart-type\"]\ngt_x = [json.loads(_) for _ in val_gt_ds[\"x\"]]\ngt_y = [json.loads(_) for _ in val_gt_ds[\"y\"]]\ngt_ids = val_gt_ds[\"id\"]\n\ni = 0\nprint(gt_chart_type[i])\nprint(gt_x[i])\nprint(gt_y[i])\nprint(gt_ids[i])\nprint(Counter(gt_chart_type))","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:58.898979Z","iopub.status.busy":"2026-03-10T08:35:58.898183Z","iopub.status.idle":"2026-03-10T08:35:59.079984Z","shell.execute_reply":"2026-03-10T08:35:59.078943Z","shell.execute_reply.started":"2026-03-10T08:35:58.898939Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Collate function to make sure the ids are all the same length in a batch","metadata":{}},{"cell_type":"code","source":"pad_token_id = processor.tokenizer.pad_token_id\n\n\ndef collate_fn(samples: List[Dict[str, Union[torch.Tensor, List[int], str]]]) -> Dict[str, Union[torch.Tensor, List[str]]]:\n    \"\"\"\n    Custom collate function for DataLoader.\n\n    This function takes a list of samples and combines them into a batch with\n    properly padded input_ids.\n\n    Args:\n        samples (List[Dict[str, Union[torch.Tensor, List[int], str]]]): \n            A list of samples, where each sample is a dictionary containing\n            \"pixel_values\" (torch.Tensor), \"input_ids\" (List[int]), and \"id\" (str).\n\n    Returns:\n        Dict[str, Union[torch.Tensor, List[str]]]: \n            A dictionary containing the combined pixel values, padded input_ids, and ids.\n    \"\"\"\n\n    batch = {}\n\n    batch[\"pixel_values\"] = torch.stack([x[\"pixel_values\"] for x in samples])\n\n    max_length = max([len(x[\"input_ids\"]) for x in samples])\n\n    # Make a multiple of 8 to efficiently use the tensor cores\n    if max_length % 8 != 0:\n        max_length = (max_length // 8 + 1) * 8\n\n    input_ids = [\n        x[\"input_ids\"] + [pad_token_id] * (max_length - len(x[\"input_ids\"]))\n        for x in samples\n    ]\n\n    labels = torch.tensor(input_ids)\n    labels[labels == pad_token_id] = -100 # ignore loss on padding tokens\n    batch[\"labels\"] = labels\n    \n    batch[\"id\"] = [x[\"id\"] for x in samples]\n\n    return batch\n","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:59.081962Z","iopub.status.busy":"2026-03-10T08:35:59.081182Z","iopub.status.idle":"2026-03-10T08:35:59.089879Z","shell.execute_reply":"2026-03-10T08:35:59.088736Z","shell.execute_reply.started":"2026-03-10T08:35:59.081931Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataloaders\n\nValidation uses generation so it is very slow. That's why I only use a small fraction of the examples.","metadata":{}},{"cell_type":"code","source":"#//数据加载模块\n#把预处理好的数据包装成PyTorch能高效读取的格式\n\n\n\n'''\n把Dataset包装成DataLoader，让训练时可以一批一批地取数据，还支持多进程加速\n\n为什么需要这一步：\n训练时不是一次处理所有数据，而是一次处理一个batch\nDataLoader负责：自动batch,打乱顺序，多进程预加载\n验证集不需要打乱，但要按顺序评估\n'''\n\n\n\n\n\n'''\n问：你是怎么加载训练数据的？\n答：我用 PyTorch 的 DataLoader 包装数据集，主要做了几件事：\n1. 批量加载：设置 batch_size=2，平衡显存和训练速度。\n2. 数据打乱：训练集 shuffle=True，防止模型记住顺序；验证集 shuffle=False，保证评估一致。\n3. 多进程加速：num_workers=2，用两个进程预加载数据，不让数据读取成为训练瓶颈。\n4. 自定义 batch 组装：通过 collate_fn 把不定长的文本 padding 到相同长度，方便模型并行计算。\n\n最后测试了一下，一个 batch 包含 [2, 3, 560, 560] 的图片 tensor 和 [2, 120] 的标签 tensor，说明 DataLoader 工作正常。\n\n\n\n问：为什么训练集要 shuffle？\nA：\"如果不 shuffle，每个 epoch 模型看到的数据顺序都一样，可能会学到数据顺序的伪规律，而不是真正的图表规律。\nshuffle 能打破这种顺序，让模型泛化能力更强。\n\n\n问：num_workers 有什么用？\n答：训练时 GPU 算得快，但数据加载慢的话，GPU 就会空等。\nnum_workers 用多个 CPU 进程提前加载数据，让 GPU 一直在算，提高训练效率。\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if CFG.debug:\n    train_ds = train_ds.select(range(100))\n    val_ds = val_ds.select(range(100))\n\n#//1.创建训练DataLoader\ntrain_dataloader = DataLoader(\n    train_ds,\n    batch_size=CFG.batch_size,  #//每次取2个样本\n    shuffle=True, #//每个epoch打乱数据,防止模型记住顺序\n    collate_fn=collate_fn, #//自定义如何把多个样本合成一个batch\n    num_workers=CFG.num_workers, #//用2个进程提前加载数据加快训练\n)\n\n#//2.创建验证DataLoader\nval_dataloader = DataLoader(\n    val_ds,\n    batch_size=CFG.batch_size,\n    shuffle=False,  #//保证每次评估的顺序一致结果可复现\n    collate_fn=collate_fn,\n    num_workers=CFG.num_workers,\n)\n\n#//3.计算总训练步数\nnum_training_steps = len(train_dataloader) * CFG.epochs // CFG.gpus\n\nbatch = next(iter(train_dataloader))\n\nbatch.keys(), [(k, v.shape) for k, v in batch.items() if k != \"id\"]","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:59.093613Z","iopub.status.busy":"2026-03-10T08:35:59.093198Z","iopub.status.idle":"2026-03-10T08:35:59.389345Z","shell.execute_reply":"2026-03-10T08:35:59.388088Z","shell.execute_reply.started":"2026-03-10T08:35:59.093577Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Functions to calculate metrics","metadata":{}},{"cell_type":"code","source":"def rmse(y_true: List[float], y_pred: List[float]) -> float:\n    \"\"\"\n    Calculate the Root Mean Square Error (RMSE) between the true and predicted values.\n\n    Args:\n        y_true (List[float]): The true values.\n        y_pred (List[float]): The predicted values.\n\n    Returns:\n        float: The Root Mean Square Error.\n    \"\"\"\n    return np.sqrt(np.mean(np.square(np.subtract(y_true, y_pred))))\n\n\ndef sigmoid(x: float) -> float:\n    \"\"\"\n    Calculate the sigmoid function for the given value.\n\n    Args:\n        x (float): The input value.\n\n    Returns:\n        float: The result of the sigmoid function.\n    \"\"\"\n    return 2 - 2 / (1 + np.exp(-x))\n\n\ndef normalized_rmse(y_true: List[float], y_pred: List[float]) -> float:\n    \"\"\"\n    Calculate the normalized Root Mean Square Error (RMSE) between the true and predicted values.\n\n    Args:\n        y_true (List[float]): The true values.\n        y_pred (List[float]): The predicted values.\n\n    Returns:\n        float: The normalized Root Mean Square Error.\n    \"\"\"\n    numerator = rmse(y_true, y_pred)\n    denominator = rmse(y_true, np.mean(y_true))\n\n    # https://www.kaggle.com/competitions/benetech-making-graphs-accessible/discussion/396947\n    if denominator == 0:\n        if numerator == 0:\n            return 1.0\n        return 0.0\n\n    return sigmoid(numerator / denominator)\n\n\ndef normalized_levenshtein_score(y_true: List[str], y_pred: List[str]) -> float:\n    \"\"\"\n    Calculate the normalized Levenshtein distance between two lists of strings.\n\n    Args:\n        y_true (List[str]): The true values.\n        y_pred (List[str]): The predicted values.\n\n    Returns:\n        float: The normalized Levenshtein distance.\n    \"\"\"\n    total_distance = np.sum([levenshtein(yt, yp) for yt, yp in zip(y_true, y_pred)])\n    length_sum = np.sum([len(yt) for yt in y_true])\n    return sigmoid(total_distance / length_sum)\n\n\ndef score_series(\n    y_true: List[Union[float, str]], y_pred: List[Union[float, str]]\n) -> float:\n    \"\"\"\n    Calculate the score for a series of true and predicted values.\n\n    Args:\n        y_true (List[Union[float, str]]): The true values.\n        y_pred (List[Union[float, str]]): The predicted values.\n\n    Returns:\n        float: The score for the series.\n    \"\"\"\n    if len(y_true) != len(y_pred):\n        return 0.0\n    if isinstance(y_true[0], str):\n        return normalized_levenshtein_score(y_true, y_pred)\n    else:\n        # Since this is a generative model, there is a chance it doesn't produce a float.\n        # In that case, we return 0.0.\n        try:\n            return normalized_rmse(y_true, list(map(float, y_pred)))\n        except:\n            return 0.0\n\n\ndef benetech_score(ground_truth: pd.DataFrame, predictions: pd.DataFrame) -> float:\n    \"\"\"Evaluate predictions using the metric from the Benetech - Making Graphs Accessible.\n\n    Parameters\n    ----------\n    ground_truth: pd.DataFrame\n        Has columns `[data_series, chart_type]` and an index `id`. Values in `data_series`\n        should be either arrays of floats or arrays of strings.\n\n    predictions: pd.DataFrame\n    \"\"\"\n    if not ground_truth.index.equals(predictions.index):\n        raise ValueError(\n            \"Must have exactly one prediction for each ground-truth instance.\"\n        )\n    if not ground_truth.columns.equals(predictions.columns):\n        raise ValueError(f\"Predictions must have columns: {ground_truth.columns}.\")\n    pairs = zip(\n        ground_truth.itertuples(index=False), predictions.itertuples(index=False)\n    )\n    scores = []\n    for (gt_series, gt_type), (pred_series, pred_type) in pairs:\n        if gt_type != pred_type:  # Check chart_type condition\n            scores.append(0.0)\n        else:  # Score with RMSE or Levenshtein as appropriate\n            scores.append(score_series(gt_series, pred_series))\n\n    ground_truth[\"score\"] = scores\n\n    grouped = ground_truth.groupby(\"chart_type\", as_index=False)[\"score\"].mean()\n\n    chart_type2score = {\n        chart_type: score\n        for chart_type, score in zip(grouped[\"chart_type\"], grouped[\"score\"])\n    }\n\n    return np.mean(scores), chart_type2score\n\n\ndef string2triplet(pred_string: str) -> Tuple[str, List[str], List[str]]:\n    \"\"\"\n    Convert a prediction string to a triplet of chart type, x values, and y values.\n\n    Args:\n        pred_string (str): The prediction string.\n\n    Returns:\n        Tuple[str, List[str], List[str]]: A triplet of chart type, x values, and y values.\n    \"\"\"\n    \n    chart_type = \"line\"\n    for tok in CHART_TYPE_TOKENS:\n        if tok in pred_string:\n            chart_type = tok.strip(\"<>\")\n\n    pred_string = re.sub(r\"<one>\", \"1\", pred_string)\n\n    x = pred_string.split(X_START)[1].split(X_END)[0].split(\";\")\n    y = pred_string.split(Y_START)[1].split(Y_END)[0].split(\";\")\n\n    if len(x) == 0 or len(y) == 0:\n        return chart_type, [], []\n\n    min_length = min(len(x), len(y))\n\n    x = x[:min_length]\n    y = y[:min_length]\n\n    return chart_type, x, y\n\n\ndef validation_metrics(val_outputs: List[str], val_ids: List[str], gt_df: pd.DataFrame) -> Dict[str, float]:\n    \"\"\"\n    Calculate validation metrics for a set of outputs, ids, and ground truth dataframe.\n\n    Args:\n        val_outputs (List[str]): A list of validation outputs.\n        val_ids (List[str]): A list of validation ids.\n        gt_df (pd.DataFrame): The ground truth dataframe.\n\n    Returns:\n        Dict[str, float]: A dictionary containing the validation scores.\n    \"\"\"\n    pred_triplets = []\n\n    for example_output in val_outputs:\n\n        if not all([x in example_output for x in [X_START, X_END, Y_START, Y_END]]):\n            pred_triplets.append((\"line\", [], []))\n        else:\n            pred_triplets.append(string2triplet(example_output))\n\n    pred_df = pd.DataFrame(\n        index=[f\"{id_}_x\" for id_ in val_ids] + [f\"{id_}_y\" for id_ in val_ids],\n        data={\n            \"data_series\": [x[1] for x in pred_triplets]\n            + [x[2] for x in pred_triplets],\n            \"chart_type\": [x[0] for x in pred_triplets] * 2,\n        },\n    )\n\n    overall_score, chart_type2score = benetech_score(\n        gt_df.loc[pred_df.index.values], pred_df\n    )\n\n    return {\n        \"val_score\": overall_score,\n        **{f\"{k}_score\": v for k, v in chart_type2score.items()},\n    }\n","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:59.391923Z","iopub.status.busy":"2026-03-10T08:35:59.391085Z","iopub.status.idle":"2026-03-10T08:35:59.413326Z","shell.execute_reply":"2026-03-10T08:35:59.412249Z","shell.execute_reply.started":"2026-03-10T08:35:59.391887Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Lightning Module\n\nIt will use generation for validation which can be very slow. If you want to utilize the two GPUs on Kaggle, you should probably not run it in a notebook because DDP gets funky when doing that.","metadata":{}},{"cell_type":"code","source":"#//模型训练封装模块\n#//用PyTorch Lightning把训练,验证,优化器都包装起来,让代码更简洁\n\n\n\n\n'''\n把Donut模型包装成一个PyTorch Lightning模块，定义好每一步要做什么，让训练过程更加规范\n'''\n\n\n\n'''\n问：你怎么组织训练代码的？\n\n答：我用 PyTorch Lightning 封装了训练流程，主要实现了几个关键方法：\n1. training_step：处理一个训练 batch，计算 loss 并记录。\n2. validation_step：处理验证 batch，用模型生成预测文本（贪婪搜索，num_beams=1），存起来。\n3. validation_epoch_end：验证结束后，调用 validation_metrics 计算指标，打印结果——你之前看到的 0.349 就是这里输出的。\n4. configure_optimizers：设置 Adam 优化器，学习率 3e-5。\n\n这样代码结构清晰，训练循环、验证循环、指标计算都分开，容易维护。\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pytorch_lightning as pl\nfrom transformers import PreTrainedTokenizerBase, PreTrainedModel\n\nclass DonutModelPLModule(pl.LightningModule):\n    def __init__(self, processor: PreTrainedTokenizerBase, model: PreTrainedModel, gt_df: pd.DataFrame, num_training_steps: int):\n        \"\"\"\n        A PyTorch Lightning module for the DonutModel.\n\n        Args:\n            processor (PreTrainedTokenizerBase): The tokenizer/processor for the model.\n            model (PreTrainedModel): The pretrained model.\n            gt_df (pd.DataFrame): The ground truth dataframe.\n            num_training_steps (int): The number of training steps.\n        \"\"\"\n        super().__init__()\n        self.processor = processor\n        self.model = model\n        self.gt_df = gt_df\n        self.num_training_steps = num_training_steps\n\n    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        pixel_values = batch[\"pixel_values\"]\n        labels = batch[\"labels\"]\n\n        outputs = self.model(pixel_values, labels=labels)\n        loss = outputs.loss\n        self.log(\"train_loss\", loss)\n        return loss\n\n    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, dataset_idx: int = 0) -> None:\n        pixel_values = batch[\"pixel_values\"]\n        batch_size = pixel_values.shape[0]\n        # we feed the prompt to the model\n        decoder_input_ids = torch.full(\n            (batch_size, 1),\n            self.model.config.decoder_start_token_id,\n            device=self.device,\n        )\n\n        outputs = self.model.generate(\n            pixel_values,\n            decoder_input_ids=decoder_input_ids,\n            max_length=CFG.max_length,\n            early_stopping=True,\n            pad_token_id=self.processor.tokenizer.pad_token_id,\n            eos_token_id=self.processor.tokenizer.eos_token_id,\n            use_cache=True,\n            num_beams=1,\n            top_k=1,\n            bad_words_ids=[[self.processor.tokenizer.unk_token_id]],\n            return_dict_in_generate=True,\n        )\n\n        self.val_outputs.extend(\n            self.processor.tokenizer.batch_decode(outputs.sequences)\n        )\n        self.val_ids.extend(batch[\"id\"])\n\n    def on_validation_start(self) -> None:\n        self.val_outputs, self.val_ids = [], []\n\n    def validation_epoch_end(self, outputs: List[Any]) -> None:\n\n        metrics = validation_metrics(self.val_outputs, self.val_ids, self.gt_df)\n        print(\"\\n\", metrics)\n\n        self.log_dict(metrics)\n\n        self.val_outputs, self.val_ids = [], []\n\n    def configure_optimizers(self) -> torch.optim.Optimizer:\n        \n        optimizer = torch.optim.Adam(self.parameters(), lr=CFG.lr)\n\n        return optimizer","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:35:59.417051Z","iopub.status.busy":"2026-03-10T08:35:59.416273Z","iopub.status.idle":"2026-03-10T08:36:00.460304Z","shell.execute_reply":"2026-03-10T08:36:00.458459Z","shell.execute_reply.started":"2026-03-10T08:35:59.417023Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create model and ground truth dataframe","metadata":{}},{"cell_type":"code","source":"'''\n问：你怎么初始化模型的？\n\n答：我做了四步：\n1. 准备验证集数据：把验证集的 x、y 数据从 JSON 字符串转回列表，按 id_x 和 id_y 的格式建好 DataFrame，\n后面算指标时用。\n2. 加载预训练模型：用 from_pretrained 加载 Donut-base，\n同时传入之前配好的 config，确保模型按 560x560 处理图片，最大生成长度 1024。\n3. 调整词表：因为加了 11 个新 token，词表变大了，\n用 resize_token_embeddings 让模型的 embedding 层适配新词表，这样新 token 才能被学习。\n4. 封装成 Lightning 模块：把所有东西（processor、model、gt_df、训练步数）包装进自定义的 LightningModule，\n方便用 Trainer 训练。\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"gt_chart_type = val_gt_ds[\"chart-type\"]\ngt_x = [json.loads(_) for _ in val_gt_ds[\"x\"]]\ngt_y = [json.loads(_) for _ in val_gt_ds[\"y\"]]\ngt_ids = val_gt_ds[\"id\"]\n\nindex = [f\"{id_}_x\" for id_ in gt_ids] + [f\"{id_}_y\" for id_ in gt_ids]\ngt_df = pd.DataFrame(\n    index=index,\n    data={\n        \"data_series\": gt_x + gt_y,\n        \"chart_type\": gt_chart_type * 2,\n    },\n)\n\n\nmodel = VisionEncoderDecoderModel.from_pretrained(\n    \"naver-clova-ix/donut-base\", config=config\n)\nmodel.decoder.resize_token_embeddings(len(processor.tokenizer))\nmodel_module = DonutModelPLModule(processor, model, gt_df, num_training_steps)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:36:00.461925Z","iopub.status.busy":"2026-03-10T08:36:00.461574Z","iopub.status.idle":"2026-03-10T08:36:07.466186Z","shell.execute_reply":"2026-03-10T08:36:07.465336Z","shell.execute_reply.started":"2026-03-10T08:36:00.461886Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train!","metadata":{}},{"cell_type":"code","source":"#//配置PyTorch Lightning训练器,启动训练,训练完成后保存模型\n\n\n\n\n'''\n问：你是怎么训练模型的？\n\n答：我用 PyTorch Lightning 的 Trainer 来管理训练流程，主要配置了：\n\n1. 硬件设置：用 1 块 T4 GPU，开启混合精度训练（precision=16），训练速度翻倍，显存占用减半。\n2. 训练参数：训练 2 个 epoch，每轮结束做一次验证。设置梯度裁剪为 1.0，防止梯度爆炸。\n3. 模型保存：用 ModelCheckpoint 自动保存最好的模型到 output 文件夹。\n\n训练了大约 6 小时后，最终验证集得分为 0.349。其中折线图和条形图表现较好（0.44），散点图只有 0.0098，说明散点图是难点\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#//1.设置回调函数和日志\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import WandbLogger\n\n#//模型检查点: 每轮训练完自动保存最好的模型\ncheckpoint_cb = ModelCheckpoint(CFG.output_path)\n\nloggers = []\nif CFG.use_wandb:\n    import wandb\n    from kaggle_secrets import UserSecretsClient\n    user_secrets = UserSecretsClient()\n    key = user_secrets.get_secret(\"wandb\")\n    wandb.login(key=key)\n    \n    loggers.append(WandbLogger(project=\"benetech\"))\n\n\n#//2.配置Trainer\ntrainer = pl.Trainer(\n        accelerator=\"gpu\", #//用GPU训练\n        devices=1,  #//只用1块GPU(T4)\n        max_epochs=CFG.epochs, #//训练2轮\n        val_check_interval=CFG.val_check_interval,\n        check_val_every_n_epoch=CFG.check_val_every_n_epoch,\n        gradient_clip_val=CFG.gradient_clip_val,\n        precision=16, # if you have tensor cores (t4, v100, a100, etc.) training will be 2x faster\n        num_sanity_val_steps=5,\n        callbacks=[checkpoint_cb], \n        logger=loggers\n)\n\n#//3.开始训练\ntrainer.fit(model_module, train_dataloaders=train_dataloader)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T08:36:07.467728Z","iopub.status.busy":"2026-03-10T08:36:07.467346Z","iopub.status.idle":"2026-03-10T14:07:54.132069Z","shell.execute_reply":"2026-03-10T14:07:54.130227Z","shell.execute_reply.started":"2026-03-10T08:36:07.467700Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"I interrupted the run because I was impatient","metadata":{}},{"cell_type":"code","source":"trainer.validate(model_module, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T14:26:32.528813Z","iopub.status.busy":"2026-03-10T14:26:32.528457Z","iopub.status.idle":"2026-03-10T14:40:39.848097Z","shell.execute_reply":"2026-03-10T14:40:39.846751Z","shell.execute_reply.started":"2026-03-10T14:26:32.528783Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Could also use the checkpoint callback\nmodel_module.model.save_pretrained(CFG.output_path)\nmodel_module.processor.save_pretrained(CFG.output_path)","metadata":{"execution":{"iopub.execute_input":"2026-03-10T14:40:44.330862Z","iopub.status.busy":"2026-03-10T14:40:44.330621Z","iopub.status.idle":"2026-03-10T14:40:45.548227Z","shell.execute_reply":"2026-03-10T14:40:45.547312Z","shell.execute_reply.started":"2026-03-10T14:40:44.330840Z"},"trusted":true},"outputs":[],"execution_count":null}]}