{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d351030",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eaa2cd3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "DATA_DIR = \"/mnt/data-ssd-1/data/supreme_corpus/processed_data/\"\n",
    "\n",
    "with open(DATA_DIR + \"segment_meta_post_norm.json\") as f:\n",
    "    segment_metas = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f4a1257",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e942270",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "14daaf91",
   "metadata": {},
   "source": [
    "## rough align"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ceeb50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# token align two strings that are very similar such as normalized and unnormalized tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "79a06a1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_list = [\n",
    "    \" \".join([ee[\"value\"] for ee in e[\"transcript\"][\"tokens\"] if ee[\"type\"] == \"text\"]) \n",
    "    for e in segment_metas[:100]\n",
    "]\n",
    "text_norm_list = [e[\"transcript\"][\"text_norm\"] for e in segment_metas[:100]]\n",
    "\n",
    "text = \" \".join(text_list)\n",
    "text_norm = \" \".join(text_norm_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69a7c8e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Thank you, Mr. Chief Justice, and may it please the Court: Justice Jackson long ago warned against giving the convict population of the country new and unprecedented opportunities to litigate until they serve their sentences or make the best of increased opportunities to escape. The Sixth Circuit here blessed precisely the sort of opportunity he warned of. It held that when a federal statute prohibits ordering a prisoner's transportation with a writ of habeas corpus, courts may instead order transportation under the All Writs Act. But courts have no such power. Every All Writs order must be agreeable to the usages and principles of law, meaning the traditional writs as altered by statute. Transportation orders must be agreeable to habeas law because habeas writs were the only traditional writs used for ordering the transportation of prisoners. So, when a federal habeas statute prohibits ordering transportation with a writ of habeas corpus in a particular situation, courts may not evade\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text[:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5059b2c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"thank you mr chief justice and may it please the court justice jackson long ago warned against giving the convict population of the country new and unprecedented opportunities to litigate until they serve their sentences or make the best of increased opportunities to escape the sixth circuit here blessed precisely the sort of opportunity he warned of it held that when a federal statute prohibits ordering a prisoner's transportation with a writ of habeas corpus courts may instead order transportation under the all writs act but courts have no such power every all writs order must be agreeable to the usages and principles of law meaning the traditional writs as altered by statute transportation orders must be agreeable to habeas law because habeas writs were the only traditional writs used for ordering the transportation of prisoners so when a federal habeas statute prohibits ordering transportation with a writ of habeas corpus in a particular situation courts may not evade that prohibit\""
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_norm[:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "80217afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = text.lower().split()\n",
    "tokens_normalized = text_norm.lower().split()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "75d47d96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5161\n",
      "5258\n"
     ]
    }
   ],
   "source": [
    "print(len(tokens))\n",
    "print(len(tokens_normalized))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "de42b142",
   "metadata": {},
   "outputs": [],
   "source": [
    "def align_norm_tokens(tokens, tokens_norm, n_idx_check=20)\n",
    "    n2 = 0\n",
    "    idx_pairs = []\n",
    "    for n1 in range(len(tokens)):\n",
    "    #     expected_n2 = int(round(n1 / len(tokens) * len(sub_tokens)))\n",
    "        t1 = tokens[n1]\n",
    "        nt = 0\n",
    "        for t2 in tokens_norm[n2:n2+n_idx_check]:\n",
    "            if t2 in t1:\n",
    "                idx_pairs.append((n1, n2))\n",
    "                n2 += nt + 1\n",
    "                break\n",
    "            nt += 1\n",
    "    #     # check if we didn't find a match\n",
    "        if nt == n_idx_check and len(idx_pairs) > 0 and n1 - idx_pairs[-1][0] > n_idx_check:\n",
    "            import pdb; pdb.set_trace()\n",
    "    if (\n",
    "        len(idx_pairs) == 0 or \n",
    "        idx_pairs[-1][0] <= len(tokens) - n_idx_check or \n",
    "        idx_pairs[-1][1] <= len(tokens_norm) - n_idx_check\n",
    "    ):\n",
    "        raise ValueError(\"alignment did not work\")\n",
    "    return idx_pairs\n",
    "\n",
    "idx_pairs = align_tokens(tokens, tokens_normalized)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "633a1563",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(5158, 5255), (5159, 5256), (5160, 5257)]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx_pairs[-3:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b81202",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e0a3bcd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90858e6b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd4572f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cda1ad9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}