<jupyter_start><jupyter_text>Token Tagging and Classification> AdaptiveModels for using Transformers and Flair for token tagging and classification<jupyter_code>#hide
from fastcore.test import test_eq, test_close
from nbverbose.showdoc import *
#export
import logging
from typing import List, Dict, Union
from collections import defaultdict, OrderedDict

import numpy as np

import torch
from torch.utils.data import TensorDataset

from flair.data import Sentence
from flair.models import SequenceTagger
from flair.tokenization import SegtokSentenceSplitter

from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    PreTrainedTokenizer,
    PreTrainedModel,
    BertForSequenceClassification,
    XLNetForSequenceClassification,
    AlbertForSequenceClassification,
)

from adaptnlp.result import DetailLevel

from adaptnlp.model import AdaptiveModel, DataLoader
from adaptnlp.model_hub import HFModelResult, FlairModelResult, FlairModelHub, HFModelHub

from fastai.torch_core import to_detach, apply, to_device

from fastcore.basics import Self, risinstance, filter_ex, listify
from fastcore.xtras import Path
#export
logger = logging.getLogger(__name__)
#export
class TokenClassificationResult:
    "A result class for Token Tagging tasks"
    def __init__(
        self, 
        inputs:str, # Original text inputs
        tokenized_inputs:torch.tensor, # Tokenized inputs
        tagged_entities:dict
    ):
        self.inputs = inputs
        self.tokenized_inputs = tokenized_inputs
        self.tagged_entities = tagged_entities
    
    def to_dict(
        self, 
        detail_level:DetailLevel=DetailLevel.Low # A level of detail to return
    ) -> dict:
        "Convert `self` to a dictionary"
        o = OrderedDict()
        o.update({
            'tags': self.tagged_entities
        })
        if detail_level == 'medium' or detail_level == 'high':
            o.update({
                'inputs':self.inputs,
                'tokenized_inputs':self.tokenized_inputs
            })
        if detail_level != 'high':
            try:
                for tag in o['tags'][0]:
                    tag.pop('offsets')
            except:
                return o
        return o
show_doc(TokenClassificationResult.to_dict)
#export
class TransformersTokenTagger(AdaptiveModel):
    "Adaptive model for Transformer's Token Tagger Model"
    def __init__(
        self, 
        tokenizer: PreTrainedTokenizer, # A tokenizer object from Huggingface's transformers (TODO) and tokenizers 
        model: PreTrainedModel # A transformers token tagger model
    ):
        # Load up model and tokenizer
        self.tokenizer = tokenizer
        super().__init__()

        # Sets the internal model
        self.set_model(model)

    @classmethod
    def load(
        cls, 
        model_name_or_path: str # A key string of one of Transformer's pre-trained Token Tagger Model or a `HFModelResult`
    ) -> AdaptiveModel:
        "Class method for loading and constructing this tagger"
        if isinstance(model_name_or_path, HFModelResult): model_name_or_path = model_name_or_path.name
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        model = AutoModelForTokenClassification.from_pretrained(model_name_or_path)
        tagger = cls(tokenizer, model)
        return tagger

    def predict(
        self,
        text: Union[List[str], str], # Sentences to run inference on
        mini_batch_size: int = 32, # Mini batch size
        grouped_entities: bool = True, # Return whole entity span strings
        detail_level:DetailLevel = DetailLevel.Low, # A level of detail to return
        **kwargs, # Optional arguments for the Transformers tagger
    ) -> List[List[Dict]]: # Returns a list of lists of tagged entities
        "Predict method for running inference using the pre-trained token tagger model"
        if isinstance(text, str):
            text = [text]
        results: List[Dict] = []

        dataset = self._tokenize(text)
        dl = DataLoader(dataset, batch_size=mini_batch_size)

        logger.info(f'Running prediction on {len(dataset)} text sequences')
        logger.info(f'Batch size = {mini_batch_size}')

        outputs,_ = super().get_preds(dl=dl)

        inputs = apply(to_device, [b for b in dl], device='cpu')
        inputs = torch.cat(*inputs)
        inputs = apply(Self.numpy(), inputs)

        outputs = torch.cat([o['logits'] for o in outputs])
        outputs = apply(to_detach, outputs, cpu=True)
        outputs = apply(Self.numpy(), outputs)

        # Iterate through batch for tagged token predictions
        for idx, pred in enumerate(outputs):
            entities = pred
            input_ids = inputs[idx]
            tagged_entities = self._generate_tagged_entities(
                entities=entities,
                input_ids=input_ids,
                grouped_entities=grouped_entities
            )
            results += tagged_entities
            
        results = TokenClassificationResult(text, inputs, results)

        return results.to_dict(detail_level) if detail_level is not None else detail_level

    def _tokenize(
        self, sentences: Union[List[Sentence], Sentence, List[str], str]
    ) -> TensorDataset:
        "Batch tokenizes text and produces a `TensorDataset` with them"

        tokenized_text = self.tokenizer.batch_encode_plus(
            sentences,
            return_tensors="pt",
            max_length=None,
        )

        # Bart, XLM, DistilBERT, RoBERTa, and XLM-RoBERTa don't use token_type_ids
        if isinstance(
            self.model,
            (
                BertForSequenceClassification,
                XLNetForSequenceClassification,
                AlbertForSequenceClassification,
            ),
        ):
            dataset = TensorDataset(
                tokenized_text["input_ids"],
                tokenized_text["attention_mask"],
                tokenized_text["token_type_ids"],
            )
        else:
            dataset = TensorDataset(
                tokenized_text["input_ids"], tokenized_text["attention_mask"]
            )

        return dataset

    # `_group_entites` and `_generate_tagged_entities` modified from pipeline code snippet from Transformers
    def _group_entities(
        self, entities: List[dict], idx_start: int, idx_end: int
    ) -> Dict:
        """Returns grouped entities"""
        # Get the last entity in the entity group
        entity = entities[-1]["entity"]
        scores = np.mean([entity["score"] for entity in entities])
        tokens = [entity["word"] for entity in entities]

        entity_group = {
            "entity": entity,
            "score": np.mean(scores),
            "word": self.tokenizer.convert_tokens_to_string(tokens),
            "offsets": (idx_start, idx_end),
        }
        return entity_group

    def _generate_tagged_entities(
        self, entities: np.ndarray, input_ids: np.ndarray, grouped_entities: bool = True
    ) -> List[Dict]:
        """Generate full list of entities given tagged token predictions and input_ids"""

        score = np.exp(entities) / np.exp(entities).sum(-1, keepdims=True)
        labels_idx = score.argmax(axis=-1)

        answers = []
        entities = []
        entity_groups = []
        entity_group_disagg = []
        # Filter to labels not in `self.ignore_labels`
        filtered_labels_idx = [
            (idx, label_idx)
            for idx, label_idx in enumerate(labels_idx)
            if self.model.config.id2label[label_idx] not in ["O"]
        ]

        for idx, label_idx in filtered_labels_idx:
            # print(tokenizer.convert_ids_to_tokens(int(input_ids[idx])))
            entity = {
                "word": self.tokenizer.convert_ids_to_tokens(int(input_ids[idx])),
                "score": score[idx][label_idx].item(),
                "entity": self.model.config.id2label[label_idx],
                "index": idx,
            }
            last_idx, _ = filtered_labels_idx[-1]
            if grouped_entities:
                if not entity_group_disagg:
                    entity_group_disagg += [entity]
                    if idx == last_idx:
                        entity_groups += [
                            self._group_entities(
                                entity_group_disagg, idx - len(entity_group_disagg), idx
                            )
                        ]
                    continue

                # If the current entity is similar and adjacent to the previous entity, append it to the disaggregated entity group
                if (
                    entity["entity"] == entity_group_disagg[-1]["entity"]
                    and entity["index"] == entity_group_disagg[-1]["index"] + 1
                ):
                    entity_group_disagg += [entity]
                    # Group the entities at the last entity
                    if idx == last_idx:
                        entity_groups += [
                            self._group_entities(
                                entity_group_disagg, idx - len(entity_group_disagg), idx
                            )
                        ]
                # If the current entity is different from the previous entity, aggregate the disaggregated entity group
                else:
                    entity_groups += [
                        self._group_entities(
                            entity_group_disagg,
                            entity_group_disagg[-1]["index"] - len(entity_group_disagg),
                            entity_group_disagg[-1]["index"],
                        )
                    ]
                    entity_group_disagg = [entity]

            entities += [entity]

        # Append
        if grouped_entities:
            answers += [entity_groups]
        else:
            answers += [entities]

        return answers
#hide
tagger = TransformersTokenTagger.load("dbmdz/bert-large-cased-finetuned-conll03-english")
pred = tagger.predict(text='Novetta Solutions is the best. <NAME> used to be employed at Novetta Solutions. The Wright brothers loved to visit the JBF headquarters, and they would have a chat with Albert.', mini_batch_size=32)
baseline = [[{'entity': 'I-ORG',
   'score': 0.998292068640391,
   'word': 'Novetta Solutions'},
  {'entity': 'I-PER',
   'score': 0.9985582232475281,
   'word': '<NAME>'},
  {'entity': 'I-ORG',
   'score': 0.9970489343007406,
   'word': 'Novetta Solutions'},
  {'entity': 'I-PER',
   'score': 0.9961656928062439,
   'word': 'Wright'},
  {'entity': 'I-ORG',
   'score': 0.9933501183986664,
   'word': 'JBF'}]]

for base, p in zip(baseline, pred['tags']):
    for base_items, p_items in zip(base, p):
        test_eq(base_items['entity'], p_items['entity'])
        test_close(base_items['score'], p_items['score'], 1e-3)
        test_eq(base_items['word'], p_items['word'])
show_doc(TransformersTokenTagger.load)
show_doc(TransformersTokenTagger.predict)
#export
class FlairTokenTagger(AdaptiveModel):
    "Adaptive Model for Flair's Token Tagger"

    def __init__(
        self, 
        model_name_or_path: str # A key string of one of Flair's pre-trained Token tagger Model, [link](https://huggingface.co/models?filter=flair)
    ):
        self.tagger = SequenceTagger.load(model_name_or_path)
        self.splitter = SegtokSentenceSplitter()

    @classmethod
    def load(
        cls, 
        model_name_or_path: str # A key string of one of Flair's pre-trained Token tagger Model, [link](https://huggingface.co/models?filter=flair)
    ) -> AdaptiveModel:
        "Class method for loading a constructing this tagger"
        tagger = cls(model_name_or_path)
        return tagger
    
    def _split_articles(self, tagged_articles):
        "Takes merged NER results and splits them back out into their original articles"
        all_payloads = []
        for article in tagged_articles:
            payload = [sentence.to_dict(tag_type=self.tagger.tag_type) for sentence in article]
            curr_pos = 0
            new_payload = {}
            new_payload['entities'] = []
            new_payload['labels'] = []
            text = ''
            for i, pay in enumerate(payload):
                entities = pay['entities']
                for e in entities:
                    labels = e['labels']
                    d = labels[0].to_dict()
                    e['value'] = d['value']
                    e['confidence'] = d['confidence']
                
                if i > 0:
                    for e in entities:
                        # Get back the original positions
                        e['start_pos'] += curr_length
                        e['end_pos'] += curr_length
                
                text += pay['text'] + ' '
                curr_length = len(text)
                new_payload['entities'] += entities
            new_payload['text'] = text
            all_payloads.append(new_payload)
        return all_payloads
    
    def decode_articles(
        self,
        original_text: Union[List[str], str], # Original text inference was run on
        sentences: List[Sentence] # Sentences inference was run on
    ) -> List[dict]:
        "Decodes `text` back into articles and extracts entities"
        if not isinstance(original_text, list):
            original_text = [original_text]
        for i,sentence in enumerate(original_text):
            if isinstance(sentence, Sentence):
                text[i] = sentence.to_original_text()
        tagged_articles = [[] for _ in range(len(original_text))]
        for sentence in sentences:
            for i, article in enumerate(original_text):
                if sentence.to_original_text() in article:
                    tagged_articles[i].append(sentence)
        return self._split_articles(tagged_articles)
            
    def predict(
        self,
        text: Union[List[Sentence], Sentence, List[str], str], # Sentences to run inference on
        mini_batch_size: int = 32, # Mini batch size
        raw:bool = False, # Whether to return a list of raw Sentences
        **kwargs, # Optional arguments for the Flair tagger
    ) -> List[Sentence]: # A list of predicted sentences
        "Predict method for running inference using the pre-trained token tagger model"
        if not isinstance(text, list):
            text = [text]
        for i,sentence in enumerate(text):
            if isinstance(sentence, Sentence):
                text[i] = sentence.to_original_text()
        tagged_articles = [[] for _ in range(len(text))]
        t = '\n'.join(text)
        sentences = self.splitter.split(t)
        self.tagger.predict(sentences, mini_batch_size=mini_batch_size, **kwargs)
        
        if not raw:
            return self.decode_articles(text, sentences)
        return sentences
#hide
tagger = FlairTokenTagger.load("flair/chunk-english-fast")
text = ['Novetta Solutions is the best.', '<NAME> used to be employed at Novetta Solutions.', 'The Wright brothers loved to visit the JBF headquarters, and they would have a chat with Albert.']
preds = tagger.predict(text=text, mini_batch_size=32)
preds[0]['entities'][0]['text'] = 'Novetta Solutions'
preds[0]['entities'][0]['value'] = 'NP'

# Make sure we calculated the positioning right
random_entity = preds[1]['entities'][3]
entity_name = random_entity['text']
start_pos, end_pos = random_entity['start_pos'], random_entity['end_pos']
test_eq(entity_name, text[1][start_pos:end_pos])
preds[0]['labels']
show_doc(FlairTokenTagger.load)
show_doc(FlairTokenTagger.predict)
#export
class EasyTokenTagger:
    "Token level classification models"

    def __init__(self):
        self.token_taggers: Dict[AdaptiveModel] = defaultdict(bool)

    def tag_text(
        self,
        text: Union[List[Sentence], Sentence, List[str], str], # Text input, it can be a string or any of Flair's `Sentence` input formats
        model_name_or_path: Union[str, FlairModelResult, HFModelResult] = "ner-ontonotes", # The hosted model name key or model path
        mini_batch_size: int = 32, # The mini batch size for running inference
        detail_level:DetailLevel = DetailLevel.Low, # The level fo detail to return on a TransformerTagger
        **kwargs, # Keyword arguments for Flair's `SequenceTagger.predict()` method
    ) -> List[Sentence]: # A list of Flair's `Sentence`'s
        "Tags tokens with labels the token classification models have been trained on"
        # Load Sequence Tagger Model and Pytorch Module into tagger dict
        name = getattr(model_name_or_path, 'name', model_name_or_path)
        if not self.token_taggers[name]:
            """
            self.token_taggers[model_name_or_path] = SequenceTagger.load(
                model_name_or_path
            )
            """
            if risinstance([FlairModelResult, HFModelResult], model_name_or_path):
                try:
                    self.token_taggers[name] = FlairTokenTagger.load(name)
                except:
                    self.token_taggers[name] = TransformersTokenTagger.load(name)
            elif risinstance([str, Path], model_name_or_path) and (Path(model_name_or_path).exists() and Path(model_name_or_path).is_dir()):
                # Load in previously existing model
                try:
                    self.token_taggers[name] = FlairTokenTagger.load(name)
                except:
                    self.token_taggers[name] = TransformersTokenTagger.load(name)
            else:
                _flair_hub = FlairModelHub()
                _hf_hub = HFModelHub()
                res = _flair_hub.search_model_by_name(name, user_uploaded=True)
                if len(res) < 1:
                    # No models found
                    res = _hf_hub.search_model_by_name(name, user_uploaded=True)
                    if len(res) < 1:
                        logger.info("Not a valid `model_name_or_path` param")
                        return [Sentence('')]
                    else:
                        res[0].name.replace('flairNLP', 'flair')
                        self.token_taggers[res[0].name] = TransformersTokenTagger.load(res[0].name)
                        name = res[0].name

                else:
                    name = res[0].name.replace('flairNLP/', '')
                    self.token_taggers[name] = FlairTokenTagger.load(name) # Returning the first should always be the non-fast option
                    
        tagger = self.token_taggers[name]
        if isinstance(tagger, TransformersTokenTagger):
            return tagger.predict(
                text=text,
                mini_batch_size=mini_batch_size,
                **kwargs
            )
        else:
            return tagger.predict(
                text=text,
                mini_batch_size=mini_batch_size,
                raw=False, 
                **kwargs,
            )

    def tag_all(
        self,
        text: Union[List[Sentence], Sentence, List[str], str], # Text input, it can be a string or any of Flair's `Sentence` input formats
        mini_batch_size: int = 32, # The mini batch size for running inference
        detail_level:DetailLevel = DetailLevel.Low, # The level of detail for a TransformersTagger to return
        **kwargs, # Keyword arguments for Flair's `SequenceTagger.predict()` method
    ) -> List[dict]: # A dictionary of Token Tagging results
        "Tags tokens with all labels from all token classification models"
        if len(self.token_taggers) == 0:
            print("No token classification models loaded...")
            return Sentence()
        if not isinstance(text, list):
            text = [text]
        results = []
        all_models = [o for o in list(self.token_taggers.keys()) if self.token_taggers[o]]
        for tagger_name in all_models:
            results.append(
                self.tag_text(
                    text,
                    model_name_or_path=tagger_name,
                    mini_batch_size=mini_batch_size,
                    **kwargs,
                )
            )
        tag_results = [{'entities':[], 'labels':[], 'text':None} for _ in range(len(results[0]))]
        for tag in results:
            for i, inp in enumerate(tag):
                if not tag_results[i]['text']:
                    tag_results[i]['text'] = inp['text']
                if 'entities' in inp:
                    if len(inp['entities']) > 0:
                        tag_results[i]['entities'] += inp['entities']
                if 'labels' in inp:
                    if len(inp['labels']) > 0:
                        tag_results[i]['labels'] += inp['labels']

        all_results = []
        for result in tag_results:
            res = []
            for entity in result['entities']:
                unique_entities = []
                new_entity = {
                    'text':entity['text'],
                    'start_pos':entity['start_pos'],
                    'end_pos':entity['end_pos']}

                if new_entity not in unique_entities:
                    unique_entities.append(new_entity)
                if len(unique_entities) > 0:
                    merged = []
                    for entity in unique_entities:
                        instances = filter_ex(
                                result['entities'],
                                lambda x: (
                                    x['text'] == entity['text'] and \
                                    x['start_pos'] == entity['start_pos'] and \
                                    x['end_pos'] == entity['end_pos']
                                )
                            )
                        attrs = ('labels','value','confidence')
                        for attr in attrs:
                            entity[attr] = []
                        for instance in instances:
                            for attr in attrs:
                                if attr in instance.keys():
                                    entity[attr] += listify(instance[attr])
                    merged = sorted(unique_entities, key=lambda x: x['start_pos'])
                    res += merged
            all_results.append(res)
        return all_results
tagger = EasyTokenTagger()
sentences = ["Jack walked through the park on a Sunday.", "Sunday was a nice and breezy afternoon.", "Jack was going to meet Jill for dinner."]
res = tagger.tag_text(text=sentences, model_name_or_path='flair/ner-english-ontonotes')
tagger.tag_text(text=sentences, model_name_or_path="flair/pos")
results = tagger.tag_all(sentences)
#hide
test_eq(len(results), 3)
test_eq(results[0][0]['text'], 'Jack')
test_eq(results[1][0]['text'], 'Sunday')
test_eq(len(results[1][0]['labels']), 2)
#hide
t = '''Novetta Solutions is the best. <NAME> used to be employed at Novetta Solutions. 
The Wright brothers loved to visit the JBF headquarters, and they would have a chat with Albert.'''
results = tagger.tag_all(t)
show_doc(EasyTokenTagger.tag_text)
show_doc(EasyTokenTagger.tag_all)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09358196010407632, 'fraction_numerical': 0.006201214223764094, 'mean_word_length': 2.0455625990491284, 'pattern_counts': {'":': 8, '<': 11, '<?xml version=': 0, '>': 25, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '158681', 'license': 'Apache-2.0', 'n_tokens_mistral': 6545, 'n_tokens_neox': 6225, 'n_words': 1621, 'path': 'nbs/05_token_classification.ipynb', 'repo_name': 'mfredriksz/adaptnlp', 'size': '23096'}