# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="aE4o0aOazcFg"
# ---
# Author: **`<NAME>`**
#
# Year: **`2021`**
#
# Date: **`2021-09-15`**
#
# Language: **`Python`**
#
# Libray: **`Pytorch`**
#
# Topic: **`Named Entity Recognition (NER)`**
#
# Main: **`Natural Language Processing (NLP)`**
#
# ---

# + [markdown] id="Fv1XmS5E0ICM"
# ### Named Entity Recognition (NER) using Bi-Directional LSTM (Bi-LSTM)
#
# In this series of notebookswe are going to have a look at an intresting topic in Natural Language Processing (NLP) known as Named Entity Recognition (NER).
#
# ### Some of the uses
# Named Entity Recognition can automatically scan entire articles and reveal which are the major people, organizations, and places discussed in them. Knowing the relevant tags for each article help in automatically categorizing the articles in defined hierarchies and enable smooth content discover.
#
# ### Imports

# + colab={"base_uri": "https://localhost:8080/", "height": 35} id="2qMKDufx0H5d" outputId="8d8e8a8f-994f-405a-baba-4279607965ed"
import time, os, torch, random, json

from torch import nn
from torch.nn import functional as F

from torchtext.legacy import data, datasets

import numpy as np

torch.__version__


# + [markdown] id="6JVcK4JL10V8"
# ### Seeds and Device

# + colab={"base_uri": "https://localhost:8080/"} id="8HJ06MWy2CVl" outputId="a014f7ce-a826-4d53-9575-1cb22e08c7f2"
SEED = 42

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

torch.backends.cudnn.deterministic = True

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

# + [markdown] id="B0PsQB171aww"
# ### Data.
#
# The data that we will be working with was found [here](https://github.com/yoseflaw/nerindo/tree/master/input). What i've done is to download the the three files and load uploaded them on my google drive so that it can be eaisly loaded in google colab. 
#
# ### Explaining the data.
# The data files are tab deliminated, `tsv` where columns are separated by a `\t`.
#
# To accommodate multi-word entities, the tags follow what is called **`BILOU`** format: **B**eginning, **I**nside, **L**ast, **O**uter, **U**nit. This is indicated by the character preceding the dash symbol. The type of the named entity is shown by the remaining part. For instance, *Universtas Gadjah Mada* is an `ORGANIZATION` and *Arie Sudjito* is a `PERSON`. 
#
# Here is what a single file may look like:
# ```
# Pengamat	O
# politik	O
# dari	O
# Universitas	B-ORGANIZATION
# Gadjah	I-ORGANIZATION
# Mada	L-ORGANIZATION
# ,	O
# Arie	B-PERSON
# Sudjito	L-PERSON
# ,	O
# menilai	O
# ,	O
# keinginan	O
# Ketua	O
# Umum	O
# Partai	B-ORGANIZATION
# Golkar	L-ORGANIZATION
# Aburizal	B-PERSON
# Bakrie	L-PERSON
# untuk	O
# maju	O
# kembali	O
# sebagai	O
# ketua	O
# umum	O
# merupakan	O
# pemaksaan	O
# kehendak	O
# .	O
#
# Menurut	O
# dia	O
# ,	O
# ada	O
# kesan	O
# bahwa	O
# Aburizal	U-PERSON
# menggunakan	O
# segala	O
# cara	O
# untuk	O
# memuluskan	O
# jalannya	O
# kembali	O
# menduduki	O
# Golkar	U-ORGANIZATION
# 1	O
# .	O
# ```

# + colab={"base_uri": "https://localhost:8080/"} id="gSkSmHpm46tY" outputId="3b6348a4-15ea-4ec6-9b8e-4b335d5260dc"
from google.colab import drive
drive.mount("/content/drive")

# + [markdown] id="bwaHCbYt_P9q"
# ### Root Path

# + colab={"base_uri": "https://localhost:8080/"} id="G8dM-4BT4r-f" outputId="2d0dab17-b365-48ee-adab-909fb78e68bd"
root = '/content/drive/My Drive/NLP Data/ner/data'
os.path.exists(root)

# + [markdown] id="2C8dEkkj-7UM"
# ### Fields

# + id="bdT24RRa0H2U"
WORD = data.Field(lower=True)
TAG = data.Field(unk_token=None)

# + id="64-oZILQ0Hzu"
fields = (
 ("word", WORD),
 ("tag", TAG)
)

# + id="KDzIWQW00Hwk"
train_dataset, valid_dataset, test_dataset = datasets.SequenceTaggingDataset.splits(
    path=root,
    train="train.tsv",
    validation="val.tsv",
    test= "test.tsv",
    fields=fields
)

# + colab={"base_uri": "https://localhost:8080/"} id="GhKm61P_0Ht8" outputId="216bb51f-dcca-4d31-9114-d53278e7ee71"
print(vars(train_dataset.examples[67]))

# + [markdown] id="hlDA2-X2FCDT"
# ### Building vocabulary

# + id="3-WtxIFm0Hqm"
WORD.build_vocab(train_dataset, min_freq=3)
TAG.build_vocab(train_dataset)

# + colab={"base_uri": "https://localhost:8080/"} id="vPbfukgdNfsF" outputId="cd077bd6-abd4-4256-e860-a6dbc9c75c53"
TAG.vocab.stoi

# + [markdown] id="0CFXDgO2FEiE"
# ### Creating iterators

# + id="P2Xb59QI0Hn7"
BATCH_SIZE = 128
train_iter, valid_iter, test_iter = data.BucketIterator.splits(
    (train_dataset,valid_dataset, test_dataset),
    device = device,
    batch_size = BATCH_SIZE,
    sort_key = lambda x: x.word,
)

# + [markdown] id="Lk1GuVZ8NyqD"
# ### Building the vocabulary
#
# Every word that apears less that 2 times in the corpus will be converted to unknown token.

# + id="5ezNs0apMxO1"
WORD.build_vocab(train_dataset, min_freq=2)
TAG.build_vocab(train_dataset)

# + [markdown] id="dAPbsNJmFeki"
# ### Counting examples

# + colab={"base_uri": "https://localhost:8080/"} id="Oogcp7Ny0HlE" outputId="4c591f41-9973-4d54-cf83-ba2ad113ec0e"
print(f"Train set: {len(train_dataset)} sentences")
print(f"Val set: {len(valid_dataset)} sentences")
print(f"Test set: {len(test_dataset)} sentences")


# + [markdown] id="LtqOIJeRNvQj"
# ### Model
#

# + id="T1W6UKqtOOa6"
class BiLSTM(nn.Module):
  def __init__(self, input_dim,
               embedding_dim,
               hidden_dim,
               output_dim,
               word_pad_idx,
               dropout=.5,
               n_layers=2,
               bidirectional = True
               ):
    super(BiLSTM, self).__init__()
    self.embedding = nn.Embedding(input_dim, embedding_dim,)
    self.lstm = nn.LSTM(
        input_size=embedding_dim,
        hidden_size=hidden_dim,
        num_layers=n_layers,
        bidirectional=bidirectional,
        dropout=dropout if n_layers > 1 else 0
    )
    self.fc = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim,
                        output_dim)
    self.dropout= nn.Dropout(.5)

  def forward(self, sentence):
    # sentence = [sentence length, batch size]
    embedded= self.dropout(self.embedding(sentence)) # [sentence length, batch size, embedding dim]
    out, _ = self.lstm(embedded) # out = [sentence length, batch size, hidden dim * 2]
    out = self.fc(self.dropout(out)) # [sentence length, batch size, output dim]
    return out



# + colab={"base_uri": "https://localhost:8080/"} id="bT56eZLyQEGs" outputId="fdd6d456-e841-4880-aed2-70c05bf35107"
INPUT_DIM=len(WORD.vocab)
EMBEDDING_DIM=300
HIDDEN_DIM=256
OUTPUT_DIM=len(TAG.vocab)
N_LAYERS=2
DROPOUT=0.1
WORD_PAD_IDX = WORD.vocab.stoi[WORD.pad_token]
TAG_PAD_IDX = TAG.vocab.stoi[TAG.pad_token]
BIDIRECTIONAL = True


model = BiLSTM(
    INPUT_DIM,
    EMBEDDING_DIM,
    HIDDEN_DIM,
    OUTPUT_DIM,
    WORD_PAD_IDX,
    DROPOUT,
    N_LAYERS,
    BIDIRECTIONAL
).to(device)
model


# + [markdown] id="Jt-MYuIAXZJR"
# ### Initializing model weights

# + colab={"base_uri": "https://localhost:8080/"} id="CNafr7nIXlDY" outputId="9db7c802-f966-42f4-e5a7-a157e078329e"
def init_weights(m):
  for name, param in m.named_parameters():
    nn.init.normal_(param.data, mean=0, std=0.1)

model.apply(init_weights)

# + [markdown] id="4GXC_SeQX2b1"
# ### initialize embedding for padding as zero

# + colab={"base_uri": "https://localhost:8080/"} id="w3AsggGuX6X2" outputId="e12af364-d502-4335-d1c8-eb2504e1d1ab"
model.embedding.weight.data[WORD_PAD_IDX] = torch.zeros(EMBEDDING_DIM)
model.embedding.weight.data


# + [markdown] id="BGZIhMCqYNNm"
# ### Counting model parameters.

# + colab={"base_uri": "https://localhost:8080/"} id="Cmk3nwuk0HfU" outputId="71cfcbcd-ffb1-4a52-af81-609c45474104"
def count_trainable_params(model):
  return sum(p.numel() for p in model.parameters()), sum(p.numel() for p in model.parameters() if p.requires_grad)

n_params, trainable_params = count_trainable_params(model)
print(f"Total number of paramaters: {n_params:,}\nTotal tainable parameters: {trainable_params:,}")


# + [markdown] id="Q7gqiihTYtbt"
# ### Training the model.

# + id="Oo1cpbd2GoCa"
def accuracy(preds, y):
  max_preds = preds.argmax(dim=1, keepdim=True).to(device)  # get the index of the max probability
  non_pad_elements = (y != TAG_PAD_IDX).nonzero().to(device)   # prepare masking for paddings
  correct = max_preds[non_pad_elements].squeeze(1).eq(y[non_pad_elements])
  return correct.sum() / torch.FloatTensor([y[non_pad_elements].shape[0]]).to(device) 



# + id="qX5v0EGJGsd9"
def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs


# + [markdown] id="U8VDeJCoZhow"
# ### Optimizer and Criterion

# + id="rpGCWQDdHPFa"
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss().to(device)


# + [markdown] id="T4_zzdE-Zmk0"
# ### Train and evalutation functions

# + id="8psaVdIUHBva"
def train(model, iterator, optimizer, criterion):
  epoch_loss = 0
  epoch_acc = 0
  model.train()
  for batch in iterator:
    # text = [sent len, batch size]
    text = batch.word
    # tags = [sent len, batch size]
    true_tags = batch.tag
    optimizer.zero_grad()
    pred_tags = model(text)
    # to calculate the loss and accuracy, we flatten both prediction and true tags
    # flatten pred_tags to [sent len, batch size, output dim]
    pred_tags = pred_tags.view(-1, pred_tags.shape[-1])
    # flatten true_tags to [sent len * batch size]
    true_tags = true_tags.view(-1)
    batch_loss = criterion(pred_tags, true_tags)
    batch_acc = accuracy(pred_tags, true_tags)
    batch_loss.backward()
    optimizer.step()
    epoch_loss += batch_loss.item()
    epoch_acc += batch_acc.item()
  return epoch_loss / len(iterator), epoch_acc / len(iterator)

def evaluate(model, iterator, criterion):
    epoch_loss = 0
    epoch_acc = 0
    model.eval()
    with torch.no_grad():
        # similar to epoch() but model is in evaluation mode and no backprop
        for batch in iterator:
            text = batch.word
            true_tags = batch.tag
            pred_tags = model(text)
            pred_tags = pred_tags.view(-1, pred_tags.shape[-1])
            true_tags = true_tags.view(-1)
            batch_loss = criterion(pred_tags, true_tags)
            batch_acc = accuracy(pred_tags, true_tags)
            epoch_loss += batch_loss.item()
            epoch_acc += batch_acc.item()
    return epoch_loss / len(iterator), epoch_acc / len(iterator)



# + [markdown] id="qECbmsYaZvcT"
# ### Training Loop
#
# We are going to create helper functions that will help us to visualize our training.
#
#
# 1. Time to string.

# + id="O5J__Uc6ZzcJ"
def hms_string(sec_elapsed):
  h = int(sec_elapsed / (60 * 60))
  m = int((sec_elapsed % (60 * 60)) / 60)
  s = sec_elapsed % 60
  return "{}:{:>02}:{:>05.2f}".format(h, m, s)


# + [markdown] id="T8KLeTjKZ1Ta"
# 2. tabulate training epoch.

# + id="EbuJE5x_Z7f4"
from prettytable import PrettyTable

def visualize_training(start, end, train_loss, train_accuracy, val_loss, val_accuracy, title):
  data = [
       ["Training", f'{train_loss:.3f}', f'{train_accuracy:.3f}', f"{hms_string(end - start)}" ],
       ["Validation", f'{val_loss:.3f}', f'{val_accuracy:.3f}', "" ],       
  ]
  table = PrettyTable(["CATEGORY", "LOSS", "ACCURACY", "ETA"])
  table.align["CATEGORY"] = 'l'
  table.align["LOSS"] = 'r'
  table.align["ACCURACY"] = 'r'
  table.align["ETA"] = 'r'
  table.title = title
  for row in data:
    table.add_row(row)
  print(table)


# + colab={"base_uri": "https://localhost:8080/"} id="A6Ht54F-GsbZ" outputId="697eab7c-8bbb-43cf-9f86-ee4c648caac8"
N_EPOCHS = 10
best_valid_loss = float('inf')
for epoch in range(N_EPOCHS):
    start = time.time()
    train_loss, train_acc = train(model, train_iter, optimizer, criterion)
    valid_loss, valid_acc = evaluate(model, valid_iter, criterion)
    title = f"EPOCH: {epoch+1:02}/{N_EPOCHS:02} {'saving best model...' if valid_loss < best_valid_loss else 'not saving...'}"
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'best-model.pt')
    end = time.time()
    visualize_training(start, end, train_loss, train_acc, valid_loss, valid_acc, title)


# + [markdown] id="9gtH2gySa-_P"
# ### Evaluating the best model.

# + colab={"base_uri": "https://localhost:8080/"} id="VYnP8R5HGsYu" outputId="ff2bdc02-3f43-4a2e-b653-c59edad35e8f"
model.load_state_dict(torch.load('best-model.pt'))

test_loss, test_acc = evaluate(model, test_iter, criterion)
print(f'Test Loss: {test_loss:.3f} | Test Acc: {test_acc*100:.2f}%')

# + [markdown] id="JyIIW748bKD_"
# ### Model Inference

# + id="z4KPkW0PbX6Z"
from spacy.lang.id import Indonesian
nlp = Indonesian()


# + colab={"base_uri": "https://localhost:8080/"} id="SAmYKbb7bIOs" outputId="6dc27f11-283b-4f90-8f61-9f1c3278e495"
def infer(model, sentence, true_tags=None):
    model.eval()
    # tokenize sentence
    tokens = [token.text.lower() for token in nlp(sentence)]
    # transform to indices based on corpus vocab
    numericalized_tokens = [WORD.vocab.stoi[t] for t in tokens]
    # find unknown words
    unk_idx = WORD.vocab.stoi[WORD.unk_token]
    unks = [t for t, n in zip(tokens, numericalized_tokens) if n == unk_idx]
    # begin prediction
    token_tensor = torch.LongTensor(numericalized_tokens)
    token_tensor = token_tensor.unsqueeze(-1).to(device)
    predictions = model(token_tensor)
    # convert results to tags
    top_predictions = predictions.argmax(-1)
    predicted_tags = [TAG.vocab.itos[t.item()] for t in top_predictions]
    # print inferred tags
    max_len_token = max([len(token) for token in tokens] + [len("word")])
    max_len_tag = max([len(tag) for tag in predicted_tags] + [len("pred")])
    print(
        f"{'word'.ljust(max_len_token)}\t{'unk'.ljust(max_len_token)}\t{'pred tag'.ljust(max_len_tag)}" 
        + ("\ttrue tag" if true_tags else "")
        )
    for i, token in enumerate(tokens):
      is_unk = "✓" if token in unks else ""
      print(
          f"{token.ljust(max_len_token)}\t{is_unk.ljust(max_len_token)}\t{predicted_tags[i].ljust(max_len_tag)}" 
          + (f"\t{true_tags[i]}" if true_tags else "")
          )
    return tokens, predicted_tags, unks

sentence = "Sementara itu, Kepala Pelaksana BPBD Luwu Utara Muslim Muchtar mengatakan, terdapat 15.000 jiwa mengungsi akibat banjir bandang."
tags = ["O", "O", "O", "O", "O", "B-ORGANIZATION", "I-ORGANIZATION", "L-ORGANIZATION", "B-PERSON", "L-PERSON", "O", "O", "O", "U-QUANTITY", "O", "O", "O", "O", "O", "O"]
words, infer_tags, unknown_tokens = infer(model, sentence=sentence, true_tags=tags)


# + id="-As2_DNEb5Rh"


StarcoderdataJupyterScriptsDedupFiltered