demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Vision-Language Tookit (VLTK)\n",
    "\n",
    "* Define FRCNN Adapter\n",
    "* Define Vision Dataset Adapters\n",
    "    * COCO and Visual Genome\n",
    "* Define Vision-Language Dataset Adapters\n",
    "    * VQA, GQA, and Visual Genome QA\n",
    "* Extract Datasets for Each User-Defined Adapter Class\n",
    "* Register User-Defined Adapters and Config  with VLTK to Superset Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Obtaining file:///home/eltoto/vltk\n",
      "Requirement already satisfied: appdirs in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.4.4)\n",
      "Requirement already satisfied: astroid in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.4.2)\n",
      "Requirement already satisfied: black in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (20.8b1)\n",
      "Requirement already satisfied: certifi in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2020.12.5)\n",
      "Requirement already satisfied: chardet in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (3.0.4)\n",
      "Requirement already satisfied: click in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (7.1.2)\n",
      "Requirement already satisfied: cycler in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.10.0)\n",
      "Requirement already satisfied: Cython in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.29.22)\n",
      "Requirement already satisfied: datasets in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.7.0)\n",
      "Requirement already satisfied: dill in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.3.3)\n",
      "Requirement already satisfied: einops in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.3.0)\n",
      "Requirement already satisfied: filelock in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (3.0.12)\n",
      "Requirement already satisfied: fire in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.3.1)\n",
      "Requirement already satisfied: flake8 in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (3.8.4)\n",
      "Requirement already satisfied: GPUtil in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.4.0)\n",
      "Requirement already satisfied: idna in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.10)\n",
      "Requirement already satisfied: isort in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (5.7.0)\n",
      "Requirement already satisfied: joblib in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.0.0)\n",
      "Requirement already satisfied: jsonlines in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.2.0)\n",
      "Requirement already satisfied: kiwisolver in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.3.1)\n",
      "Requirement already satisfied: lazy-object-proxy in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.4.3)\n",
      "Requirement already satisfied: matplotlib in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (3.3.3)\n",
      "Requirement already satisfied: mccabe in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.6.1)\n",
      "Requirement already satisfied: multiprocess in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.70.11.1)\n",
      "Requirement already satisfied: numpy in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.19.4)\n",
      "Requirement already satisfied: opencv-python in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (4.4.0.46)\n",
      "Requirement already satisfied: packaging in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (20.9)\n",
      "Requirement already satisfied: pandas in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.1.5)\n",
      "Requirement already satisfied: pathspec in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.8.1)\n",
      "Requirement already satisfied: pep8 in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.7.1)\n",
      "Requirement already satisfied: Pillow in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (8.2.0)\n",
      "Requirement already satisfied: pyaml in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (20.4.0)\n",
      "Requirement already satisfied: pyarrow in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.0.0)\n",
      "Requirement already satisfied: pycocotools in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.0.2)\n",
      "Requirement already satisfied: pycodestyle in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.6.0)\n",
      "Requirement already satisfied: pyflakes in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.2.0)\n",
      "Requirement already satisfied: pylint in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.6.0)\n",
      "Requirement already satisfied: pyparsing in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.4.7)\n",
      "Requirement already satisfied: python-dateutil in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.8.1)\n",
      "Requirement already satisfied: pytz in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2020.4)\n",
      "Requirement already satisfied: PyYAML in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (5.3.1)\n",
      "Requirement already satisfied: regex in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2020.11.13)\n",
      "Requirement already satisfied: requests in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.23.0)\n",
      "Requirement already satisfied: sacremoses in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.0.43)\n",
      "Requirement already satisfied: scipy in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.6.0)\n",
      "Requirement already satisfied: six in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.15.0)\n",
      "Requirement already satisfied: termcolor in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.1.0)\n",
      "Requirement already satisfied: timm in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.3.4)\n",
      "Requirement already satisfied: tokenizers in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.10.3)\n",
      "Requirement already satisfied: toml in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.10.2)\n",
      "Requirement already satisfied: torch in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.9.0)\n",
      "Requirement already satisfied: torchvision in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.10.0)\n",
      "Requirement already satisfied: tqdm in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (4.49.0)\n",
      "Requirement already satisfied: transformers in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (4.8.1)\n",
      "Requirement already satisfied: typed-ast in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.4.2)\n",
      "Requirement already satisfied: urllib3 in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.25.11)\n",
      "Requirement already satisfied: wget in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (3.2)\n",
      "Requirement already satisfied: wrapt in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (1.12.1)\n",
      "Requirement already satisfied: xxhash in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (2.0.0)\n",
      "Requirement already satisfied: scikit-image in /home/eltoto/mlnv/lib/python3.8/site-packages (from vltk==1.0.0) (0.18.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4 in /home/eltoto/mlnv/lib/python3.8/site-packages (from black->vltk==1.0.0) (3.7.4.3)\n",
      "Requirement already satisfied: mypy-extensions>=0.4.3 in /home/eltoto/mlnv/lib/python3.8/site-packages (from black->vltk==1.0.0) (0.4.3)\n",
      "Requirement already satisfied: huggingface-hub<0.1.0 in /home/eltoto/mlnv/lib/python3.8/site-packages (from datasets->vltk==1.0.0) (0.0.12)\n",
      "Requirement already satisfied: fsspec in /home/eltoto/mlnv/lib/python3.8/site-packages (from datasets->vltk==1.0.0) (2021.5.0)\n",
      "Requirement already satisfied: setuptools>=18.0 in /home/eltoto/mlnv/lib/python3.8/site-packages (from pycocotools->vltk==1.0.0) (50.3.2)\n",
      "Requirement already satisfied: networkx>=2.0 in /home/eltoto/mlnv/lib/python3.8/site-packages (from scikit-image->vltk==1.0.0) (2.5)\n",
      "Requirement already satisfied: imageio>=2.3.0 in /home/eltoto/mlnv/lib/python3.8/site-packages (from scikit-image->vltk==1.0.0) (2.9.0)\n",
      "Requirement already satisfied: PyWavelets>=1.1.1 in /home/eltoto/mlnv/lib/python3.8/site-packages (from scikit-image->vltk==1.0.0) (1.1.1)\n",
      "Requirement already satisfied: tifffile>=2019.7.26 in /home/eltoto/mlnv/lib/python3.8/site-packages (from scikit-image->vltk==1.0.0) (2021.3.17)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: decorator>=4.3.0 in /home/eltoto/mlnv/lib/python3.8/site-packages (from networkx>=2.0->scikit-image->vltk==1.0.0) (4.4.2)\n",
      "Installing collected packages: vltk\n",
      "  Attempting uninstall: vltk\n",
      "    Found existing installation: vltk 1.0.0\n",
      "    Uninstalling vltk-1.0.0:\n",
      "      Successfully uninstalled vltk-1.0.0\n",
      "  Running setup.py develop for vltk\n",
      "Successfully installed vltk-1.0.0\n",
      "\u001b[33mWARNING: You are using pip version 21.1.2; however, version 21.1.3 is available.\n",
      "You should consider upgrading via the '/home/eltoto/mlnv/bin/python -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vltk\n",
    "from vltk import Features, adapters\n",
    "from vltk.configs import DataConfig, VisionConfig, LangConfig\n",
    "from vltk.dataset import build\n",
    "from vltk.utils.adapters import rescale_box, clean_label, soft_score\n",
    "\n",
    "from pprint import pprint\n",
    "from PIL import Image\n",
    "import torch\n",
    "from collections import Counter, defaultdict\n",
    "\n",
    "DATADIR = '/home/eltoto/demodata/'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define FRCNN Adapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FRCNN(adapters.VisnExtraction):\n",
    "\n",
    "    # TODO: currently, this image preprocessing config is not correct\n",
    "    default_processor = VisionConfig(\n",
    "        **{\n",
    "            \"transforms\": [\"FromFile\", \"Resize\", \"ToTensor\", \"Normalize\"],\n",
    "            \"size\": 800,\n",
    "            \"max_size\": 1333,\n",
    "            \"mode\": Image.BILINEAR,\n",
    "            \"pad_value\": 0.0,\n",
    "            \"mean\": [0.404, 0.455,0.482],\n",
    "            \"std\": [1.0, 1.0, 1.0],\n",
    "        }\n",
    "    )\n",
    "\n",
    "    def setup():\n",
    "        from vltk import compat\n",
    "        from vltk.modeling.frcnn import FRCNN as FasterRCNN\n",
    "\n",
    "        weights = \"unc-nlp/frcnn-vg-finetuned\"\n",
    "        model_config = compat.Config.from_pretrained(\"unc-nlp/frcnn-vg-finetuned\")\n",
    "        return FasterRCNN.from_pretrained(weights, model_config), model_config\n",
    "\n",
    "    def schema(max_detections=36, visual_dim=2048):\n",
    "        return {\n",
    "            \"attr_ids\": Features.Ids(),\n",
    "            \"object_ids\": Features.Ids(),\n",
    "            vltk.features: Features.Features3D(max_detections, visual_dim),\n",
    "            vltk.box: Features.Box(),\n",
    "        }\n",
    "\n",
    "    def forward(model, entry):\n",
    "\n",
    "        size = entry[vltk.size]\n",
    "        scale_hw = entry[vltk.scale]\n",
    "        image = entry[vltk.img]\n",
    "\n",
    "        model_out = model(\n",
    "            images=image.unsqueeze(0),\n",
    "            image_shapes=size.unsqueeze(0),\n",
    "            scales_yx=scale_hw.unsqueeze(0),\n",
    "            padding=\"max_detections\",\n",
    "            pad_value=0.0,\n",
    "            location=\"cpu\",\n",
    "        )\n",
    "        normalized_boxes = torch.round(rescale_box(model_out[\"boxes\"][0], 1 / scale_hw))\n",
    "\n",
    "        return {\n",
    "            \"object_ids\": [model_out[\"obj_ids\"][0].tolist()],\n",
    "            \"attr_ids\": [model_out[\"attr_ids\"][0].tolist()],\n",
    "            vltk.box: [normalized_boxes.tolist()],\n",
    "            vltk.features: [model_out[\"roi_features\"][0]],\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Vision Dataset Adapters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### COCO and Visual Genome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Coco2014(adapters.VisnDataset):\n",
    "    def schema():\n",
    "        return {\n",
    "            vltk.box: Features.Box(),\n",
    "            vltk.polygons: Features.Polygons(),\n",
    "            vltk.objects: Features.StringList(),\n",
    "        }\n",
    "\n",
    "    def forward(json_files, splits):\n",
    "\n",
    "        total_annos = {}\n",
    "        id_to_cat = {}\n",
    "        file_to_id_to_stem = defaultdict(dict)\n",
    "        for file, json in json_files.items():\n",
    "            if \"instance\" not in file:\n",
    "                continue\n",
    "            info = json[\"images\"]\n",
    "            for i in info:\n",
    "                img_id = i[\"file_name\"].split(\".\")[0]\n",
    "                file_to_id_to_stem[file][i[\"id\"]] = img_id\n",
    "        for file, json in json_files.items():\n",
    "            if \"instance\" not in file:\n",
    "                continue\n",
    "\n",
    "            categories = json[\"categories\"]\n",
    "            for cat in categories:\n",
    "                id_to_cat[cat[\"id\"]] = cat[\"name\"]\n",
    "\n",
    "            for entry in json[\"annotations\"]:\n",
    "                # TODO: change this image ID thing later\n",
    "\n",
    "                img_id = str(file_to_id_to_stem[file][entry[\"image_id\"]])\n",
    "                bbox = entry[\"bbox\"]\n",
    "                segmentation = entry[\"segmentation\"]\n",
    "                category_id = id_to_cat[entry[\"category_id\"]]\n",
    "                if entry[\"iscrowd\"]:\n",
    "                    seg_mask = []\n",
    "                else:\n",
    "                    seg_mask = segmentation\n",
    "                    if not isinstance(seg_mask[0], list):\n",
    "                        seg_mask = [seg_mask]\n",
    "                img_data = total_annos.get(img_id, None)\n",
    "                if img_data is None:\n",
    "                    img_entry = defaultdict(list)\n",
    "                    img_entry[vltk.objects].append(category_id)\n",
    "                    img_entry[vltk.box].append(bbox)\n",
    "                    img_entry[vltk.polygons].append(seg_mask)\n",
    "                    total_annos[img_id] = img_entry\n",
    "                else:\n",
    "                    total_annos[img_id][vltk.box].append(bbox)\n",
    "                    total_annos[img_id][vltk.objects].append(category_id)\n",
    "                    total_annos[img_id][vltk.polygons].append(seg_mask)\n",
    "\n",
    "        return [{vltk.imgid: img_id, **entry} for img_id, entry in total_annos.items()]\n",
    "\n",
    "\n",
    "class VisualGenome(adapters.VisnDataset):\n",
    "    def schema():\n",
    "        return {}\n",
    "\n",
    "    def forward(json_files, splits):\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Adapters for Vision-Language Datasets "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VQA, GQA, and Visual Genome QA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GQA(adapters.VisnLangDataset):\n",
    "    data_info = {\n",
    "        \"dev\": {\"coco2014\": [\"test\"]},\n",
    "        \"train\": {\"visualgenome\": [\"train\"]},\n",
    "        \"val\": {\"visualgenome\": [\"train\"]},\n",
    "        \"test\": {\"coco2014\": [\"test\"]},\n",
    "        \"testdev\": {\"coco2014\": [\"val\"]},\n",
    "    }\n",
    "\n",
    "    filters = [\"unbalanced\", \"train\"]\n",
    "\n",
    "    def schema():\n",
    "        return {vltk.label: Features.StringList(), \"layout\": Features.StringList()}\n",
    "\n",
    "    def forward(json_files, split, min_label_frequency=2):\n",
    "        label_frequencies = Counter()\n",
    "        batch_entries = []\n",
    "\n",
    "        for filename, data in json_files.items():\n",
    "            for i, (k, v) in enumerate(data.items()):\n",
    "                if \"answer\" in v:\n",
    "                    answer = clean_label(v[\"answer\"])\n",
    "                    label_frequencies.update([answer])\n",
    "\n",
    "            for i, (k, v) in enumerate(data.items()):\n",
    "                if split == \"test\":\n",
    "                    answer = None\n",
    "                    layout = None\n",
    "                elif label_frequencies[v[\"answer\"]] < min_label_frequency:\n",
    "                    continue\n",
    "                else:\n",
    "                    answer = clean_label(v[\"answer\"])\n",
    "                    layout = [layout[\"operation\"] for layout in v[\"semantic\"]]\n",
    "\n",
    "                text = v[\"question\"]\n",
    "                img_id = v[\"imageId\"].lstrip(\"n\")\n",
    "\n",
    "                entry = {\n",
    "                    vltk.text: text,\n",
    "                    vltk.imgid: img_id,\n",
    "                    vltk.label: [answer],\n",
    "                    \"layout\": layout,\n",
    "                }\n",
    "\n",
    "                batch_entries.append(entry)\n",
    "\n",
    "        return batch_entries\n",
    "    \n",
    "class VQA(adapters.VisnLangDataset):\n",
    "    data_info = {\n",
    "        \"val\": {\"coco2014\": [\"val\"]},\n",
    "        \"train\": {\"coco2014\": [\"train\"]},\n",
    "        \"test\": {\"coco2014\": [\"test\"]},\n",
    "    }\n",
    "\n",
    "    def schema():\n",
    "        return {\n",
    "            vltk.qid: Features.String(),\n",
    "            vltk.label: Features.StringList(),\n",
    "            vltk.score: Features.FloatList(),\n",
    "        }\n",
    "\n",
    "    def adjust_imgid(imgid, vdset_name, vdset_split):\n",
    "        imgid = f'{\"COCO\"}_{vdset_split[0].lower()}{2014}_{\"\".join([\"0\"] * (12 - len(imgid)))}{imgid}'\n",
    "        return imgid\n",
    "\n",
    "    def forward(json_files, split, min_label_frequency=9):\n",
    "        batch_entries = []\n",
    "        all_questions = []\n",
    "        qid2answers = {}\n",
    "        label_frequencies = Counter()\n",
    "        for filename, x in json_files.items():\n",
    "            if \"questions\" in x:\n",
    "                all_questions.extend(x[\"questions\"])\n",
    "            else:\n",
    "                annotations = x[\"annotations\"]\n",
    "                accepted_answers = {\n",
    "                    clean_label(anno[\"multiple_choice_answer\"]) for anno in annotations\n",
    "                }\n",
    "                for anno in annotations:\n",
    "                    qid = str(anno[\"question_id\"])\n",
    "                    answers = anno[\"answers\"]\n",
    "                    label_frequencies.update(\n",
    "                        [clean_label(anno[\"multiple_choice_answer\"])]\n",
    "                    )\n",
    "                    answer_counter = Counter()\n",
    "                    for ans_dict in answers:\n",
    "                        ans = ans_dict[\"answer\"]\n",
    "                        if ans not in accepted_answers:\n",
    "                            pass\n",
    "                        else:\n",
    "                            ans = clean_label(ans)\n",
    "                            answer_counter.update([ans])\n",
    "                    qid2answers[qid] = {\n",
    "                        k: soft_score(v) for k, v in answer_counter.items()\n",
    "                    }\n",
    "\n",
    "        for entry in all_questions:\n",
    "            try:\n",
    "                entry[vltk.imgid] = str(entry.pop(\"image_id\"))\n",
    "            except Exception:\n",
    "                raise Exception(entry.keys())\n",
    "            entry[vltk.text] = entry.pop(\"question\")\n",
    "            # entry.pop(\"question_id\")\n",
    "            entry[\"qid\"] = str(entry.pop(\"question_id\"))\n",
    "            try:\n",
    "                entry[vltk.label] = qid2answers[entry[\"qid\"]]\n",
    "                labels = {\n",
    "                    l: s\n",
    "                    for l, s in entry[vltk.label].items()\n",
    "                    if label_frequencies[l] > min_label_frequency\n",
    "                }\n",
    "                if not labels:\n",
    "                    continue\n",
    "\n",
    "                labels, scores = adapters.VisnLangDataset._label_handler(labels)\n",
    "                entry[vltk.score] = scores\n",
    "                entry[vltk.label] = labels\n",
    "            except KeyError:\n",
    "                pass\n",
    "\n",
    "            batch_entries.append(entry)\n",
    "        return batch_entries\n",
    "    \n",
    "class COCOCaptions(adapters.VisnLangDataset):\n",
    "    data_info = {\n",
    "        \"train\": {\"coco2014\": [\"train\"]},\n",
    "        \"val\": {\"coco2014\": [\"val\"]},\n",
    "    }\n",
    "\n",
    "    def schema():\n",
    "        return {}\n",
    "\n",
    "    def forward(json_files, split, min_label_frequency=2):\n",
    "        batch_entries = []\n",
    "        id2imgid = {}\n",
    "        for filename, data in json_files.items():\n",
    "            if \"annotations\" not in data:\n",
    "                continue\n",
    "            if \"caption\" not in data[\"annotations\"][0]:\n",
    "                continue\n",
    "            for img in data[\"images\"]:\n",
    "                id2imgid[img[\"id\"]] = img[\"file_name\"]\n",
    "            for item in (data[\"annotations\"]):\n",
    "                imgid = id2imgid[item[\"image_id\"]].split(\".\")[0]\n",
    "                entry = {vltk.imgid: imgid, vltk.text: item[\"caption\"]}\n",
    "\n",
    "                batch_entries.append(entry)\n",
    "\n",
    "        return batch_entries\n",
    "    \n",
    "    \n",
    "class VGQA(adapters.VisnLangDataset):\n",
    "    data_info = {\n",
    "        \"train\": {\"visualgenome\": [\"train\"]},\n",
    "    }\n",
    "\n",
    "    def schema():\n",
    "        return {\n",
    "            vltk.qid: Features.String(),\n",
    "            vltk.label: Features.StringList(),\n",
    "        }\n",
    "\n",
    "    def forward(json_files, split, min_label_frequency=9):\n",
    "        batch_entries = []\n",
    "        answer_counts = Counter()\n",
    "        for filename, data in json_files.items():\n",
    "            for y in data:\n",
    "                for x in y[\"qas\"]:\n",
    "                    answer_counts.update([clean_label(x[\"answer\"])]),\n",
    "\n",
    "        for filename, data in json_files.items():\n",
    "            for y in data:\n",
    "                for x in y[\"qas\"]:\n",
    "                    if answer_counts[clean_label(x[\"answer\"])] >= min_label_frequency:\n",
    "                        entry = {\n",
    "                            vltk.qid: str(x[\"qa_id\"]),\n",
    "                            vltk.imgid: str(x[\"image_id\"]),\n",
    "                            vltk.text: x[\"question\"],\n",
    "                            vltk.label: [clean_label(x[\"answer\"])],\n",
    "                        }\n",
    "                        batch_entries.append(entry)\n",
    "        return batch_entries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract Datasets for Each Defined Adapter Class "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "will write to /home/eltoto/demodata/coco2014/frcnn\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/eltoto/.local/lib/python3.8/site-packages/torchvision/transforms/transforms.py:257: UserWarning: Argument interpolation should be of type InterpolationMode instead of int. Please, use InterpolationMode enum.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading weights file https://cdn.huggingface.co/unc-nlp/frcnn-vg-finetuned/pytorch_model.bin from cache at /home/eltoto/.cache/torch/transformers/57f6df6abe353be2773f2700159c65615babf39ab5b48114d2b49267672ae10f.77b59256a4cf8343ae0f923246a81489fc8d82f98d082edc2d2037c977c0d9d0\n",
      "All model checkpoint weights were used when initializing FRCNN.\n",
      "\n",
      "All the weights of FRCNN were initialized from the model checkpoint at unc-nlp/frcnn-vg-finetuned.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use FRCNN for predictions without further training.\n",
      "extracting from ['/home/eltoto/demodata/coco2014/test', '/home/eltoto/demodata/coco2014/train', '/home/eltoto/demodata/coco2014/val']\n",
      "100%|██████████| 10/10 [01:03<00:00,  6.39s/it]\n",
      "saving...\n",
      "Success! You wrote 8 entry(s) and 2 mb\n",
      "Located: /home/eltoto/demodata/coco2014/frcnn/val.arrow\n",
      "Success! You wrote 1 entry(s) and 0 mb\n",
      "Located: /home/eltoto/demodata/coco2014/frcnn/test.arrow\n",
      "Success! You wrote 1 entry(s) and 0 mb\n",
      "Located: /home/eltoto/demodata/coco2014/frcnn/train.arrow\n",
      "will write to /home/eltoto/demodata/visualgenome/frcnn\n",
      "loading weights file https://cdn.huggingface.co/unc-nlp/frcnn-vg-finetuned/pytorch_model.bin from cache at /home/eltoto/.cache/torch/transformers/57f6df6abe353be2773f2700159c65615babf39ab5b48114d2b49267672ae10f.77b59256a4cf8343ae0f923246a81489fc8d82f98d082edc2d2037c977c0d9d0\n",
      "All model checkpoint weights were used when initializing FRCNN.\n",
      "\n",
      "All the weights of FRCNN were initialized from the model checkpoint at unc-nlp/frcnn-vg-finetuned.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use FRCNN for predictions without further training.\n",
      "extracting from ['/home/eltoto/demodata/visualgenome/train']\n",
      "100%|██████████| 11/11 [01:04<00:00,  5.87s/it]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/6 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "saving...\n",
      "Success! You wrote 11 entry(s) and 3 mb\n",
      "Located: /home/eltoto/demodata/visualgenome/frcnn/train.arrow\n",
      "loading annotations...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6/6 [00:19<00:00,  3.18s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing to Datasets/Arrow object\n",
      "saving...\n",
      "Success! You wrote 122218 entry(s) and 203 mb\n",
      "Located: /home/eltoto/demodata/coco2014/annotations.arrow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "searching for input files for splits: {'test', 'train', 'val', 'validation', 'evaluation', 'eval', 'dev'}\n",
      "loading json files from: ['/home/eltoto/demodata/vqa/v2_mscoco_train2014_annotations.json', '/home/eltoto/demodata/vqa/v2_OpenEnded_mscoco_train2014_questions.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:04<00:00,  2.16s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n",
      "Success! You wrote 432628 entry(s) and 46 mb\n",
      "Located: /home/eltoto/demodata/vqa/train.arrow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading json files from: ['/home/eltoto/demodata/vqa/v2_mscoco_val2014_annotations.json', '/home/eltoto/demodata/vqa/v2_OpenEnded_mscoco_val2014_questions.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:01<00:00,  1.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n",
      "Success! You wrote 206454 entry(s) and 21 mb\n",
      "Located: /home/eltoto/demodata/vqa/val.arrow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 12.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "searching for input files for splits: {'test', 'train', 'val', 'validation', 'evaluation', 'eval', 'dev'}\n",
      "loading json files from: ['/home/eltoto/demodata/gqa/test_balanced_questions.json']\n",
      "begin extraction\n",
      "writing rows to arrow dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success! You wrote 95336 entry(s) and 5 mb\n",
      "Located: /home/eltoto/demodata/gqa/test.arrow\n",
      "loading json files from: ['/home/eltoto/demodata/gqa/val_balanced_questions.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:03<00:00,  3.26s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success! You wrote 129563 entry(s) and 13 mb\n",
      "Located: /home/eltoto/demodata/gqa/val.arrow\n",
      "loading json files from: ['/home/eltoto/demodata/gqa/testdev_balanced_questions.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  4.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success! You wrote 11025 entry(s) and 1 mb\n",
      "Located: /home/eltoto/demodata/gqa/dev.arrow\n",
      "searching for input files for splits: {'test', 'train', 'val', 'validation', 'evaluation', 'eval', 'dev'}\n",
      "No files pattern matched with corresponding split, falling back to searching all json in top level directory\n",
      "loading json files from: ['/home/eltoto/demodata/vgqa/question_answers.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:06<00:00,  6.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n",
      "Success! You wrote 1185566 entry(s) and 78 mb\n",
      "Located: /home/eltoto/demodata/vgqa/train.arrow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/3 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "searching for input files for splits: {'test', 'train', 'val', 'validation', 'evaluation', 'eval', 'dev'}\n",
      "loading json files from: ['/home/eltoto/demodata/cococaptions/instances_train2014.json', '/home/eltoto/demodata/cococaptions/captions_train2014.json', '/home/eltoto/demodata/cococaptions/person_keypoints_train2014.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:12<00:00,  4.17s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n",
      "Success! You wrote 414113 entry(s) and 34 mb\n",
      "Located: /home/eltoto/demodata/cococaptions/train.arrow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/3 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading json files from: ['/home/eltoto/demodata/cococaptions/instances_val2014.json', '/home/eltoto/demodata/cococaptions/person_keypoints_val2014.json', '/home/eltoto/demodata/cococaptions/captions_val2014.json']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:05<00:00,  1.81s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin extraction\n",
      "writing rows to arrow dataset\n",
      "Success! You wrote 202654 entry(s) and 16 mb\n",
      "Located: /home/eltoto/demodata/cococaptions/val.arrow\n"
     ]
    }
   ],
   "source": [
    "cocofeats = FRCNN.extract(DATADIR, dataset=\"coco2014\")\n",
    "vgfeats = FRCNN.extract(DATADIR, dataset=\"visualgenome\")\n",
    "coco2014 = Coco2014.extract(DATADIR)\n",
    "vqa = VQA.extract(DATADIR)\n",
    "gqa = GQA.extract(DATADIR)\n",
    "vgqa = VGQA.extract(DATADIR)\n",
    "cococaptions = COCOCaptions.extract(DATADIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Register User-Defined Adapters and Define Config with VLTK to Superset Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added VisnLangDataset cococaptions: train\n",
      "Added VisnDataset coco2014: train\n",
      "Added VisnLangDataset cococaptions: val\n",
      "Added VisnDataset coco2014: val\n",
      "Added VisnLangDataset gqa: train\n",
      "Warning: No Annotations for visualgenome\n",
      "Added VisnDataset visualgenome: train\n",
      "Added VisnLangDataset gqa: val\n",
      "Added VisnLangDataset vgqa: train\n",
      "Added VisnLangDataset vqa: train\n",
      "Added VisnLangDataset vqa: val\n",
      "Max spanning column names for each batch: {'label', 'attr_ids', 'object_ids', 'imgid', 'layout', 'box', 'text', 'objects', 'features', 'qid', 'score'} (not including extra columns/features from processors)\n",
      "resizing datasets to account for 228610 missing image IDs\n"
     ]
    }
   ],
   "source": [
    "# add adapters to library\n",
    "adapters.Adapters().add(VQA, GQA, Coco2014, VisualGenome, FRCNN, COCOCaptions, VGQA)\n",
    "\n",
    "#config\n",
    "config = DataConfig(\n",
    "    # choose which dataset and dataset split for train and eval\n",
    "    train_datasets=[[\"vqa\", \"trainval\"],[\"cococaptions\", 'trainval'], ['vgqa', 'train'], ['gqa', 'trainval']],\n",
    "    # choose which feature extractor to use\n",
    "    extractor=\"frcnn\",\n",
    "    datadir=DATADIR,\n",
    "    train_batch_size=1,\n",
    "    # iterate with through datasets via images first versus text\n",
    "    img_first=True,\n",
    "    #ignore segmentation annotations from being prcoessed with the COCO dataset\n",
    "    ignore_segmentation=True\n",
    ")\n",
    "train, val = build(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# First Batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['layout', 'text_attention_mask', 'input_ids', 'type_ids', 'label', 'score', 'attr_ids', 'box', 'features', 'imgid', 'object_ids'])\n",
      "\n",
      "{'attr_ids': tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 7., 0., 7., 0., 0., 0., 0.,\n",
      "         0., 0., 0., 0., 0., 7., 0., 0., 7., 7., 0., 0., 7., 0., 7., 0., 0., 0.]]),\n",
      " 'box': tensor([[[ 32.,  40., 800., 634.],\n",
      "         [ 53.,   1., 800., 461.],\n",
      "         [  2., 211., 800., 781.],\n",
      "         [100.,   2., 800., 315.],\n",
      "         [185., 108., 800., 683.],\n",
      "         [  1.,   5., 526., 603.],\n",
      "         [  3., 118., 514., 725.],\n",
      "         [  2.,   4., 573., 388.],\n",
      "         [  1.,  14., 402., 684.],\n",
      "         [ 68., 314., 800., 846.],\n",
      "         [  2.,   5., 460., 465.],\n",
      "         [  5.,   3., 562., 263.],\n",
      "         [281.,  50., 800., 621.],\n",
      "         [  4.,   4., 447., 311.],\n",
      "         [280., 172., 800., 742.],\n",
      "         [415., 132., 800., 716.],\n",
      "         [435.,  28., 800., 629.],\n",
      "         [310.,   8., 800., 490.],\n",
      "         [219., 281., 800., 800.],\n",
      "         [305.,   6., 800., 323.],\n",
      "         [507.,  69., 800., 711.],\n",
      "         [ 84., 117., 535., 844.],\n",
      "         [  4., 245., 525., 808.],\n",
      "         [443.,   6., 800., 307.],\n",
      "         [443.,   9., 800., 444.],\n",
      "         [  3.,  74., 296., 709.],\n",
      "         [  4.,   6., 322., 403.],\n",
      "         [515.,   5., 800., 265.],\n",
      "         [  3., 141., 345., 821.],\n",
      "         [ 71., 418., 800., 943.],\n",
      "         [  3.,   8., 306., 570.],\n",
      "         [512.,  12., 800., 553.],\n",
      "         [555.,   9., 800., 349.],\n",
      "         [576.,  17., 800., 674.],\n",
      "         [425., 250., 800., 808.],\n",
      "         [  4., 342., 462., 844.]]]),\n",
      " 'features': tensor([[[2.0876e-01, 0.0000e+00, 3.4234e-04,  ..., 0.0000e+00,\n",
      "          1.3988e+00, 0.0000e+00],\n",
      "         [2.0218e-01, 0.0000e+00, 3.1476e-04,  ..., 0.0000e+00,\n",
      "          1.4383e+00, 0.0000e+00],\n",
      "         [8.1444e-02, 0.0000e+00, 1.2307e-04,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 3.7016e-03],\n",
      "         ...,\n",
      "         [8.9323e-03, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 0.0000e+00],\n",
      "         [3.1045e-03, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 7.8256e-03],\n",
      "         [3.4366e-03, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
      "          7.4322e-02, 0.0000e+00]]]),\n",
      " 'imgid': ['1009'],\n",
      " 'input_ids': tensor([[[ 101,  100, 2003, 2006, 1996, 2795, 1029,  102,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2785, 1997, 7390, 2003, 1996, 2338, 2006, 1029,  102,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2029, 2217, 1997, 1996, 3746, 2003, 1996, 3242, 1029,\n",
      "           102,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2003, 1996, 3538, 1997, 7390, 2006, 1996, 2723, 1029,\n",
      "           102,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2785, 1997, 7390, 2003, 2006, 1996, 2723, 1029,  102,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2785, 1997, 7390, 2079, 2017, 2228, 2003, 2006, 1996,\n",
      "          2723, 1029,  102,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2029, 2112, 1997, 1996, 3861, 2003, 1996, 2330, 2338,\n",
      "          1010, 1996, 2327, 2030, 1996, 3953, 1029,  102,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 1996, 2338, 2000, 1996, 2157, 1997, 1037, 3242, 1029,\n",
      "           102,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 2045, 1037, 4624, 2030, 1037, 3242, 1999, 1996, 3861,\n",
      "          1029,  102,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0],\n",
      "         [ 101,  100, 1996, 3242, 2006, 1996, 2157, 2217, 1997, 1996, 6302,\n",
      "          1029,  102,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "             0,    0,    0]]]),\n",
      " 'label': tensor([[ 137, 1303, 1060,  259,  259,  259,  146,  838, 1488, 1488]]),\n",
      " 'layout': [[tensor([95., 90., 89.]),\n",
      "             tensor([95., 90., 89.]),\n",
      "             tensor([95., 89.]),\n",
      "             tensor([95., 90., 89.]),\n",
      "             tensor([95., 90., 89.]),\n",
      "             tensor([95., 90., 89.]),\n",
      "             tensor([95., 50., 40.]),\n",
      "             tensor([ 95., 120.]),\n",
      "             tensor([95., 49., 95., 49., 88.]),\n",
      "             tensor([ 95., 110.])]],\n",
      " 'object_ids': tensor([[72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72.,\n",
      "         72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72., 72.,\n",
      "         72., 72., 72., 72., 72., 72., 72., 72.]]),\n",
      " 'score': tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]),\n",
      " 'text_attention_mask': tensor([[[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]]),\n",
      " 'type_ids': tensor([[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
      "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]])}\n"
     ]
    }
   ],
   "source": [
    "for batch in train:\n",
    "    pprint(batch.keys())\n",
    "    print()\n",
    "    pprint(batch)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}