notebook_templates/Table 1 of NMI paper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from pathlib import Path\n",
    "import os\n",
    "import sys\n",
    "import mne\n",
    "import torch\n",
    "import numpy as np\n",
    "import bm\n",
    "from bm import play\n",
    "from bm.train import main\n",
    "from bm.events import Word\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython import display as disp\n",
    "\n",
    "mne.set_log_level(False)\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "os.chdir(main.dora.dir.parent)\n",
    "os.environ['NO_DOWNLOAD'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigs = ['34219380', '6e3bf7d7', '557f5f8a', '4395629c']\n",
    "# sigs = ['889cf0fa']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:bm.play:Loading solver from XP 34219380. Overrides used: ['model=clip_conv', 'dset.selections=[\"audio_mous\"]', 'seed=2036', 'dset.force_uid_assignement=true']\n",
      "WARNING:bm.dataset:Requested 1000 recordings but only found 96\n",
      "INFO:bm.dataset:Loading Subjects | 19/96 | 0.85 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 38/96 | 1.00 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 57/96 | 1.08 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 76/96 | 1.12 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 95/96 | 1.10 it/sec\n",
      "INFO:bm.train:Model hash: 3502acedd4c0aad6ce5666c554cf3c70065bec93\n",
      "INFO:bm.play:Loading solver from XP 6e3bf7d7. Overrides used: ['model=clip_conv', 'dset.selections=[\"gwilliams2022\"]', 'seed=2036', 'optim.lr=0.0003', 'optim.batch_size=128']\n",
      "WARNING:bm.dataset:Requested 1000 recordings but only found 196\n",
      "INFO:bm.dataset:Loading Subjects | 39/196 | 7.01 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 78/196 | 7.13 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 117/196 | 7.20 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 156/196 | 7.19 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 195/196 | 7.23 it/sec\n",
      "INFO:bm.train:Model hash: 2e99ee0c09517dd1bcbe22e4ffb18feef688b7ef\n",
      "INFO:bm.play:Loading solver from XP 557f5f8a. Overrides used: ['model=clip_conv', 'dset.selections=[\"broderick2019\"]', 'seed=2036', 'test.wer_recordings=100']\n",
      "WARNING:bm.dataset:Requested 1000 recordings but only found 380\n",
      "INFO:bm.dataset:Loading Subjects | 76/380 | 13.21 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 152/380 | 15.13 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 228/380 | 15.93 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 304/380 | 16.37 it/sec\n",
      "INFO:bm.train:Model hash: e2f2a287e9031960614d2ebb4057e95c0888495b\n",
      "INFO:bm.play:Loading solver from XP 4395629c. Overrides used: ['model=clip_conv', 'dset.selections=[\"brennan2019\"]', 'seed=2036']\n",
      "WARNING:bm.dataset:Requested 1000 recordings but only found 33\n",
      "INFO:bm.dataset:Loading Subjects | 6/33 | 10.73 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 12/33 | 10.53 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 18/33 | 10.67 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 24/33 | 10.85 it/sec\n",
      "INFO:bm.dataset:Loading Subjects | 30/33 | 10.88 it/sec\n",
      "INFO:bm.train:Model hash: 0b43108cca9ff1485aa1dd1116aec35ee9c82ea4\n",
      "ALL SOLVERS LOADED\n",
      "now the table.\n",
      "audio_mous 273 & 96 & 80.9 h& 5497 & 1754& 1270 & 745\\\\\n",
      "Vocab overlap: 85.2%\n",
      "gwilliams2022 208 & 27 & 56.2 h& 3266 & 1810& 1284 & 846\\\\\n",
      "Vocab overlap: 63.9%\n",
      "broderick2019 128 & 19 & 19.2 h& 1189 & 1418& 1054 & 764\\\\\n",
      "Vocab overlap: 67.3%\n",
      "brennan2019 60 & 33 & 6.7 h& 1211 & 513& 190 & 148\\\\\n",
      "Vocab overlap: 59.5%\n"
     ]
    }
   ],
   "source": [
    "def _get_segments_and_vocabs(solver):\n",
    "    from scripts.run_eval_probs import _get_extra_info\n",
    "    per_split = {}\n",
    "    for split in ['train', 'test']:\n",
    "        segments = set()\n",
    "        sentences = set()\n",
    "        vocab = set()\n",
    "        dset = getattr(solver.datasets, split)\n",
    "        loader = solver.make_loader(dset, shuffle=False)\n",
    "        for idx, batch in enumerate(loader):\n",
    "            data, *_ = _get_extra_info(batch, solver.args.dset.sample_rate)\n",
    "            time_to_main_word = 0 - solver.args.dset.tmin  # location of main word relative to segment start\n",
    "            # e.g. with MNE we have tmin=-0.5 so the main word is 0.5 seconds after start of MNE Epoch.\n",
    "            margin = 2 # we need to look a bit after 0.5 due to rounding error, this is in time steps.\n",
    "            look_at_index = int(time_to_main_word * solver.args.dset.sample_rate + margin)\n",
    "            word_index = data[:, 0, look_at_index]\n",
    "            sequence_id = data[:, 1, look_at_index]\n",
    "            segment_ids = list(zip(word_index.tolist(), sequence_id.tolist()))\n",
    "            \n",
    "            segment_duration = data.shape[-1] / solver.args.dset.sample_rate\n",
    "            for events in batch._event_lists:\n",
    "                for event in events:\n",
    "                    if isinstance(event, Word):\n",
    "                        start = event.start - events[0].start\n",
    "                        end = start + event.duration\n",
    "                        if end > 0.02 and start < segment_duration - 0.02:\n",
    "                            # due to rounding errors, retrieval of related events\n",
    "                            # can sometime overlap in a non meaningful way, e.g. less than 20ms.\n",
    "                            # we only consider an event if it overlaps for at least 20ms.\n",
    "                            sentences.add(event.word_sequence)\n",
    "                            vocab.add(event.word)\n",
    "            segments |= set(segment_ids)\n",
    "#             print(idx, len(loader), end='\\r')\n",
    "#         print(split, \"done\", \" \" * 400)\n",
    "        per_split[split] = (segments, vocab, sentences)\n",
    "    return per_split\n",
    "\n",
    "\n",
    "def print_table_line(solver):\n",
    "    channels = solver.datasets.train[0].meg.shape[0]\n",
    "    n_subjects = len(set([dataset.recording.subject_uid for dataset in solver.datasets.train.datasets]))\n",
    "    per_split = _get_segments_and_vocabs(solver)\n",
    "    assert len(solver.args.dset.selections) == 1\n",
    "    name = solver.args.dset.selections[0]\n",
    "    duration = 0.\n",
    "    for dset in solver.datasets.train.datasets:\n",
    "        events = dset.recording.events()\n",
    "        duration += (events.start + events.duration).max()\n",
    "    \n",
    "    print(name, channels, '&' , n_subjects, '&', format(duration/ 3600, '.1f') + ' h', end='')\n",
    "    for split in ('train', 'test'):\n",
    "        segments, vocab, sentences = per_split[split]\n",
    "        print('&', len(segments), '&', len(vocab), end='')\n",
    "    vocab_train = per_split['train'][1]\n",
    "    vocab_test = per_split['test'][1]\n",
    "    vocab_overlap = len(vocab_train & vocab_test) / len(vocab_test)\n",
    "#     print('&', format(vocab_overlap, '.1%'), end='')\n",
    "    print(r'\\\\')\n",
    "    print(\"Vocab overlap:\", format(vocab_overlap, '.1%'))\n",
    "    \n",
    "solvers = [play.get_solver_from_sig(sig) for sig in sigs]\n",
    "print(\"ALL SOLVERS LOADED\")\n",
    "print(\"now the table.\")\n",
    "for solver in solvers:\n",
    "    print_table_line(solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attention_map(solver):\n",
    "    loader = solver.make_loader(solver.datasets.train)\n",
    "    batch = next(iter(loader)).to(solver.device)\n",
    "    model = solver.model\n",
    "    merger = model.merger\n",
    "    positions = merger.position_getter.get_positions(batch)\n",
    "    embedding = merger.embedding(positions)\n",
    "    meg = batch.meg\n",
    "    B, C, T = meg.shape\n",
    "    score_offset = torch.zeros(B, C, device=meg.device)\n",
    "    score_offset[merger.position_getter.is_invalid(positions)] = float('-inf')\n",
    "    heads = merger.heads[None].expand(B, -1, -1)\n",
    "    scores = torch.einsum(\"bcd,bod->boc\", embedding, heads)\n",
    "    scores += score_offset[:, None]\n",
    "    weights = torch.softmax(scores, dim=2)\n",
    "    \n",
    "    # Weights is of shape [Virtual Channels, Input Channels]\n",
    "    # Each Virtual Channel is a weighted sum over the input channels.\n",
    "    # Positions give the normalized 2d position for each Input channel.\n",
    "    # To get an overall weight for a given input sensor you can for instance do\n",
    "    # weights[0].sum(dim=0)\n",
    "    return weights[0], positions[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights.shape, positions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}