NER_ML_BACKEND.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca69e07e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:30:54.168896Z",
     "start_time": "2022-09-15T19:30:54.161209Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import re\n",
    "import os\n",
    "import io\n",
    "import logging\n",
    "\n",
    "from functools import partial\n",
    "from itertools import groupby\n",
    "from operator import itemgetter\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm, trange\n",
    "from tensorboardX import SummaryWriter\n",
    "from collections import deque\n",
    "\n",
    "from transformers import (\n",
    "    BertTokenizer, BertForTokenClassification, BertConfig,\n",
    "    RobertaConfig, RobertaForTokenClassification, RobertaTokenizer,\n",
    "    DistilBertConfig, DistilBertForTokenClassification, DistilBertTokenizer,\n",
    "    CamembertConfig, CamembertForTokenClassification, CamembertTokenizer,\n",
    "    AutoConfig, AutoModelForTokenClassification, AutoTokenizer,\n",
    "    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,\n",
    ")\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "from label_studio_ml.model import LabelStudioMLBase\n",
    "import utils \n",
    "\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd084d65",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:30:56.157873Z",
     "start_time": "2022-09-15T19:30:56.153923Z"
    }
   },
   "outputs": [],
   "source": [
    "# pip install utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839b7486",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:30:56.413929Z",
     "start_time": "2022-09-15T19:30:56.410164Z"
    }
   },
   "outputs": [],
   "source": [
    "# pip install -U requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64014514",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77f20d89",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:30:56.770808Z",
     "start_time": "2022-09-15T19:30:56.761643Z"
    }
   },
   "outputs": [],
   "source": [
    "ALL_MODELS = sum(\n",
    "    [list(conf.keys()) for conf in (BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP)],\n",
    "    [])\n",
    "\n",
    "MODEL_CLASSES = {\n",
    "    'bert': (BertConfig, BertForTokenClassification, BertTokenizer),\n",
    "    'roberta': (RobertaConfig, RobertaForTokenClassification, RobertaTokenizer),\n",
    "    'distilbert': (DistilBertConfig, DistilBertForTokenClassification, DistilBertTokenizer),\n",
    "    'camembert': (CamembertConfig, CamembertForTokenClassification, CamembertTokenizer),\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e470b9c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:39:29.981850Z",
     "start_time": "2022-09-15T19:39:29.962064Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['bert-base-uncased',\n",
       " 'bert-large-uncased',\n",
       " 'bert-base-cased',\n",
       " 'bert-large-cased',\n",
       " 'bert-base-multilingual-uncased',\n",
       " 'bert-base-multilingual-cased',\n",
       " 'bert-base-chinese',\n",
       " 'bert-base-german-cased',\n",
       " 'bert-large-uncased-whole-word-masking',\n",
       " 'bert-large-cased-whole-word-masking',\n",
       " 'bert-large-uncased-whole-word-masking-finetuned-squad',\n",
       " 'bert-large-cased-whole-word-masking-finetuned-squad',\n",
       " 'bert-base-cased-finetuned-mrpc',\n",
       " 'bert-base-german-dbmdz-cased',\n",
       " 'bert-base-german-dbmdz-uncased',\n",
       " 'cl-tohoku/bert-base-japanese',\n",
       " 'cl-tohoku/bert-base-japanese-whole-word-masking',\n",
       " 'cl-tohoku/bert-base-japanese-char',\n",
       " 'cl-tohoku/bert-base-japanese-char-whole-word-masking',\n",
       " 'TurkuNLP/bert-base-finnish-cased-v1',\n",
       " 'TurkuNLP/bert-base-finnish-uncased-v1',\n",
       " 'wietsedv/bert-base-dutch-cased',\n",
       " 'roberta-base',\n",
       " 'roberta-large',\n",
       " 'roberta-large-mnli',\n",
       " 'distilroberta-base',\n",
       " 'roberta-base-openai-detector',\n",
       " 'roberta-large-openai-detector',\n",
       " 'distilbert-base-uncased',\n",
       " 'distilbert-base-uncased-distilled-squad',\n",
       " 'distilbert-base-cased',\n",
       " 'distilbert-base-cased-distilled-squad',\n",
       " 'distilbert-base-german-cased',\n",
       " 'distilbert-base-multilingual-cased',\n",
       " 'distilbert-base-uncased-finetuned-sst-2-english']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ALL_MODELS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76aace71",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:31:24.016542Z",
     "start_time": "2022-09-15T19:31:23.935739Z"
    }
   },
   "outputs": [],
   "source": [
    "class SpanLabeledTextDataset(Dataset):\n",
    "\n",
    "    def __init__(\n",
    "        self, list_of_strings, list_of_spans=None, tokenizer=None, tag_idx_map=None,\n",
    "        cls_token='[CLS]', sep_token='[SEP]', pad_token_label_id=-1, max_seq_length=128, sep_token_extra=False,\n",
    "        cls_token_at_end=False, sequence_a_segment_id=0, cls_token_segment_id=1, mask_padding_with_zero=True\n",
    "    ):\n",
    "        self.list_of_strings = list_of_strings\n",
    "        self.list_of_spans = list_of_spans or [[] * len(list_of_strings)]\n",
    "        self.tokenizer = tokenizer\n",
    "        self.cls_token = cls_token\n",
    "        self.sep_token = sep_token\n",
    "        self.pad_token_label_id = pad_token_label_id\n",
    "        self.max_seq_length = max_seq_length\n",
    "        self.sep_token_extra = sep_token_extra\n",
    "        self.cls_token_at_end = cls_token_at_end\n",
    "        self.sequence_a_segment_id = sequence_a_segment_id\n",
    "        self.cls_token_segment_id = cls_token_segment_id\n",
    "        self.mask_padding_with_zero = mask_padding_with_zero\n",
    "\n",
    "        (self.original_list_of_tokens, self.original_list_of_tags, tag_idx_map_,\n",
    "         original_list_of_tokens_start_map) = self._prepare_data()\n",
    "\n",
    "        if tag_idx_map is None:\n",
    "            self.tag_idx_map = tag_idx_map_\n",
    "        else:\n",
    "            self.tag_idx_map = tag_idx_map\n",
    "\n",
    "        (self.list_of_tokens, self.list_of_token_ids, self.list_of_labels, self.list_of_label_ids,\n",
    "         self.list_of_segment_ids, self.list_of_token_start_map) = [], [], [], [], [], []\n",
    "\n",
    "        for original_tokens, original_tags, original_token_start_map in zip(\n",
    "            self.original_list_of_tokens,\n",
    "            self.original_list_of_tags,\n",
    "            original_list_of_tokens_start_map\n",
    "        ):\n",
    "            tokens, token_ids, labels, label_ids, segment_ids, token_start_map = self._convert_to_features(\n",
    "                original_tokens, original_tags, self.tag_idx_map, original_token_start_map)\n",
    "            self.list_of_token_ids.append(token_ids)\n",
    "            self.list_of_tokens.append(tokens)\n",
    "            self.list_of_labels.append(labels)\n",
    "            self.list_of_segment_ids.append(segment_ids)\n",
    "            self.list_of_label_ids.append(label_ids)\n",
    "            self.list_of_token_start_map.append(token_start_map)\n",
    "\n",
    "    def get_params_dict(self):\n",
    "        return {\n",
    "            'cls_token': self.cls_token,\n",
    "            'sep_token': self.sep_token,\n",
    "            'pad_token_label_id': self.pad_token_label_id,\n",
    "            'max_seq_length': self.max_seq_length,\n",
    "            'sep_token_extra': self.sep_token_extra,\n",
    "            'cls_token_at_end': self.cls_token_at_end,\n",
    "            'sequence_a_segment_id': self.sequence_a_segment_id,\n",
    "            'cls_token_segment_id': self.cls_token_segment_id,\n",
    "            'mask_padding_with_zero': self.mask_padding_with_zero\n",
    "        }\n",
    "\n",
    "    def dump(self, output_file):\n",
    "        with io.open(output_file, mode='w') as f:\n",
    "            for tokens, labels in zip(self.list_of_tokens, self.list_of_labels):\n",
    "                for token, label in zip(tokens, labels):\n",
    "                    f.write(f'{token} {label}\\n')\n",
    "                f.write('\\n')\n",
    "\n",
    "    def _convert_to_features(self, words, labels, label_map, list_token_start_map):\n",
    "        tokens, out_labels, label_ids, tokens_idx_map = [], [], [], []\n",
    "        for i, (word, label, token_start) in enumerate(zip(words, labels, list_token_start_map)):\n",
    "            word_tokens = self.tokenizer.tokenize(word)\n",
    "            tokens.extend(word_tokens)\n",
    "            tokens_idx_map.extend([token_start] * len(word_tokens))\n",
    "            # Use the real label id for the first token of the word, and padding ids for the remaining tokens\n",
    "            label_ids.extend([label_map[label]] + [self.pad_token_label_id] * (len(word_tokens) - 1))\n",
    "            out_labels.extend([label] + ['X'] * (len(word_tokens) - 1))\n",
    "\n",
    "        # Account for [CLS] and [SEP] with \"- 2\" and with \"- 3\" for RoBERTa.\n",
    "        special_tokens_count = 3 if self.sep_token_extra else 2\n",
    "        if len(tokens) > self.max_seq_length - special_tokens_count:\n",
    "            tokens = tokens[:(self.max_seq_length - special_tokens_count)]\n",
    "            label_ids = label_ids[:(self.max_seq_length - special_tokens_count)]\n",
    "            out_labels = out_labels[:(self.max_seq_length - special_tokens_count)]\n",
    "            tokens_idx_map = tokens_idx_map[:(self.max_seq_length - special_tokens_count)]\n",
    "\n",
    "        # The convention in BERT is:\n",
    "        # (a) For sequence pairs:\n",
    "        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]\n",
    "        #  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1\n",
    "        # (b) For single sequences:\n",
    "        #  tokens:   [CLS] the dog is hairy . [SEP]\n",
    "        #  type_ids:   0   0   0   0  0     0   0\n",
    "        #\n",
    "        # Where \"type_ids\" are used to indicate whether this is the first\n",
    "        # sequence or the second sequence. The embedding vectors for `type=0` and\n",
    "        # `type=1` were learned during pre-training and are added to the wordpiece\n",
    "        # embedding vector (and position vector). This is not *strictly* necessary\n",
    "        # since the [SEP] token unambiguously separates the sequences, but it makes\n",
    "        # it easier for the model to learn the concept of sequences.\n",
    "        #\n",
    "        # For classification tasks, the first vector (corresponding to [CLS]) is\n",
    "        # used as as the \"sentence vector\". Note that this only makes sense because\n",
    "        # the entire model is fine-tuned.\n",
    "        tokens += [self.sep_token]\n",
    "        label_ids += [self.pad_token_label_id]\n",
    "        out_labels += ['X']\n",
    "        tokens_idx_map += [-1]\n",
    "        if self.sep_token_extra:\n",
    "            # roberta uses an extra separator b/w pairs of sentences\n",
    "            tokens += [self.sep_token]\n",
    "            label_ids += [self.pad_token_label_id]\n",
    "            out_labels += ['X']\n",
    "            tokens_idx_map += [-1]\n",
    "        segment_ids = [self.sequence_a_segment_id] * len(tokens)\n",
    "        if self.cls_token_at_end:\n",
    "            tokens += [self.cls_token]\n",
    "            label_ids += [self.pad_token_label_id]\n",
    "            out_labels += ['X']\n",
    "            segment_ids += [self.cls_token_segment_id]\n",
    "            tokens_idx_map += [-1]\n",
    "        else:\n",
    "            tokens = [self.cls_token] + tokens\n",
    "            label_ids = [self.pad_token_label_id] + label_ids\n",
    "            out_labels = ['X'] + out_labels\n",
    "            segment_ids = [self.cls_token_segment_id] + segment_ids\n",
    "            tokens_idx_map = [-1] + tokens_idx_map\n",
    "\n",
    "        token_ids = self.tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "        return tokens, token_ids, out_labels, label_ids, segment_ids, tokens_idx_map\n",
    "\n",
    "    def _apply_tokenizer(self, original_tokens, original_tags):\n",
    "        out_tokens, out_tags, out_maps = [], [], []\n",
    "        for i, (original_token, original_tag) in enumerate(zip(original_tokens, original_tags)):\n",
    "            tokens = self.tokenizer.tokenize(original_token)\n",
    "            out_tokens.extend(tokens)\n",
    "            out_maps.extend([i] * len(tokens))\n",
    "            start_tag = original_tag.startswith('B-')\n",
    "            for j in range(len(tokens)):\n",
    "                if (j == 0 and start_tag) or original_tag == 'O':\n",
    "                    out_tags.append(original_tag)\n",
    "                else:\n",
    "                    out_tags.append(f'I-{original_tag[2:]}')\n",
    "        return out_tokens, out_tags, out_maps\n",
    "\n",
    "    def _prepare_data(self):\n",
    "        list_of_tokens, list_of_tags, list_of_token_idx_maps = [], [], []\n",
    "        tag_idx_map = {'O': 0}\n",
    "        for text, spans in zip(self.list_of_strings, self.list_of_spans):\n",
    "            if not text:\n",
    "                continue\n",
    "\n",
    "            tokens = []\n",
    "            start = 0\n",
    "            for t in text.split():\n",
    "                tokens.append((t, start))\n",
    "                start += len(t) + 1\n",
    "\n",
    "            if spans:\n",
    "                spans = list(sorted(spans, key=itemgetter('start')))\n",
    "                span = spans.pop(0)\n",
    "                prefix = 'B-'\n",
    "                tags = []\n",
    "                for token, token_start in tokens:\n",
    "                    token_end = token_start + len(token) - 1\n",
    "\n",
    "                    # token precedes current span\n",
    "                    if not span or token_end < span['start']:\n",
    "                        tags.append('O')\n",
    "                        continue\n",
    "\n",
    "                    # token jumps over the span (it could happens\n",
    "                    # when prev label ends with whitespaces, e.g. \"cat \" \"too\" or span created for whitespace)\n",
    "                    if token_start > span['end']:\n",
    "\n",
    "                        prefix = 'B-'\n",
    "                        no_more_spans = False\n",
    "                        while token_start > span['end']:\n",
    "                            if not len(spans):\n",
    "                                no_more_spans = True\n",
    "                                break\n",
    "                            span = spans.pop(0)\n",
    "\n",
    "                        if no_more_spans:\n",
    "                            tags.append('O')\n",
    "                            span = None\n",
    "                            continue\n",
    "\n",
    "                        if token_end < span['start']:\n",
    "                            tags.append('O')\n",
    "                            continue\n",
    "\n",
    "                    label = span['label']\n",
    "                    if label.startswith(prefix):\n",
    "                        tag = label\n",
    "                    else:\n",
    "                        tag = f'{prefix}{label}'\n",
    "                    tags.append(tag)\n",
    "                    if tag not in tag_idx_map:\n",
    "                        tag_idx_map[tag] = len(tag_idx_map)\n",
    "                    if span['end'] > token_end:\n",
    "                        prefix = 'I-'\n",
    "                    elif len(spans):\n",
    "                        span = spans.pop(0)\n",
    "                        prefix = 'B-'\n",
    "                    else:\n",
    "                        span = None\n",
    "            else:\n",
    "                tags = ['O'] * len(tokens)\n",
    "\n",
    "            list_of_tokens.append([t[0] for t in tokens])\n",
    "            list_of_token_idx_maps.append([t[1] for t in tokens])\n",
    "            list_of_tags.append(tags)\n",
    "\n",
    "        return list_of_tokens, list_of_tags, tag_idx_map, list_of_token_idx_maps\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.list_of_token_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            'tokens': self.list_of_token_ids[idx],\n",
    "            'labels': self.list_of_label_ids[idx],\n",
    "            'segments': self.list_of_segment_ids[idx],\n",
    "            'token_start_map': self.list_of_token_start_map[idx],\n",
    "            'string': self.list_of_strings[idx]\n",
    "        }\n",
    "\n",
    "    @property\n",
    "    def num_labels(self):\n",
    "        return len(self.tag_idx_map)\n",
    "\n",
    "    @classmethod\n",
    "    def pad_sequences(cls, batch, mask_padding_with_zero, pad_on_left, pad_token, pad_token_segment_id, pad_token_label_id):\n",
    "        # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "        # tokens are attended to.\n",
    "        max_seq_length = max(map(len, (sample['tokens'] for sample in batch)))\n",
    "        batch_input_ids, batch_label_ids, batch_segment_ids, batch_input_mask, batch_token_start_map = [], [], [], [], []\n",
    "        batch_strings = []\n",
    "        for sample in batch:\n",
    "            input_ids = sample['tokens']\n",
    "            label_ids = sample['labels']\n",
    "            segment_ids = sample['segments']\n",
    "            # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "            # tokens are attended to.\n",
    "            input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)\n",
    "            padding_length = max_seq_length - len(input_ids)\n",
    "            if pad_on_left:\n",
    "                input_ids = ([pad_token] * padding_length) + input_ids\n",
    "                input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask\n",
    "                segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids\n",
    "                label_ids = ([pad_token_label_id] * padding_length) + label_ids\n",
    "            else:\n",
    "                input_ids += ([pad_token] * padding_length)\n",
    "                input_mask += ([0 if mask_padding_with_zero else 1] * padding_length)\n",
    "                segment_ids += ([pad_token_segment_id] * padding_length)\n",
    "                label_ids += ([pad_token_label_id] * padding_length)\n",
    "            batch_input_ids.append(input_ids)\n",
    "            batch_label_ids.append(label_ids)\n",
    "            batch_segment_ids.append(segment_ids)\n",
    "            batch_input_mask.append(input_mask)\n",
    "            batch_token_start_map.append(sample['token_start_map'])\n",
    "            batch_strings.append(sample['string'])\n",
    "\n",
    "        return {\n",
    "            'input_ids': torch.tensor(batch_input_ids, dtype=torch.long),\n",
    "            'label_ids': torch.tensor(batch_label_ids, dtype=torch.long),\n",
    "            'segment_ids': torch.tensor(batch_segment_ids, dtype=torch.long),\n",
    "            'input_mask': torch.tensor(batch_input_mask, dtype=torch.long),\n",
    "            'token_start_map': batch_token_start_map,\n",
    "            'strings': batch_strings\n",
    "        }\n",
    "\n",
    "    @classmethod\n",
    "    def get_padding_function(cls, model_type, tokenizer, pad_token_label_id):\n",
    "        return partial(\n",
    "            cls.pad_sequences,\n",
    "            mask_padding_with_zero=True,\n",
    "            pad_on_left=model_type in ['xlnet'],\n",
    "            pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],\n",
    "            pad_token_segment_id=4 if model_type in ['xlnet'] else 0,\n",
    "            pad_token_label_id=pad_token_label_id\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ff846b8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-09-15T19:31:26.614787Z",
     "start_time": "2022-09-15T19:31:26.536738Z"
    }
   },
   "outputs": [],
   "source": [
    "class TransformersBasedTagger(LabelStudioMLBase):\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super(TransformersBasedTagger, self).__init__(**kwargs)\n",
    "\n",
    "        assert len(self.parsed_label_config) == 1\n",
    "        self.from_name, self.info = list(self.parsed_label_config.items())[0]\n",
    "        assert self.info['type'] == 'Labels'\n",
    "\n",
    "        # the model has only one textual input\n",
    "        assert len(self.info['to_name']) == 1\n",
    "        assert len(self.info['inputs']) == 1\n",
    "        assert self.info['inputs'][0]['type'] == 'Text'\n",
    "        self.to_name = self.info['to_name'][0]\n",
    "        self.value = self.info['inputs'][0]['value']\n",
    "\n",
    "        if not self.train_output:\n",
    "            self.labels = self.info['labels']\n",
    "        else:\n",
    "            self.load(self.train_output)\n",
    "\n",
    "    def load(self, train_output):\n",
    "        pretrained_model = train_output['model_path']\n",
    "        self._model_type = train_output['model_type']\n",
    "        _, model_class, tokenizer_class = MODEL_CLASSES[train_output['model_type']]\n",
    "\n",
    "        self._tokenizer = AutoTokenizer.from_pretrained(pretrained_model)\n",
    "        self._model = AutoModelForTokenClassification.from_pretrained(pretrained_model)\n",
    "        self._batch_size = train_output['batch_size']\n",
    "        self._pad_token = self._tokenizer.convert_tokens_to_ids([self._tokenizer.pad_token])[0]\n",
    "        self._pad_token_label_id = train_output['pad_token_label_id']\n",
    "        self._label_map = train_output['label_map']\n",
    "        self._mask_padding_with_zero = True\n",
    "        self._dataset_params_dict = train_output['dataset_params_dict']\n",
    "\n",
    "        self._batch_padding = SpanLabeledTextDataset.get_padding_function(\n",
    "            self._model_type, self._tokenizer, self._pad_token_label_id)\n",
    "\n",
    "    def predict(self, tasks, **kwargs):\n",
    "        if not hasattr(self, \"_tokenizer\"):\n",
    "            logger.error(\"No model loaded! Please train model before getting predictions!\")\n",
    "            return {}\n",
    "        texts = [task['data'][self.value] for task in tasks]\n",
    "        predict_set = SpanLabeledTextDataset(texts, tokenizer=self._tokenizer, **self._dataset_params_dict)\n",
    "        from_name = self.from_name\n",
    "        to_name = self.to_name\n",
    "        predict_loader = DataLoader(\n",
    "            dataset=predict_set,\n",
    "            batch_size=self._batch_size,\n",
    "            collate_fn=self._batch_padding\n",
    "        )\n",
    "\n",
    "        results = []\n",
    "        for batch in tqdm(predict_loader, desc='Prediction'):\n",
    "            inputs = {\n",
    "                'input_ids': batch['input_ids'],\n",
    "                'attention_mask': batch['input_mask'],\n",
    "                'token_type_ids': batch['segment_ids']\n",
    "            }\n",
    "            if self._model_type == 'distilbert':\n",
    "                inputs.pop('token_type_ids')\n",
    "            with torch.no_grad():\n",
    "                model_output = self._model(**inputs)\n",
    "                logits = model_output[0]\n",
    "\n",
    "            batch_preds = logits.detach().cpu().numpy()\n",
    "            argmax_batch_preds = np.argmax(batch_preds, axis=-1)\n",
    "            max_batch_preds = np.max(batch_preds, axis=-1)\n",
    "            input_mask = batch['input_mask'].detach().cpu().numpy()\n",
    "            batch_token_start_map = batch['token_start_map']\n",
    "            batch_strings = batch['strings']\n",
    "\n",
    "            for max_preds, argmax_preds, mask_tokens, token_start_map, string in zip(\n",
    "                max_batch_preds, argmax_batch_preds, input_mask, batch_token_start_map, batch_strings\n",
    "            ):\n",
    "                preds, scores, starts = [], [], []\n",
    "                for max_pred, argmax_pred, mask_token, token_start in zip(max_preds, argmax_preds, mask_tokens, token_start_map):\n",
    "                    if token_start != -1:\n",
    "                        preds.append(self._label_map[str(argmax_pred)])\n",
    "                        scores.append(max_pred)\n",
    "                        starts.append(token_start)\n",
    "                mean_score = np.mean(scores) if len(scores) > 0 else 0\n",
    "\n",
    "                result = []\n",
    "\n",
    "                for label, group in groupby(zip(preds, starts, scores), key=lambda i: re.sub('^(B-|I-)', '', i[0])):\n",
    "                    _, group_start, _ = list(group)[0]\n",
    "                    if len(result) > 0:\n",
    "                        if group_start == 0:\n",
    "                            result.pop(-1)\n",
    "                        else:\n",
    "                            result[-1]['value']['end'] = group_start - 1\n",
    "                    if label != 'O':\n",
    "                        result.append({\n",
    "                            'from_name': from_name,\n",
    "                            'to_name': to_name,\n",
    "                            'type': 'labels',\n",
    "                            'value': {\n",
    "                                'labels': [label],\n",
    "                                'start': group_start,\n",
    "                                'end': None, \n",
    "                                'text': '...'\n",
    "                            }\n",
    "                        })\n",
    "                if result and result[-1]['value']['end'] is None:\n",
    "                    result[-1]['value']['end'] = len(string)\n",
    "                results.append({\n",
    "                    'result': result,\n",
    "                    'score': float(mean_score),\n",
    "                    'cluster': None\n",
    "                })\n",
    "        return results\n",
    "\n",
    "    def get_spans(self, completion):\n",
    "        spans = []\n",
    "        for r in completion['result']:\n",
    "            if r['from_name'] == self.from_name and r['to_name'] == self.to_name:\n",
    "                labels = r['value'].get('labels')\n",
    "                if not isinstance(labels, list) or len(labels) == 0:\n",
    "                    logger.warning(f'Error while parsing {r}: list type expected for \"labels\"')\n",
    "                    continue\n",
    "                label = labels[0]\n",
    "                start, end = r['value'].get('start'), r['value'].get('end')\n",
    "                if start is None or end is None:\n",
    "                    logger.warning(f'Error while parsing {r}: \"labels\" should contain \"start\" and \"end\" fields')\n",
    "                spans.append({\n",
    "                    'label': label,\n",
    "                    'start': start,\n",
    "                    'end': end\n",
    "                })\n",
    "        return spans\n",
    "\n",
    "    def fit(\n",
    "        self, completions, workdir=None, model_type='bert', pretrained_model='hemangjoshi37a/autotrain-ratnakar_1000_sample_curated-1474454086',\n",
    "        batch_size=32, learning_rate=5e-5, adam_epsilon=1e-8, num_train_epochs=100, weight_decay=0.0, logging_steps=1,\n",
    "        warmup_steps=0, save_steps=50, dump_dataset=True, cache_dir='~/.heartex/cache', train_logs=None,\n",
    "        **kwargs\n",
    "    ):\n",
    "        train_logs = train_logs or os.path.join(workdir, 'train_logs')\n",
    "        os.makedirs(train_logs, exist_ok=True)\n",
    "        logger.debug('Prepare models')\n",
    "        cache_dir = os.path.expanduser(cache_dir)\n",
    "        os.makedirs(cache_dir, exist_ok=True)\n",
    "\n",
    "        model_type = model_type.lower()\n",
    "        # assert model_type in MODEL_CLASSES.keys(), f'Input model type {model_type} not in {MODEL_CLASSES.keys()}'\n",
    "        # assert pretrained_model in ALL_MODELS, f'Pretrained model {pretrained_model} not in {ALL_MODELS}'\n",
    "\n",
    "        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)\n",
    "\n",
    "        logger.debug('Read data')\n",
    "        # read input data stream\n",
    "        texts, list_of_spans = [], []\n",
    "        for item in completions:\n",
    "            texts.append(item['data'][self.value])\n",
    "            list_of_spans.append(self.get_spans(item['annotations'][0]))\n",
    "\n",
    "        logger.debug('Prepare dataset')\n",
    "        pad_token_label_id = CrossEntropyLoss().ignore_index\n",
    "        train_set = SpanLabeledTextDataset(\n",
    "            texts, list_of_spans, tokenizer,\n",
    "            cls_token_at_end=model_type in ['xlnet'],\n",
    "            cls_token_segment_id=2 if model_type in ['xlnet'] else 0,\n",
    "            sep_token_extra=model_type in ['roberta'],\n",
    "            pad_token_label_id=pad_token_label_id\n",
    "        )\n",
    "\n",
    "        if dump_dataset:\n",
    "            dataset_file = os.path.join(workdir, 'train_set.txt')\n",
    "            train_set.dump(dataset_file)\n",
    "\n",
    "        # config = config_class.from_pretrained(pretrained_model, num_labels=train_set.num_labels, cache_dir=cache_dir)\n",
    "        config = AutoConfig.from_pretrained(pretrained_model, num_labels=train_set.num_labels, cache_dir=cache_dir)\n",
    "        # model = model_class.from_pretrained(pretrained_model, config=config, cache_dir=cache_dir)\n",
    "        model = AutoModelForTokenClassification.from_pretrained(pretrained_model, config=config, cache_dir=cache_dir)\n",
    "\n",
    "        batch_padding = SpanLabeledTextDataset.get_padding_function(model_type, tokenizer, pad_token_label_id)\n",
    "\n",
    "        train_loader = DataLoader(\n",
    "            dataset=train_set,\n",
    "            batch_size=batch_size,\n",
    "            shuffle=True,\n",
    "            collate_fn=batch_padding\n",
    "        )\n",
    "\n",
    "        no_decay = ['bias', 'LayerNorm.weight']\n",
    "        optimizer_grouped_parameters = [\n",
    "            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
    "             'weight_decay': weight_decay},\n",
    "            {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "        ]\n",
    "\n",
    "        num_training_steps = len(train_loader) * num_train_epochs\n",
    "        optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)\n",
    "        scheduler = get_linear_schedule_with_warmup(\n",
    "            optimizer, num_warmup_steps=warmup_steps, num_training_steps=num_training_steps)\n",
    "\n",
    "        tr_loss, logging_loss = 0, 0\n",
    "        global_step = 0\n",
    "        if train_logs:\n",
    "            tb_writer = SummaryWriter(logdir=os.path.join(train_logs, os.path.basename(workdir)))\n",
    "        epoch_iterator = trange(num_train_epochs, desc='Epoch')\n",
    "        loss_queue = deque(maxlen=10)\n",
    "        for _ in epoch_iterator:\n",
    "            batch_iterator = tqdm(train_loader, desc='Batch')\n",
    "            for step, batch in enumerate(batch_iterator):\n",
    "\n",
    "                model.train()\n",
    "                inputs = {\n",
    "                    'input_ids': batch['input_ids'],\n",
    "                    'attention_mask': batch['input_mask'],\n",
    "                    'labels': batch['label_ids'],\n",
    "                    'token_type_ids': batch['segment_ids']\n",
    "                }\n",
    "                if model_type == 'distilbert':\n",
    "                    inputs.pop('token_type_ids')\n",
    "\n",
    "                model_output = model(**inputs)\n",
    "                loss = model_output[0]\n",
    "                loss.backward()\n",
    "                tr_loss += loss.item()\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "                model.zero_grad()\n",
    "                global_step += 1\n",
    "                if global_step % logging_steps == 0:\n",
    "                    last_loss = (tr_loss - logging_loss) / logging_steps\n",
    "                    loss_queue.append(last_loss)\n",
    "                    if train_logs:\n",
    "                        tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)\n",
    "                        tb_writer.add_scalar('loss', last_loss, global_step)\n",
    "                    logging_loss = tr_loss\n",
    "\n",
    "            # slope-based early stopping\n",
    "            if len(loss_queue) == loss_queue.maxlen:\n",
    "                slope = calc_slope(loss_queue)\n",
    "                if train_logs:\n",
    "                    tb_writer.add_scalar('slope', slope, global_step)\n",
    "                if abs(slope) < 1e-2:\n",
    "                    break\n",
    "\n",
    "        if train_logs:\n",
    "            tb_writer.close()\n",
    "\n",
    "        model_to_save = model.module if hasattr(model, \"module\") else model  # Take care of distributed/parallel training\n",
    "        model_to_save.save_pretrained(workdir)\n",
    "        tokenizer.save_pretrained(workdir)\n",
    "        label_map = {i: t for t, i in train_set.tag_idx_map.items()}\n",
    "\n",
    "        return {\n",
    "            'model_path': workdir,\n",
    "            'batch_size': batch_size,\n",
    "            'pad_token_label_id': pad_token_label_id,\n",
    "            'dataset_params_dict': train_set.get_params_dict(),\n",
    "            'model_type': model_type,\n",
    "            'pretrained_model': pretrained_model,\n",
    "            'label_map': label_map\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f98b3d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a1f67ed",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a1dc286",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129871d1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e827984",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d95bd88",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.10.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": false,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": false,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}