# -*- coding: utf-8 -*-
"""Drug Name Detection spaCy v3 NER.ipynb

Automatically generated by Colaboratory.

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

!pip install -U spacy
!python -m spacy download en_core_web_lg

!pip install typer
#!pip install spacy[cuda110]

import spacy
import srsly
import pandas as pd
spacy.prefer_gpu()

# download training set
!wget https://raw.githubusercontent.com/explosion/projects/v3/tutorials/ner_drugs/assets/drugs_training.jsonl

# download dev set
!wget https://raw.githubusercontent.com/explosion/projects/v3/tutorials/ner_drugs/assets/drugs_eval.jsonl

# download script converting jsonl to .spacy format
!wget https://raw.githubusercontent.com/explosion/projects/v3/tutorials/ner_drugs/scripts/preprocess.py

"""#### Convert JSONL to .spaCy format"""

!python preprocess.py ./drugs_training.jsonl ./drugs_training.spacy
!python preprocess.py ./drugs_eval.jsonl ./drugs_eval.spacy

"""#### Generate config"""

!python -m spacy init config --lang en --pipeline ner configs/config.cfg --force --optimize accuracy

"""#### Training"""

!python -m spacy train configs/config.cfg --output models/ --paths.train ./drugs_training.spacy --paths.dev ./drugs_eval.spacy --gpu-id 0 --components.tok2vec.model.encode.width 64 --components.tok2vec.model.encode.depth 4

!python -m spacy evaluate ./models/model-last/ ./drugs_eval.spacy --gpu-id -1 --output evaluate.json

"""#### Retrain a new model focusing on Recall"""

!python -m spacy train configs/config.cfg --output models_recall/ --paths.train ./drugs_training.spacy --paths.dev ./drugs_eval.spacy --gpu-id 0 --components.tok2vec.model.encode.width 64 --components.tok2vec.model.encode.depth 4 --training.score_weights.ents_r 1.0 --training.score_weights.ents_f 0.0

!python -m spacy evaluate ./models_recall/model-last/ ./drugs_eval.spacy --gpu-id -1 --output evaluate.json

nlp = spacy.load("./models/model-last/")

data = srsly.read_jsonl("./drugs_eval.jsonl")

for idx in range(25):
    test_case = next(data)
    sentence = test_case["text"] 
    doc = nlp(sentence)
    y_pred = {ent: ent.text for ent in doc.ents}
    print(f"\n\n[Sentence]: \n\t{sentence}")
    print(f"\n[Predicted]: \n\t {y_pred}")
    print("-"*100)

