# -*- coding: utf-8 -*-
"""Toxic Comment Multilabel Classification using spaCy v3 TextCategorizer.ipynb

Automatically generated by Colaboratory.

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

!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/tianqwang/Toxic-Comment-Classification-Challenge/master/data/train.csv train.csv

# download test set
!wget https://raw.githubusercontent.com/tianqwang/Toxic-Comment-Classification-Challenge/master/data/test.csv test.csv

# download script converting jsonl to .spacy format
!wget https://raw.githubusercontent.com/explosion/projects/v3/pipelines/textcat_multilabel_demo/scripts/convert.py

labels = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']

def df_csv_to_jsonl(df, output_path):
    x = list(df['comment_text'].to_dict().values())
    y = df[labels].to_dict(orient='records')

    data = [{"text": _x, "cats": _y} for _x, _y in zip(x, y)]
    srsly.write_jsonl(output_path, data)

train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")

# split train_df into train and val df
train = train_df.sample(frac=0.8)
validation = train_df.drop(train.index)

# convert from dataframe to jsonl
df_csv_to_jsonl(train, "./train.json")
df_csv_to_jsonl(validation, "./validation.json")

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

!python convert.py en ./train.json ./train.spacy

!python convert.py en ./validation.json ./validation.spacy

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

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

"""#### Training"""

# edit config
# --components.tok2vec.model.encode.width 32 --components.tok2vec.model.encode.depth 5
# --nlp.batch_size 32 --training.batcher.start 32 --training.batcher.stop 128
#

!python -m spacy train configs/config.cfg --output models/ --paths.train ./train.spacy --paths.dev ./validation.spacy --gpu-id 0 --components.tok2vec.model.encode.width 32 --components.tok2vec.model.encode.depth 5

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

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

for idx, row in test_df.iterrows():
    sentence = row['comment_text']
    doc = nlp(sentence)
    y_pred = {label: pred_prob >= 0.5 for label, pred_prob in doc.cats.items()}
    print(f"\n\n[Sentence]: \n\t{sentence}")
    print(f"\n[Predicted]: \n\t {y_pred}")
    print("-"*100)

    if idx >= 15:
        break

