{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -Uqq datasets transformers polyleven ","metadata":{"execution":{"iopub.status.busy":"2023-05-19T16:05:16.132354Z","iopub.execute_input":"2023-05-19T16:05:16.133813Z","iopub.status.idle":"2023-05-19T16:05:42.321595Z","shell.execute_reply.started":"2023-05-19T16:05:16.133753Z","shell.execute_reply":"2023-05-19T16:05:42.320226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\nimport json\nfrom collections import Counter\nfrom itertools import chain\nfrom pathlib import Path\nfrom typing import List, Union\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom datasets import Dataset, concatenate_datasets\nfrom PIL import Image\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import DataLoader\nimport pytorch_lightning as pl\nfrom transformers import (\n DonutProcessor,\n VisionEncoderDecoderConfig,\n VisionEncoderDecoderModel,\n)\nfrom datasets import Image as ds_img\nfrom polyleven import levenshtein # a faster version of levenshtein\ndata_dir = Path(\"/kaggle/input/benetech-making-graphs-accessible/train\")\nimages_path = data_dir / \"images\"\ntrain_json_files = list((data_dir / \"annotations\").glob(\"*.json\"))\nclass CFG:\n debug = False\n num_proc = 2\n num_workers = 2\n gpus = 2\n max_length = 1024\n image_height = 560\n image_width = 560\n epochs = 1\n val_check_interval = 1.0\n check_val_every_n_epoch = 1\n gradient_clip_val = 1.0\n lr = 3e-5\n seed = 42\n warmup_steps = 300\n output_path = \"output\"\n log_steps = 200\n batch_size = 3\nBOS_TOKEN = \"<|BOS|>\"\nX_START = \"\"\nX_END = \"\"\nY_START = \"\"\nY_END = \"\"\nnew_tokens = [\n \"\",\n \"\",\n \"\",\n \"\",\n \"\",\n X_START,\n X_END,\n Y_START,\n Y_END,\n BOS_TOKEN,\n]\ndef get_gt_string_and_xy(filepath):\n with open(filepath) as fp:\n data = json.load(fp)\n data_series = data[\"data-series\"]\n all_x, all_y = [], []\n for d in data_series:\n x = d[\"x\"]\n y = d[\"y\"]\n if isinstance(x, float):\n x = str(x)\n if \".\" in x:\n integer, decimal = x.split(\".\")\n if abs(float(integer)) > 1:\n decimal = decimal[:1]\n else:\n decimal = decimal[:4]\n x = float(integer + \".\" + decimal)\n if isinstance(y, float):\n y = str(y)\n if \".\" in y:\n integer, decimal = y.split(\".\")\n if abs(float(integer)) > 1:\n decimal = decimal[:1]\n else:\n decimal = decimal[:4]\n y = float(integer + \".\" + decimal)\n if (str(x) == \"nan\" and isinstance(x, float)) or (\n str(y) == \"nan\" and isinstance(y, float)\n ):\n continue\n all_x.append(x)\n all_y.append(y)\n chart_type = \"<\" + 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 = BOS_TOKEN + chart_type + x_str + y_str\n\n return {\n \"ground_truth\": gt_string,\n \"x\": json.dumps(\n all_x\n ), # must have all lists be str rather than some str and some float\n \"y\": json.dumps(all_y),\n \"chart-type\": data[\"chart-type\"],\n \"id\": filepath.stem,\n \"source\": data[\"source\"],\n }\ndef gen_data(files):\n for f in files:\n yield {\n **get_gt_string_and_xy(f),\n \"image_path\": str(images_path / f\"{f.stem}.jpg\"),\n }\nds = Dataset.from_generator(\n gen_data, gen_kwargs={\"files\": train_json_files}, num_proc=CFG.num_proc\n)\nBOS_TOKEN = \"<|BOS|>\"\nX_START = \"\"\nX_END = \"\"\nY_START = \"\"\nY_END = \"\"\nnew_tokens = [\n \"\",\n \"\",\n \"\",\n \"\",\n \"\",\n X_START,\n X_END,\n Y_START,\n Y_END,\n BOS_TOKEN,\n]\ndef get_gt_string_and_xy(filepath):\n with open(filepath) as fp:\n data = json.load(fp)\n data_series = data[\"data-series\"]\n all_x, all_y = [], []\n for d in data_series:\n x = d[\"x\"]\n y = d[\"y\"]\n if isinstance(x, float):\n x = str(x)\n if \".\" in x:\n integer, decimal = x.split(\".\")\n if abs(float(integer)) > 1:\n decimal = decimal[:1]\n else:\n decimal = decimal[:4]\n x = float(integer + \".\" + decimal)\n if isinstance(y, float):\n y = str(y)\n if \".\" in y:\n integer, decimal = y.split(\".\")\n if abs(float(integer)) > 1:\n decimal = decimal[:1]\n else:\n decimal = decimal[:4]\n y = float(integer + \".\" + decimal)\n if (str(x) == \"nan\" and isinstance(x, float)) or (\n str(y) == \"nan\" and isinstance(y, float)\n ):\n continue\n all_x.append(x)\n all_y.append(y)\n chart_type = \"<\" + 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 gt_string = BOS_TOKEN + chart_type + x_str + y_str\n return {\n \"ground_truth\": gt_string,\n \"x\": json.dumps(\n all_x\n ),\n \"y\": json.dumps(all_y),\n \"chart-type\": data[\"chart-type\"],\n \"id\": filepath.stem,\n \"source\": data[\"source\"],\n }\ndef add_image_sizes(examples):\n sizes = [Image.open(x).size for x in examples[\"image_path\"]]\n width, height = list(zip(*sizes))\n return {\n \"width\": list(width),\n \"height\": list(height),\n }\nds = ds.map(add_image_sizes, batched=True, num_proc=CFG.num_proc)\nconfig = VisionEncoderDecoderConfig.from_pretrained(\"naver-clova-ix/donut-base\")\nconfig.encoder.image_size = (CFG.image_height, CFG.image_width)\nconfig.decoder.max_length = CFG.max_length\nprocessor = DonutProcessor.from_pretrained(\"naver-clova-ix/donut-base\")\nprocessor.image_processor.size = {\n \"height\": CFG.image_height,\n \"width\": CFG.image_width,\n}\ndef check_for_unk(examples):\n texts = examples[\"ground_truth\"]\n ids = processor.tokenizer(texts).input_ids\n tokens = [processor.tokenizer.tokenize(x, add_special_tokens=True) for x in texts]\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 unk_tokens.append(example_unk_tokens)\n return {\"unk_tokens\": unk_tokens}\nunk = ds.map(check_for_unk, batched=True, num_proc=CFG.num_proc)\nunk = unk.filter(lambda x: len(x[\"unk_tokens\"]) > 0, num_proc=CFG.num_proc)\nall_unk_tokens = [x for y in unk[\"unk_tokens\"] for x in y]\nnum_added = processor.tokenizer.add_tokens([\"\"] + new_tokens)\nconfig.pad_token_id = processor.tokenizer.pad_token_id\nconfig.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([BOS_TOKEN])[\n 0\n]\nconfig.bos_token_id = processor.tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\none_token_id = processor.tokenizer(\"\", add_special_tokens=False).input_ids[0]\nunk_token_id = processor.tokenizer.unk_token_id\ndef preprocess(examples):\n pixel_values = []\n texts = examples[\"ground_truth\"]\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 if isinstance(texts, str):\n texts = [texts]\n tokens = []\n for text in texts:\n tokens.append(processor.tokenizer.tokenize(text, add_special_tokens=False))\n final_ids = []\n for example_ids, example_tokens in zip(ids, tokens):\n if unk_token_id in example_ids:\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 final_ids.append(temp_ids)\n else:\n final_ids.append(example_ids)\n\n for sample in examples[\"image_path\"]:\n pixel_values.append(processor(sample, random_padding=True).pixel_values)\n return {\n \"pixel_values\": torch.tensor(np.vstack(pixel_values)),\n \"input_ids\": final_ids,\n \"id\": examples[\"id\"],\n }\nimage_ds = ds.cast_column(\"image_path\", ds_img())\nimage_ds.set_transform(preprocess)\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)\nchart_types = extracted_ds[\"chart-type\"]\nskf = StratifiedKFold(n_splits=4)\nfold_idxs = []\nfor _, val_idxs in skf.split(chart_types, y=chart_types):\n fold_idxs.append(val_idxs)\nfor n, idxs in enumerate(fold_idxs):\n print(Counter([chart_types[i] for i in idxs]))\n if n > 3:\n break\nfold = 0\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(preprocess)\nval_gt_ds = extracted_ds.select(fold_idxs[fold])\nval_ds = val_gt_ds.cast_column(\"image_path\", ds_img())\nval_ds.set_transform(preprocess)\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\"]\npad_token_id = processor.tokenizer.pad_token_id\ndef collate_fn(samples):\n batch = {}\n batch[\"pixel_values\"] = torch.stack([x[\"pixel_values\"] for x in samples])\n max_length = max([len(x[\"input_ids\"]) for x in samples])\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 input_ids = [\n x[\"input_ids\"] + [pad_token_id] * (max_length - len(x[\"input_ids\"]))\n for x in samples\n ]\n batch[\"labels\"] = torch.tensor(input_ids)\n batch[\"id\"] = [x[\"id\"] for x in samples]\n return batch\nif CFG.debug:\n train_ds = train_ds.select(range(100))\n val_ds = val_ds.select(range(100))\ntrain_dataloader = DataLoader(\n train_ds,\n batch_size=CFG.batch_size,\n shuffle=True,\n collate_fn=collate_fn,\n num_workers=CFG.num_workers,\n)\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)\nbatch = next(iter(train_dataloader))\ndef rmse(y_true: List[float], y_pred: List[float]) -> float:\n return np.sqrt(np.mean(np.square(np.subtract(y_true, y_pred))))\ndef sigmoid(x: float) -> float:\n return 2 - 2 / (1 + np.exp(-x))\ndef normalized_rmse(y_true: List[float], y_pred: List[float]) -> float:\n numerator = rmse(y_true, y_pred)\n denominator = rmse(y_true, np.mean(y_true))\n if denominator == 0:\n if numerator == 0:\n return 1.0\n return 0.0\n return sigmoid(numerator / denominator)\ndef normalized_levenshtein_score(y_true: List[str], y_pred: List[str]):\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)\ndef score_series(\n y_true: List[Union[float, str]], y_pred: List[Union[float, str]]\n) -> float:\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 try:\n return normalized_rmse(y_true, list(map(float, y_pred)))\n except:\n return 0.0\ndef benetech_score(ground_truth: pd.DataFrame, predictions: pd.DataFrame) -> float:\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:\n scores.append(0.0)\n else:\n scores.append(score_series(gt_series, pred_series))\n ground_truth[\"score\"] = scores\n grouped = ground_truth.groupby(\"chart_type\", as_index=False)[\"score\"].mean()\n chart_type2score = {\n chart_type: score\n for chart_type, score in zip(grouped[\"chart_type\"], grouped[\"score\"])\n }\n return np.mean(scores), chart_type2score\ndef string2triplet(pred_string: str):\n if \"\" in pred_string:\n chart_type = \"dot\"\n elif \"\" in pred_string:\n chart_type = \"horizontal_bar\"\n elif \"\" in pred_string:\n chart_type = \"vertical_bar\"\n elif \"\" in pred_string:\n chart_type = \"scatter\"\n elif \"\" in pred_string:\n chart_type = \"line\"\n else:\n return \"line\", [], []\n pred_string = re.sub(r\"\", \"1\", pred_string)\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 if len(x) == 0 or len(y) == 0:\n return chart_type, [], []\n min_length = min(len(x), len(y))\n x = x[:min_length]\n y = y[:min_length]\n return chart_type, x, y\ndef validation_metrics(val_outputs, val_ids, gt_df):\n pred_triplets = []\n for example_output in val_outputs:\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 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 overall_score, chart_type2score = benetech_score(\n gt_df.loc[pred_df.index.values], pred_df\n )\n return {\n \"val_score\": overall_score,\n **{f\"{k}_score\": v for k, v in chart_type2score.items()},\n }\nclass DonutModelPLModule(pl.LightningModule):\n def __init__(self, processor, model, gt_df):\n super().__init__()\n self.processor = processor\n self.model = model\n self.gt_df = gt_df\n def training_step(self, batch, batch_idx):\n pixel_values = batch[\"pixel_values\"]\n labels = batch[\"labels\"]\n outputs = self.model(pixel_values, labels=labels)\n loss = outputs.loss\n self.log(\"train_loss\", loss, sync_dist=True)\n return loss\n def validation_step(self, batch, batch_idx, dataset_idx=0):\n pixel_values = batch[\"pixel_values\"]\n batch_size = pixel_values.shape[0]\n decoder_input_ids = torch.full(\n (batch_size, 1),\n self.model.config.decoder_start_token_id,\n device=self.device,\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=10,\n bad_words_ids=[[self.processor.tokenizer.unk_token_id]],\n return_dict_in_generate=True,\n )\n self.val_outputs.extend(\n self.processor.tokenizer.batch_decode(outputs.sequences)\n )\n self.val_ids.extend(batch[\"id\"])\n def on_validation_start(self) -> None:\n self.val_outputs, self.val_ids = [], []\n def on_validation_epoch_end(self):\n val_outputs = self.val_outputs\n val_ids = self.val_ids\n metrics = validation_metrics(val_outputs, val_ids, self.gt_df)\n print(\"\\n\", metrics)\n self.log_dict(metrics, sync_dist=True)\n val_outputs = self.val_outputs\n self.val_outputs, self.val_ids = [], []\n def configure_optimizers(self):\n optimizer = torch.optim.Adam(self.parameters(), lr=CFG.lr)\n return optimizer\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\"]\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)\nmodel = VisionEncoderDecoderModel.from_pretrained(\n \"naver-clova-ix/donut-base\",ignore_mismatched_sizes=True, config=config\n)\nmodel.decoder.resize_token_embeddings(len(processor.tokenizer))\nmodel_module = DonutModelPLModule(processor, model, gt_df)\ntorch.set_float32_matmul_precision('medium')\ntrainer = pl.Trainer(\n accelerator=\"gpu\",\n devices=2,\n max_epochs=CFG.epochs,\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-mixed',\n num_sanity_val_steps=0,\n callbacks=[],\n strategy='ddp_notebook',\n)\ntrainer.fit(model_module, train_dataloaders=train_dataloader,\n val_dataloaders=val_dataloader)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del trainer, index, gt_df, gt_x, gt_y, gt_ids, gt_chart_type, batch, image_ds, val_ds, generated_ds, extracted_ds, ds\nmodel_module.model.save_pretrained(CFG.output_path)\nmodel_module.processor.save_pretrained(CFG.output_path)","metadata":{},"execution_count":null,"outputs":[]}]}