# ---
# 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
# ---

# ![JohnSnowLabs](https://nlp.johnsnowlabs.com/assets/images/logo.png)

# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/models_hub/Train_a_Spark_NLP_Model.ipynb)

#

# + id="WZEkgBuabyGF"
# %%capture
 
# Setup Spark NLP on Colab
# !wget https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp-workshop/master/colab_setup.sh -O - | bash

# + colab={"base_uri": "https://localhost:8080/"} id="7i4Z9OrgcE4b" outputId="579e109c-2eb9-4a6c-aadb-8ad4c14b76c1"
import sparknlp
 
spark = sparknlp.start() # for GPU training >> sparknlp.start(gpu = True) # for Spark 2.3 =>> sparknlp.start(spark23 = True)
import pyspark.sql.functions as F
from sparknlp.annotator import *
from sparknlp.base import *
import sparknlp
from sparknlp.pretrained import PretrainedPipeline
 
print("Spark NLP version", sparknlp.version())
 
print("Apache Spark version:", spark.version)

# + id="ooWKiaQEcUPB"
#download training data
# !wget -q https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp/master/src/test/resources/conll2003/eng.train
# !wget -q https://raw.githubusercontent.com/JohnSnowLabs/spark-nlp/master/src/test/resources/conll2003/eng.testa

# + id="jn72rDuTcW_l"
from sparknlp.training import CoNLL

training_data = CoNLL().readDataset(spark, './eng.train')
testing_data= CoNLL().readDataset(spark, './eng.testa')

# + colab={"base_uri": "https://localhost:8080/"} id="1RW2pBP_cdgP" outputId="a70f0a9d-b7b8-4d4c-8bf9-98dd85e06917"
import pyspark.sql.functions as F

training_data.select(F.explode(F.arrays_zip('token.result', 'pos.result',  'label.result')).alias("cols")) \
.select(F.expr("cols['0']").alias("token"),
        F.expr("cols['1']").alias("pos"),
        F.expr("cols['2']").alias("ner_label")).show(truncate=50)

# + [markdown] id="cLule_H4rDmv"
# ## 1. Create Spark NLP train pipeline

# + id="Z6mvq8Avcjyx" colab={"base_uri": "https://localhost:8080/"} outputId="3b029733-eaf7-4eb4-d1ad-2b1dd3c2761a"
# !mkdir ner_logs

# + id="zijeZRPrcmE4" colab={"base_uri": "https://localhost:8080/"} outputId="4bf81eeb-db55-4229-d792-08a93b81fc6a"
# You can use any word embeddings you want (Glove, Elmo, Bert, custom etc.)

embeddings = WordEmbeddingsModel.pretrained('glove_100d')\
          .setInputCols(["document", "token"])\
          .setOutputCol("embeddings")

nerTagger = NerDLApproach()\
      .setInputCols(["sentence", "token", "embeddings"])\
      .setLabelColumn("label")\
      .setOutputCol("ner")\
      .setMaxEpochs(1)\
      .setLr(0.003)\
      .setBatchSize(32)\
      .setRandomSeed(0)\
      .setVerbose(1)\
      .setValidationSplit(0.2)\
      .setEvaluationLogExtended(True) \
      .setEnableOutputLogs(True)\
      .setIncludeConfidence(True)\
      .setOutputLogsPath('ner_logs') # if not set, logs will be written to ~/annotator_logs
 #    .setGraphFolder('graphs') >> put your graph file (pb) under this folder if you are using a custom graph generated thru 4.1 NerDL-Graph.ipynb notebook
 #    .setEnableMemoryOptimizer() >> if you have a limited memory and a large conll file, you can set this True to train batch by batch 
    
ner_converter = NerConverter() \
    .setInputCols(['document', 'token', 'ner']) \
    .setOutputCol('ner_chunk')

ner_pipeline = Pipeline(stages=[
      embeddings,
      nerTagger,
      ner_converter
 ])

# + [markdown] id="6lJ8fCjmrLtw"
# ## 2. Train model

# + colab={"base_uri": "https://localhost:8080/"} id="rsJO74W-czVS" outputId="74055cac-d50c-470b-e420-4178278e873d"
# %%time
ner_model = ner_pipeline.fit(training_data)
ner_model.stages[-1].write().overwrite().save('outputs/ner_wiki_glove100d_en')

# + colab={"base_uri": "https://localhost:8080/"} id="3TIDdUYHdG7y" outputId="7eb373a6-f632-4e28-a31d-e069926b5fdf"
import pyspark.sql.functions as F
predictions = ner_model.transform(testing_data) 
predictions.select(F.explode(F.arrays_zip('token.result','label.result','ner.result')).alias("cols")) \
           .select(F.expr("cols['0']").alias("token"),
                   F.expr("cols['1']").alias("ground_truth"),
                   F.expr("cols['2']").alias("prediction")).show(truncate=False)

# + [markdown] id="M7gTMzBXSJY1"
# ## 3. Benchmark

# + id="hV8fcONKdMgF" colab={"base_uri": "https://localhost:8080/"} outputId="07be85e0-501f-447f-cdb9-38dca72d7fd3"
from sklearn.metrics import classification_report
 
preds_df = predictions.select(F.explode(F.arrays_zip('token.result','label.result','ner.result')).alias("cols")) \
                      .select(F.expr("cols['0']").alias("token"),
                              F.expr("cols['1']").alias("ground_truth"),
                              F.expr("cols['2']").alias("prediction")).toPandas()
 
print(classification_report(preds_df['ground_truth'], preds_df['prediction']))

# + [markdown] id="562zPiw4SlRg"
# ## 4. Saving model and Zipping it

# + id="cwIo061j8KPm" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="e5ae8d70-69f5-4ab3-cf93-9280aadfdc1d"
import shutil

shutil.make_archive("/content/outputs/ner_wiki_glove100d_en", 'zip', "/content/outputs/ner_wiki_glove100d_en")

StarcoderdataJupyterScriptsDedupFiltered