{"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":"%%capture\n! pip install git+https://github.com/huggingface/transformers.git","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:42:29.450853Z","iopub.execute_input":"2023-06-19T15:42:29.451119Z","iopub.status.idle":"2023-06-19T15:43:17.184195Z","shell.execute_reply.started":"2023-06-19T15:42:29.451094Z","shell.execute_reply":"2023-06-19T15:43:17.182763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install git+https://github.com/huggingface/accelerate.git","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:43:17.186464Z","iopub.execute_input":"2023-06-19T15:43:17.186779Z","iopub.status.idle":"2023-06-19T15:43:42.643865Z","shell.execute_reply.started":"2023-06-19T15:43:17.186750Z","shell.execute_reply":"2023-06-19T15:43:42.642632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install peft","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:43:42.650975Z","iopub.execute_input":"2023-06-19T15:43:42.651508Z","iopub.status.idle":"2023-06-19T15:43:54.529713Z","shell.execute_reply.started":"2023-06-19T15:43:42.651461Z","shell.execute_reply":"2023-06-19T15:43:54.528529Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nConfig = {\n    'IMAGE_DIR': '/kaggle/input/benetech-making-graphs-accessible/train/images/',\n    'MAX_PATCHES': 1024,\n    #'MODEL_NAME': 'ybelkada/pix2struct-base',\n    'MODEL_NAME': 'google/deplot',\n    'IMG_SIZE': (256, 256),\n    'MAX_LEN': 256,\n    'LR': 3e-5,\n    'NB_EPOCHS': 1,\n    'TRAIN_BS': 8,\n    'VALID_BS': 2,\n    '_wandb_kernel': 'tanaym',\n}\n","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:43:54.531225Z","iopub.execute_input":"2023-06-19T15:43:54.531536Z","iopub.status.idle":"2023-06-19T15:43:54.537885Z","shell.execute_reply.started":"2023-06-19T15:43:54.531506Z","shell.execute_reply":"2023-06-19T15:43:54.536893Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import GradScaler, autocast\nfrom transformers.optimization import Adafactor\n#from datasets import load_dataset\nfrom torch.utils.data import DataLoader, Dataset\nfrom peft import PeftModel, get_peft_config\nfrom transformers import (\n    AutoProcessor,\n    Pix2StructConfig,\n    Pix2StructForConditionalGeneration,\n    get_linear_schedule_with_warmup,\n)\nfrom random import shuffle\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nimport warnings\nwarnings.simplefilter(\"ignore\")\n# Let's define the LoraConfig\n\n    \nconfig = {\n             \"peft_type\": \"LORA\",\n             \"task_type\": \"VISION_2_SEQ\",\n             \"inference_mode\": False,\n             \"r\": 8,\n             \"target_modules\": [\"query\", \"value\"],\n             \"lora_alpha\": 32,\n             \"lora_dropout\": 0.01,\n             \"fan_in_fan_out\": False,\n             \"bias\": \"none\",\n            }\n\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-19T15:51:34.499758Z","iopub.execute_input":"2023-06-19T15:51:34.500143Z","iopub.status.idle":"2023-06-19T15:51:34.537400Z","shell.execute_reply.started":"2023-06-19T15:51:34.500115Z","shell.execute_reply":"2023-06-19T15:51:34.535715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def augments():\n    return A.Compose([\n        #A.Resize(width=Config['IMG_SIZE'][0], height=Config['IMG_SIZE'][1]),\n       #A.Normalize(\n        #mean=[0, 0, 0],\n        #std=[1, 1, 1],\n        #max_pixel_value=255,\n        #),\n        ToTensorV2(),\n    ])","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:44:08.533061Z","iopub.execute_input":"2023-06-19T15:44:08.533380Z","iopub.status.idle":"2023-06-19T15:44:08.539444Z","shell.execute_reply.started":"2023-06-19T15:44:08.533350Z","shell.execute_reply":"2023-06-19T15:44:08.538485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nclass BeneTechDataset(Dataset):\n    def __init__(self, dataset, processor, augments=None):\n        self.dataset = dataset\n        self.processor = processor\n        self.augments = augments\n\n    def __len__(self):\n        return len(self.dataset)\n\n    def __getitem__(self, idx):\n        item = self.dataset[idx]\n\n        image = cv2.imread('/kaggle/input/benetech-making-graphs-accessible/train/images/'+item['image'])\n        \n        if self.augments:\n            image = self.augments(image=image)['image']\n            \n       \n        # encoding = self.processor(\n            \n         #   im#ages=image,\n          #  retu#rn_tensors=\"pt\", \n           # add_sp#ecial_tokens=True, \n           # max_pat#ches=Config['MAX_PATCHES']\n            #,text=\"Generate underlying data table of the figure below:\"\n        #)\n        encoding= self.processor(\n                images=image, return_tensors=\"pt\", max_patches=Config['MAX_PATCHES'],text=\"Generate underlying data table of the figure below:\",is_vqa=True,do_normalize=True,font_path=\"/kaggle/input/deplot-fonts/Arial.TTF\")\n            \n        \n        encoding = {k:v.squeeze() for k,v in encoding.items()}\n        \n        encoding[\"text\"] = '<pad>'+item[\"label\"]+'</s>'\n        \n        return encoding","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:44:08.541137Z","iopub.execute_input":"2023-06-19T15:44:08.541826Z","iopub.status.idle":"2023-06-19T15:44:08.555043Z","shell.execute_reply.started":"2023-06-19T15:44:08.541791Z","shell.execute_reply":"2023-06-19T15:44:08.553963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model():\n    processor = AutoProcessor.from_pretrained(Config['MODEL_NAME'])\n    model = Pix2StructForConditionalGeneration.from_pretrained(Config['MODEL_NAME'],is_decoder=True)\n    \n    return processor, model\n\ndef collator(batch):\n    new_batch = {\"flattened_patches\":[], \"attention_mask\":[]}\n    texts = [item[\"text\"] for item in batch]\n    \n    \n    text_inputs =processor.tokenizer(\n                text=texts,\n                add_special_tokens=True,\n                padding=\"max_length\",\n                truncation=True,\n                max_length=Config['MAX_LEN'],\n                return_tensors=\"pt\")\n        \n    new_batch[\"labels\"] = text_inputs.input_ids\n    for item in batch:\n        new_batch[\"flattened_patches\"].append(item[\"flattened_patches\"])\n        new_batch[\"attention_mask\"].append(item[\"attention_mask\"])\n    new_batch[\"flattened_patches\"] = torch.stack(new_batch[\"flattened_patches\"])\n    new_batch[\"attention_mask\"] = torch.stack(new_batch[\"attention_mask\"])\n\n    return new_batch\n","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:48:41.221989Z","iopub.execute_input":"2023-06-19T15:48:41.222376Z","iopub.status.idle":"2023-06-19T15:48:41.230509Z","shell.execute_reply.started":"2023-06-19T15:48:41.222348Z","shell.execute_reply":"2023-06-19T15:48:41.229569Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from random import shuffle\n\nwith open(\"/kaggle/input/benetechjsonfromjson/out\", \"r\") as fl:\n    dataset = json.load(fl)\n    dataset1=[d for d in dataset if (d['chart-type'] != 'scatter'and d['source']=='extracted')]    \n    # Shuffle the dataset and select however samples you want for training\n    shuffle(dataset)\n    \n    dataset = [d for d in dataset if (d['chart-type'] != 'scatter')]\n    \n    # We are splitting the data naively for now\n    split = 0.003\n    train_samples = int(len(dataset) * split)\n    train_ds = dataset[:train_samples]\n    valid_ds = dataset[train_samples:]\n    shuffle(train_ds)\n    processor, model = get_model()\n    #model.to(device)\n    \n    train_dataset = BeneTechDataset(train_ds, processor, augments=augments())\n    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=Config['TRAIN_BS'], collate_fn=collator)\n    \n    valid_dataset = BeneTechDataset(valid_ds, processor, augments=augments())\n    valid_dataloader = DataLoader(valid_dataset, shuffle=False, batch_size=Config['VALID_BS'], collate_fn=collator)\n    \n    nb_train_steps = int(train_samples / Config['TRAIN_BS'] * Config['NB_EPOCHS'])\n    \n    print(len(dataset1))\n    \n    print(f\"Training on {len(train_ds)} samples, Validating on {len(valid_ds)} samples\")","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:49:11.746194Z","iopub.execute_input":"2023-06-19T15:49:11.746568Z","iopub.status.idle":"2023-06-19T15:49:15.976555Z","shell.execute_reply.started":"2023-06-19T15:49:11.746538Z","shell.execute_reply":"2023-06-19T15:49:15.975586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import wandb\ndef wandb_log(**kwargs):\n    for k, v in kwargs.items():\n        wandb.log({k: v})\n\n# Start W&B logging\n# W&B Login\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"wandb\")\n\nwandb.login(key=wb_key)\n\nrun = wandb.init(\n    project='pytorch',\n    config=Config,\n    group='deplot-peft',\n    job_type='train',\n)","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:49:15.979554Z","iopub.execute_input":"2023-06-19T15:49:15.980300Z","iopub.status.idle":"2023-06-19T15:49:50.833240Z","shell.execute_reply.started":"2023-06-19T15:49:15.980265Z","shell.execute_reply":"2023-06-19T15:49:50.832387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"peft_config = get_peft_config(config)\n\n#peft_model = PeftModel(model, peft_config,'/kaggle/input/continuinglora/adapter_model (1)')\npeft_model = PeftModel(model, peft_config)\npeft_model.print_trainable_parameters()","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:51:38.394890Z","iopub.execute_input":"2023-06-19T15:51:38.395254Z","iopub.status.idle":"2023-06-19T15:51:39.384352Z","shell.execute_reply.started":"2023-06-19T15:51:38.395227Z","shell.execute_reply":"2023-06-19T15:51:39.383434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model=peft_model","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:51:41.019431Z","iopub.execute_input":"2023-06-19T15:51:41.019806Z","iopub.status.idle":"2023-06-19T15:51:41.028649Z","shell.execute_reply.started":"2023-06-19T15:51:41.019774Z","shell.execute_reply":"2023-06-19T15:51:41.027575Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2\nimport transformers\nfrom transformers import Trainer\nclass CordTrainer(Trainer):\n    \n    def get_train_dataloader(self):\n        return train_dataloader\n    def get_test_dataloader(self):\n        return valid_dataloader\n    \ntrainer = CordTrainer(\n    model=model, \n    args=transformers.TrainingArguments(\n        per_device_train_batch_size=4, \n        gradient_accumulation_steps=1,\n        warmup_steps=50, \n        max_steps=200, \n        learning_rate=6e-5, \n        fp16=True,\n        save_steps=500,\n        logging_steps=1, \n        output_dir='outputs',\n        #save_strategy='epoch'\n    ),\n    data_collator=collator,\n)\nmodel.config.use_cache = False  # silence the warnings. Please re-enable for inference!\ntrainer.train()","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:51:43.644145Z","iopub.execute_input":"2023-06-19T15:51:43.644513Z","iopub.status.idle":"2023-06-19T15:53:44.379173Z","shell.execute_reply.started":"2023-06-19T15:51:43.644485Z","shell.execute_reply":"2023-06-19T15:53:44.377622Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.save_pretrained(\"/kaggle/working/output_dir\") ","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:44:08.828325Z","iopub.status.idle":"2023-06-19T15:44:08.829009Z","shell.execute_reply.started":"2023-06-19T15:44:08.828769Z","shell.execute_reply":"2023-06-19T15:44:08.828792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nmodel.config.use_cache=True\nimage = Image.open('/kaggle/input/benetech-making-graphs-accessible/test/images/000b92c3b098.jpg')\n\ninputs = processor(images=image, text=\"Generate underlying data table of the figure below:\", return_tensors=\"pt\").to(device)\npredictions = model.generate(**inputs, max_new_tokens=512)\nprint(processor.decode(predictions[0], skip_special_tokens=True))","metadata":{"execution":{"iopub.status.busy":"2023-06-19T15:44:08.830243Z","iopub.status.idle":"2023-06-19T15:44:08.830955Z","shell.execute_reply.started":"2023-06-19T15:44:08.830688Z","shell.execute_reply":"2023-06-19T15:44:08.830740Z"},"trusted":true},"execution_count":null,"outputs":[]}]}