# -*- coding: utf-8 -*-
"""Conll2003 NER Classification using DistilBERT HF Pytorch Script.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1K55Dlto68jOmnijtcN3rXTMLaakVGY1a

### 0. Install and load library
"""

#Restart kernel after installation: Runtime -> Restart runtime

#!pip install -U sentencepiece datasets
#!pip install -U transformers sentencepiece datasets

# Commented out IPython magic to ensure Python compatibility.
!git clone https://github.com/huggingface/transformers
# %cd transformers
!python setup.py install

# Commented out IPython magic to ensure Python compatibility.
# %cd "/content/transformers/examples/pytorch/token-classification"
!pip install -r requirements.txt

"""### 1. Fine-tuning DistillBERT model on CONLL2003 NER dataset."""

!python /content/transformers/examples/pytorch/token-classification/run_ner.py \
  --model_name_or_path distilbert-base-cased \
  --label_all_tokens True \
  --return_entity_level_metrics True \
  --dataset_name conll2003 \
  --output_dir /tmp/distilbert-base-cased-finetuned-conll03-english \
  --do_train \
  --do_eval

"""[INFO|tokenization_utils_base.py:1954] 2021-07-09 08:05:47,971 >> Special tokens file saved in /tmp/distilbert-base-cased-finetuned-conll03-english/special_tokens_map.json
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     0.0706
  train_runtime            = 0:05:10.33
  train_samples            =      14041
  train_samples_per_second =    135.734
  train_steps_per_second   =     16.975
07/09/2021 08:05:48 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:521] 2021-07-09 08:05:48,010 >> The following columns in the evaluation set  don't have a corresponding argument in `DistilBertForTokenClassification.forward` and have been ignored: chunk_tags, tokens, pos_tags, id, ner_tags.
[INFO|trainer.py:2159] 2021-07-09 08:05:48,030 >> ***** Running Evaluation *****
[INFO|trainer.py:2161] 2021-07-09 08:05:48,031 >>   Num examples = 3250
[INFO|trainer.py:2164] 2021-07-09 08:05:48,031 >>   Batch size = 8
100% 407/407 [00:06<00:00, 60.22it/s]
***** eval metrics *****
  epoch                   =        3.0
  eval_LOC_f1             =     0.9592
  eval_LOC_number         =       3635
  eval_LOC_precision      =     0.9639
  eval_LOC_recall         =     0.9546
  eval_MISC_f1            =     0.8856
  eval_MISC_number        =       1480
  eval_MISC_precision     =     0.8896
  eval_MISC_recall        =     0.8818
  eval_ORG_f1             =     0.9053
  eval_ORG_number         =       2702
  eval_ORG_precision      =     0.8931
  eval_ORG_recall         =     0.9178
  eval_PER_f1             =       0.96
  eval_PER_number         =       3329
  eval_PER_precision      =     0.9584
  eval_PER_recall         =     0.9616
  eval_loss               =     0.0782
  eval_overall_accuracy   =     0.9838
  eval_overall_f1         =     0.9365
  eval_overall_precision  =     0.9349
  eval_overall_recall     =     0.9381
  eval_runtime            = 0:00:06.78
  eval_samples            =       3250
  eval_samples_per_second =    479.029
  eval_steps_per_second   =     59.989

### 2. Inference using pipeline
"""

!ls tmp/distilbert-base-cased-finetuned-conll03-english/

# Update NER label2id and id2label in model config file
"""
  "id2label": {
    "0": "O",
    "1": "B-PER",
    "2": "I-PER",
    "3": "B-ORG",
    "4": "I-ORG",
    "5": "B-LOC",
    "6": "I-LOC",
    "7": "B-MISC",
    "8": "I-MISC"
  },


  "label2id": {
    "B-LOC": 5,
    "B-MISC": 7,
    "B-ORG": 3,
    "B-PER": 1,
    "I-MISC": 8,
    "I-ORG": 4,
    "I-PER": 2,
    "O": 0
  },

"""

from transformers import pipeline
distilbert_ner = pipeline('ner', model="/tmp/distilbert-base-cased-finetuned-conll03-english/", tokenizer="distilbert-base-cased", aggregation_strategy='first')

article = """
KUALA LUMPUR - Malaysian Prime Minister Muhyiddin Yassin's party said on Thursday (July 8) that his government would continue to function despite Umno withdrawing its backing. 
Amid uncertainty over whether Tan Sri Muhyiddin continues to command majority support without Umno, the largest party in the Perikatan Nasional (PN) ruling pact, 
his Parti Pribumi Bersatu Malaysia said Umno's decision "had no effect on the workings of government"."""


results = distilbert_ner(article)

print("Predicted:")
for tag in results:
    print(f"{tag['entity_group']:<5} as {tag['word']}")


"""
Output
------

Predicted:
ORG   as KUALA LUMPUR
MISC  as Malaysian
PER   as Muhyiddin Yassin
PER   as Umno
PER   as Tan Sri Muhyiddin
PER   as Umno
ORG   as Perikatan Nasional
ORG   as PN
ORG   as Parti Pribumi
PER   as Bersatu
ORG   as Malaysia
PER   as Umno

"""