spanbert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2ac33e44-4e7f-4cbd-9e4f-a710c78aa594",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normal stuff\n",
    "import pandas as pd\n",
    "from config import OUTPUT_DIR\n",
    "from os.path import join as pjoin\n",
    "# transformers\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
    "from transformers import pipeline\n",
    "\n",
    "# importing the models\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"abhibisht89/spanbert-large-cased-finetuned-ade_corpus_v2\")\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"abhibisht89/spanbert-large-cased-finetuned-ade_corpus_v2\")\n",
    "ner = pipeline(\"ner\",model=model,tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e5d06cb9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exception: list indices must be integers or slices, not str\n",
      "> \u001b[0;32m/home/james/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m(97)\u001b[0;36mget_tokens_from_ner_specific\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m     96 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m---> 97 \u001b[0;31m    \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m     98 \u001b[0;31m        \u001b[0mall_stacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n",
      "\u001b[1;32m     83 \u001b[0m\u001b[0;32mdef\u001b[0m \u001b[0mget_tokens_from_ner_specific\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentity_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     84 \u001b[0m    \u001b[0;34m'''gets proper tokens from nlp pipeline (merges subwords as well)'''\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     85 \u001b[0m    \u001b[0;31m# output is a list of {entity : str, score : float, index : int, word : str, start : int, end : int}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     86 \u001b[0m    \u001b[0;31m# goal: merge consecutive indices into one word\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     87 \u001b[0m    outputs = filter(lambda output: output['entity'].endswith(\n",
      "\u001b[1;32m     88 \u001b[0m        entity_type), raw_outputs)\n",
      "\u001b[1;32m     89 \u001b[0m    \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     90 \u001b[0m        word_stack = reduce(word_accumulator_fn, outputs,\n",
      "\u001b[1;32m     91 \u001b[0m                            dict(word_list=[], curr_stack=[]))\n",
      "\u001b[1;32m     92 \u001b[0m        \u001b[0mall_stacks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'word_list'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     93 \u001b[0m    \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     94 \u001b[0m        \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"exception:\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     95 \u001b[0m        \u001b[0;32mimport\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     96 \u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m---> 97 \u001b[0;31m    \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[1;32m     98 \u001b[0m        \u001b[0mall_stacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m     99 \u001b[0m    \u001b[0;31m# little assert statement for sanity check\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m    100 \u001b[0m    \u001b[0;34m[\u001b[0m\u001b[0mverify_sub_words_are_consecutive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mstack\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mall_stacks\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m    101 \u001b[0m    \u001b[0;32mreturn\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mel\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mel\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m''\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcondense_stack\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'word_list'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m    102 \u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n",
      "*** NameError: name 'word_stack' is not defined\n"
     ]
    },
    {
     "ename": "BdbQuit",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mBdbQuit\u001b[0m                                   Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_535920/167074687.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mraw_text\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msample_entry\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mraw_text\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mannotated_tsv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msample_entry\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mannotated_tsv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0mextracted_entities\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mextract_entities\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_text\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mner\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m \u001b[0mtrue_drugs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mextract_drug_names_truth\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mannotated_tsv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0mtrue_advs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mextract_adv_names_truth\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mannotated_tsv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36mextract_entities\u001b[0;34m(text, ner, batch_size)\u001b[0m\n\u001b[1;32m    118\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mextract_entities\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mner\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    119\u001b[0m     \u001b[0;34m'''Extracts the entities, and return them grouped by drug and adverse events'''\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 120\u001b[0;31m     \u001b[0mextracted_entities\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mextract_info\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mner\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    121\u001b[0m     \u001b[0madvs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfilter_entities\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mextracted_entities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mADVERSE_EVENT_TAG\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    122\u001b[0m     \u001b[0mdrugs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfilter_entities\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mextracted_entities\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mDRUG_EVENT_TAG\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36mextract_info\u001b[0;34m(text, ner, batch_size)\u001b[0m\n\u001b[1;32m    108\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mextract_info\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mner\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    109\u001b[0m     \u001b[0;34m'''Extracts Adverse Events and Drug names given a text (using batching and GPU)'''\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 110\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mget_tokens_from_ner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mner_inference\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mner\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    111\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36mget_tokens_from_ner\u001b[0;34m(raw_outputs, entity_list)\u001b[0m\n\u001b[1;32m    103\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_tokens_from_ner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentity_list\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'ADR'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'DRUG'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 105\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mentity_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mget_tokens_from_ner_specific\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentity_type\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mentity_type\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mentity_list\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    103\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_tokens_from_ner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentity_list\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'ADR'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'DRUG'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 105\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mentity_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mget_tokens_from_ner_specific\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mraw_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentity_type\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mentity_type\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mentity_list\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36mget_tokens_from_ner_specific\u001b[0;34m(raw_outputs, entity_type)\u001b[0m\n\u001b[1;32m     95\u001b[0m         \u001b[0;32mimport\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 97\u001b[0;31m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     98\u001b[0m         \u001b[0mall_stacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m     \u001b[0;31m# little assert statement for sanity check\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/projects/mimiciii/a2_spanbert.py\u001b[0m in \u001b[0;36mget_tokens_from_ner_specific\u001b[0;34m(raw_outputs, entity_type)\u001b[0m\n\u001b[1;32m     95\u001b[0m         \u001b[0;32mimport\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mipdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 97\u001b[0;31m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     98\u001b[0m         \u001b[0mall_stacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mword_stack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'curr_stack'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m     \u001b[0;31m# little assert statement for sanity check\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/programs/miniconda3/envs/torch-gpu-env/lib/python3.8/bdb.py\u001b[0m in \u001b[0;36mtrace_dispatch\u001b[0;34m(self, frame, event, arg)\u001b[0m\n\u001b[1;32m     86\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;31m# None\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'line'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'call'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/programs/miniconda3/envs/torch-gpu-env/lib/python3.8/bdb.py\u001b[0m in \u001b[0;36mdispatch_line\u001b[0;34m(self, frame)\u001b[0m\n\u001b[1;32m    111\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbreak_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    112\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquitting\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mBdbQuit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    114\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_dispatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mBdbQuit\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from a2_spanbert import load_csv, load_ner, extract_entities, extract_drug_names_truth, extract_adv_names_truth, load_entry\n",
    "dataset = load_csv()\n",
    "ner = load_ner()\n",
    "sample_index = 2\n",
    "sample_entry = load_entry(dataset,sample_index)\n",
    "raw_text = sample_entry.raw_text\n",
    "annotated_tsv = sample_entry.annotated_tsv\n",
    "extracted_entities = extract_entities(raw_text,ner)\n",
    "true_drugs = extract_drug_names_truth(annotated_tsv)\n",
    "true_advs = extract_adv_names_truth(annotated_tsv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b71a2ba7-8931-4e93-ac91-10f6f09ced11",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Can't mix strings and bytes in path components",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_535920/2144369818.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mcsv_file\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mOUTPUT_DIR\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34mb\"n2c2.csv\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mdf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcsv_file\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/programs/miniconda3/envs/torch-gpu-env/lib/python3.8/posixpath.py\u001b[0m in \u001b[0;36mjoin\u001b[0;34m(a, *p)\u001b[0m\n\u001b[1;32m     88\u001b[0m                 \u001b[0mpath\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0msep\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     89\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mTypeError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mAttributeError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBytesWarning\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 90\u001b[0;31m         \u001b[0mgenericpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_arg_types\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'join'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     91\u001b[0m         \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     92\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/programs/miniconda3/envs/torch-gpu-env/lib/python3.8/genericpath.py\u001b[0m in \u001b[0;36m_check_arg_types\u001b[0;34m(funcname, *args)\u001b[0m\n\u001b[1;32m    153\u001b[0m                             f'os.PathLike object, not {s.__class__.__name__!r}') from None\n\u001b[1;32m    154\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mhasstr\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mhasbytes\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 155\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Can't mix strings and bytes in path components\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: Can't mix strings and bytes in path components"
     ]
    }
   ],
   "source": [
    "csv_file = pjoin(OUTPUT_DIR,b\"n2c2.csv\")\n",
    "df = pd.read_csv(csv_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "33fc9703-42f7-4cf9-83a0-90f467ec37c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>type</th>\n",
       "      <th>name</th>\n",
       "      <th>value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>T1</td>\n",
       "      <td>Drug 19078 19085</td>\n",
       "      <td>Aspirin</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>T2</td>\n",
       "      <td>Strength 19086 19092</td>\n",
       "      <td>325 mg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>T3</td>\n",
       "      <td>Route 19093 19095</td>\n",
       "      <td>PO</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>T4</td>\n",
       "      <td>Drug 19465 19478;19479 19482</td>\n",
       "      <td>Ciprofloxacin HCl</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>T5</td>\n",
       "      <td>Strength 19483 19489</td>\n",
       "      <td>500 mg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>168</th>\n",
       "      <td>R69</td>\n",
       "      <td>Reason-Drug Arg1:T102 Arg2:T59</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>169</th>\n",
       "      <td>R70</td>\n",
       "      <td>Reason-Drug Arg1:T103 Arg2:T16</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>170</th>\n",
       "      <td>R71</td>\n",
       "      <td>ADE-Drug Arg1:T94 Arg2:T8</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>171</th>\n",
       "      <td>T93</td>\n",
       "      <td>Route 12791 12793</td>\n",
       "      <td>IV</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>172</th>\n",
       "      <td>R63</td>\n",
       "      <td>Route-Drug Arg1:T93 Arg2:T92</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>173 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    type                            name              value\n",
       "0     T1                Drug 19078 19085            Aspirin\n",
       "1     T2            Strength 19086 19092             325 mg\n",
       "2     T3               Route 19093 19095                 PO\n",
       "3     T4    Drug 19465 19478;19479 19482  Ciprofloxacin HCl\n",
       "4     T5            Strength 19483 19489             500 mg\n",
       "..   ...                             ...                ...\n",
       "168  R69  Reason-Drug Arg1:T102 Arg2:T59                NaN\n",
       "169  R70  Reason-Drug Arg1:T103 Arg2:T16                NaN\n",
       "170  R71       ADE-Drug Arg1:T94 Arg2:T8                NaN\n",
       "171  T93               Route 12791 12793                 IV\n",
       "172  R63    Route-Drug Arg1:T93 Arg2:T92                NaN\n",
       "\n",
       "[173 rows x 3 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from io import StringIO\n",
    "sample_index = 0\n",
    "row = df.iloc[sample_index]\n",
    "raw_text = row['raw']\n",
    "annotated_tsv = pd.read_csv(StringIO(row['annotated']),sep='\\t',header=None,names=['type','name','value'])\n",
    "annotated_tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cf2d189c-e473-4c4e-b4d7-19f7d57c35a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tokenizer(\"Hello world i have a slight fever\",return_tensors=\"pt\")\n",
    "outputs = model(**inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31b8fe56-e4f7-4f04-8f55-6ffc66801c26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TokenClassifierOutput(loss=None, logits=tensor([[[ 1.9963, -0.7155, -0.7707, -0.1233, -0.4926],\n",
       "         [ 3.8603, -1.0510, -1.3438, -1.0383, -1.1842],\n",
       "         [ 3.9350, -1.2033, -1.1814, -1.3157, -1.0649],\n",
       "         [ 3.7329, -0.9513, -1.3155, -1.1470, -1.1257],\n",
       "         [ 3.9507, -1.3179, -1.1036, -1.3013, -1.0481],\n",
       "         [ 3.9218, -1.3843, -1.2534, -1.0409, -1.0559],\n",
       "         [ 2.6338, -1.6201, -1.6524,  1.0794, -1.1400],\n",
       "         [-0.9270, -1.1644, -1.2274,  3.3766, -0.2071],\n",
       "         [ 1.9958, -0.7154, -0.7709, -0.1228, -0.4926]]],\n",
       "       grad_fn=<AddBackward0>), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "65d5c27a-f71c-4304-8a2e-76c727041078",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'entity': 'B-ADR',\n",
       "  'score': 0.94164777,\n",
       "  'index': 7,\n",
       "  'word': 'fever',\n",
       "  'start': 28,\n",
       "  'end': 33}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ner(\"Hello world i have a slight fever\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bd8ed6a7-9ddf-49ec-aa76-e56e6898d443",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = ner(raw_text[:512])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1d01963e-6feb-4b4b-9b28-71265da65e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from functools import reduce\n",
    "from typing import List, Dict\n",
    "def verify_sub_words_are_consecutive(stack):\n",
    "    indices = [el['index'] for subword in stack]\n",
    "    assert reduce(lambda holds_true_prev, curr_pair: curr_pair[0] + 1 == curr_pair[1],zip(indices,indices[1:]),True)\n",
    "\n",
    "def word_accumulator_fn(info_dict : Dict[str,List],curr_tok):\n",
    "    if curr_tok['entity'].startswith('I'):\n",
    "        if len(info_dict['curr_stack']) == 0:\n",
    "            return info_dict # just don't modify it\n",
    "        try:\n",
    "            assert info_dict['curr_stack'][0]['entity'].startswith('B'), \"WHY DOES IT TAG SOMETHING WITHOUT A BEGINNING WTF?\"\n",
    "        except AssertionError:\n",
    "            import ipdb; ipdb.set_trace()\n",
    "        info_dict['curr_stack'].append(curr_tok)\n",
    "    else: #starts with B\n",
    "        assert curr_tok['entity'].startswith('B'), \"THERE IS SOME SUSSY STUFF GOING ON WITH THE TAGGER. OUTPUT OTHER THAN I- OR B-???\"\n",
    "        info_dict['word_list'].append(info_dict['curr_stack'])\n",
    "        info_dict['curr_stack'] = [curr_tok]\n",
    "    return info_dict\n",
    "        \n",
    "def condense_stack(stack):\n",
    "    condensed_word = ''.join([part['word'][2:] if part['word'].startswith(\"##\") else ' ' + part['word'] for part in stack]).lstrip()\n",
    "    avg_score = np.mean([part['score'] for part in stack])\n",
    "    return condensed_word, avg_score\n",
    "\n",
    "def get_tokens_from_ner_specific(raw_outputs,entity_type):\n",
    "    '''gets proper tokens from nlp pipeline (merges subwords as well)'''\n",
    "    # output is a list of {entity : str, score : float, index : int, word : str, start : int, end : int}\n",
    "    # goal: merge consecutive indices into one word\n",
    "    outputs = filter(lambda output: output['entity'].endswith(entity_type),raw_outputs)\n",
    "    word_stack = reduce(word_accumulator_fn,outputs,dict(word_list=[],curr_stack=[]))\n",
    "    all_stacks = word_stack['word_list']\n",
    "    if len(word_stack['curr_stack']) > 0: \n",
    "        all_stacks.append(word_stack['curr_stack'])\n",
    "    (verify_sub_words_are_consecutive(stack) for stack in all_stacks) # little assert statement for sanity check \n",
    "    return list(filter(lambda el: el[0] != '',map(condense_stack,word_stack['word_list'])))\n",
    "\n",
    "def get_tokens_from_ner(raw_outputs,entity_list=['ADR','DRUG']):\n",
    "    return {entity_type:get_tokens_from_ner_specific(raw_outputs,entity_type) for entity_type in entity_list} \n",
    "\n",
    "def identify_info(text):\n",
    "    return get_tokens_from_ner(ner(text))\n",
    "\n",
    "def identify_infos(texts):\n",
    "    return [get_tokens_from_ner(output) for output in ner(texts)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "31403bdd-3472-4385-96a3-368cd5bac943",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ADR': [('abdominal pain', 0.93423426),\n",
       "  ('chest and abdominal pain', 0.8709032),\n",
       "  ('odynophagia', 0.962869)],\n",
       " 'DRUG': [('ativan', 0.96955997)]}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "identify_info(raw_text[:512])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bed63a6e-e255-4399-99b1-8810e5702472",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['Aspirin', '325 mg', 'PO', 'Ciprofloxacin HCl', '500 mg', 'Q12H',\n",
       "       'PO', 'ativan', 'butalbital -acetaminophen-caff',\n",
       "       'Metoprolol Tartrate', '12.5 mg', 'BID', 'PO', 'Omeprazole',\n",
       "       '40 mg', 'Ativan', 'DiCYCLOmine', '10 mg', 'QID', 'PO', 'Aspirin',\n",
       "       'Ativan', '50 mg -325 mg-40 mg', 'Omeprazole', '40 mg', 'PO',\n",
       "       'Acetaminophen-Caff- Butalbital', 'TAB', 'Q6H :PRN', 'PO',\n",
       "       'ciprofloxacin [Cipro]', '500 mg', 'tablet (s)', '1', 'by mouth',\n",
       "       'Aspirin EC', '325 mg', 'PO', 'metoprolol tartrate', '25 mg',\n",
       "       'tablet (s)', 'by mouth', 'ATIVAN', 'Aspirin', '325mg',\n",
       "       'once daily', 'Atorvastatin', '20 mg', 'PO', 'DiCYCLOmine',\n",
       "       '10 mg', 'QID', 'PO', 'Atorvastatin', '20 mg', 'PO', 'Tylenol',\n",
       "       'PPI', 'GI cocktail', 'delerious', 'Hadol', 'agitation', 'PO',\n",
       "       'DAILY', 'Donnatol', '10 mL', 'PO', 'BID:PRN', 'abdominal pain',\n",
       "       'DAILY', 'DAILY', 'DAILY', 'DAILY', 'DAILY', 'headache',\n",
       "       'tablet(s)', 'by mouth', 'q 4 hours', '7 Days', 'twice a day',\n",
       "       '0.5 (One half)', 'twice a day', 'Donnatol', '10 mL', 'PO',\n",
       "       'BID:PRN', 'abdominal pain', 'very confused', nan, nan, nan, nan,\n",
       "       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,\n",
       "       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,\n",
       "       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,\n",
       "       nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,\n",
       "       nan, nan, nan, nan, nan, nan, 'CONTRAST', 'contrast', 'contrast',\n",
       "       'contrast', 'acute delerium', 'pain medications',\n",
       "       'pain medication', 'stool softener', 'To avoid constipation',\n",
       "       'abdominal pain', 'odynophagia', 'agitated', 'post procedure pain',\n",
       "       nan, nan, nan, nan, nan, nan, nan, nan, 'IV', nan], dtype=object)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "annotated_tsv['value'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "af5d9322-e67c-4303-80f8-1b4f6a62bbed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.tokenize import sent_tokenize\n",
    "\n",
    "def predict_text(text,key):\n",
    "    return [ad for sentence_output in identify_infos(sent_tokenize(text)) for ad in sentence_output[key]]\n",
    "\n",
    "def extract_entities_text(text):\n",
    "    return {key: predict_text(text,key) for key in ['ADR','DRUG']}\n",
    "\n",
    "def extract_drug_names_pred(text):\n",
    "    return predict_text(text,'DRUG')\n",
    "\n",
    "def extract_adv_names_pred(text):\n",
    "    return predict_text(text,'ADR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "cdc856b8-22b3-4dfd-aa0c-b66d2cb52862",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>type</th>\n",
       "      <th>name</th>\n",
       "      <th>value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>154</th>\n",
       "      <td>T94</td>\n",
       "      <td>ADE 17894 17908</td>\n",
       "      <td>acute delerium</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    type             name           value\n",
       "154  T94  ADE 17894 17908  acute delerium"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "annotated_tsv[annotated_tsv['type'] == \"T94\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "1ed751b0-2a05-4abd-b961-934ff0861bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_drug(row):\n",
    "    return row['name'].startswith(\"Drug\") and row['type'].startswith(\"T\")\n",
    "\n",
    "def is_adv(row):\n",
    "    return row['name'].startswith(\"ADE\") and row['type'].startswith(\"T\")\n",
    "def extract_drug_names_truth(annotated_tsv):\n",
    "    return [row['value'].lower() for i , row in annotated_tsv.iterrows() if is_drug(row)]\n",
    "\n",
    "def extract_adv_names_truth(annotated_tsv):\n",
    "    return [row['value'].lower() for i , row in annotated_tsv.iterrows() if is_adv(row)]\n",
    "# matching_rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6967e74-fedb-478f-ad6f-d39cfc3aba1e",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}