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

# + papermill={"duration": 7.929453, "end_time": "2022-02-18T19:58:58.405664", "exception": false, "start_time": "2022-02-18T19:58:50.476211", "status": "completed"} tags=[]
from tqdm.auto import tqdm
import os

import numpy as np
import pandas as pd 

import matplotlib.pyplot as plt
import seaborn as sns 

from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split , StratifiedKFold


import tensorflow as tf 
import tensorflow.keras.backend as K
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import Model, load_model, save_model
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.layers import Input,Dense, LSTM, RNN, Bidirectional, GlobalAveragePooling2D , Dropout, Conv1D, Flatten

from transformers import TFAutoModel , AutoTokenizer
# import ray
# from ray import tune



# + papermill={"duration": 0.031405, "end_time": "2022-02-18T19:58:58.457299", "exception": false, "start_time": "2022-02-18T19:58:58.425894", "status": "completed"} tags=[]
class config:
    #train_path = "../input/dravidianlangtech2022-personal/Train_Data_Combined.csv"
    #val_path = "../input/dravidianlangtech2022-personal/Validation_Data_Combined.csv"
    test_path = "../input/test-for-dravid-lang-tech-new/ta-misogyny-test.csv"
    save_dir = "./result"
    seed = 55
    try:
        AUTOTUNE = tf.data.AUTOTUNE   
    except:
        AUTOTUNE = tf.data.experimental.AUTOTUNE 
    epochs = 50
    max_len = 64
    batch_size = 32
    hf_path = "google/muril-base-cased"
    tokenizer_path = "../input/with-n-abusive-comment-detection-tamil-dravidianl/result/muril_tokenizer"  
    model_weights = "../input/with-n-abusive-comment-detection-tamil-dravidianl/result"
def seed_everything(seed = config.seed):
    print(f"seeded everything to seed {seed}")
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)
if not os.path.exists(config.save_dir):
    os.makedirs(config.save_dir)
seed_everything()


# + papermill={"duration": 0.05603, "end_time": "2022-02-18T19:58:58.533176", "exception": false, "start_time": "2022-02-18T19:58:58.477146", "status": "completed"} tags=[]
col_names = ['text']
df_test = pd.read_csv(config.test_path,names=col_names,sep='\t')

# + papermill={"duration": 0.048078, "end_time": "2022-02-18T19:58:58.601649", "exception": false, "start_time": "2022-02-18T19:58:58.553571", "status": "completed"} tags=[]
df_test.info()

# + papermill={"duration": 3.25411, "end_time": "2022-02-18T19:59:01.877389", "exception": false, "start_time": "2022-02-18T19:58:58.623279", "status": "completed"} tags=[]
#Tokenization Process
tokenizer = AutoTokenizer.from_pretrained(config.hf_path)
tokenizer.save_pretrained(os.path.join(config.save_dir , "muril_tokenizer"))


# + papermill={"duration": 0.034419, "end_time": "2022-02-18T19:59:01.935196", "exception": false, "start_time": "2022-02-18T19:59:01.900777", "status": "completed"} tags=[]
def fast_encode(texts, tokenizer, chunk_size=512, maxlen=config.max_len):
    
    input_ids = []
    tt_ids = []
    at_ids = []
    
    for i in tqdm(range(0, len(texts), chunk_size)):
        text_chunk = texts[i:i+chunk_size]
        encs = tokenizer(
                    text_chunk,
                    max_length = config.max_len,
                    padding='max_length',
                    truncation=True
                    )
        
        input_ids.extend(encs['input_ids'])
        tt_ids.extend(encs['token_type_ids'])
        at_ids.extend(encs['attention_mask'])
    
    return {'input_ids': input_ids, 'token_type_ids': tt_ids, 'attention_mask':at_ids}


# + papermill={"duration": 0.113987, "end_time": "2022-02-18T19:59:02.071772", "exception": false, "start_time": "2022-02-18T19:59:01.957785", "status": "completed"} tags=[]
test_token_data = fast_encode(list(df_test['text'].values), tokenizer)
#train_token_data['label'] = list(df_train['label'].values)

# + papermill={"duration": 0.049698, "end_time": "2022-02-18T19:59:02.145001", "exception": false, "start_time": "2022-02-18T19:59:02.095303", "status": "completed"} tags=[]
df_tokenized_test = pd.DataFrame(test_token_data)
len(df_tokenized_test['input_ids'][0])

# + papermill={"duration": 0.035817, "end_time": "2022-02-18T19:59:02.209855", "exception": false, "start_time": "2022-02-18T19:59:02.174038", "status": "completed"} tags=[]
del test_token_data 


# + papermill={"duration": 0.03299, "end_time": "2022-02-18T19:59:02.267257", "exception": false, "start_time": "2022-02-18T19:59:02.234267", "status": "completed"} tags=[]
#preparing dataset
def test_prep_function(embeddings):
  input_ids = embeddings['input_ids']
  attention_mask = embeddings['attention_mask']
  token_type_ids = embeddings['token_type_ids']

  #target = tf.cast(target, tf.int32)
  
  return {'input_ids': input_ids ,'token_type_ids':token_type_ids,'attention_mask': attention_mask}


# + papermill={"duration": 6.154511, "end_time": "2022-02-18T19:59:08.445075", "exception": false, "start_time": "2022-02-18T19:59:02.290564", "status": "completed"} tags=[]
# Detect hardware, return appropriate distribution strategy
try:
    # TPU detection. No parameters necessary if TPU_NAME environment variable is
    # set: this is always the case on Kaggle.
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print('Running on TPU ', tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.
    strategy = tf.distribute.get_strategy()

print("REPLICAS: ", strategy.num_replicas_in_sync)


# + papermill={"duration": 0.035582, "end_time": "2022-02-18T19:59:08.505089", "exception": false, "start_time": "2022-02-18T19:59:08.469507", "status": "completed"} tags=[]
def create_model(transformer_model):
  input_id_layer = Input(shape=(config.max_len,) ,dtype = tf.int32 , name = 'input_ids')
  attention_mask_layer = Input(shape=(config.max_len,) , dtype = tf.int32 , name = 'attention_mask')
  token_type_layer = Input(shape=(config.max_len,) , dtype = tf.int32 , name = 'token_type_ids')
  transformer = transformer_model(input_ids = input_id_layer ,token_type_ids=token_type_layer,attention_mask = attention_mask_layer)[0]
  

  x = Dropout(0.5)(transformer)
  x = Conv1D(1,1)(x)
  x = Flatten()(x) 
  predictions = Dense(8, activation = "softmax")(x)

  model = Model(inputs=[input_id_layer ,token_type_layer, attention_mask_layer], outputs = predictions)
  model.compile(
      optimizer = Adam(learning_rate= 0.01),
      metrics = ['accuracy'],
      loss = 'sparse_categorical_crossentropy'
  )

  return model

# + papermill={"duration": 75.051693, "end_time": "2022-02-18T20:00:23.581194", "exception": false, "start_time": "2022-02-18T19:59:08.529501", "status": "completed"} tags=[]
with strategy.scope():
  transformer_model = TFAutoModel.from_pretrained(config.hf_path)
  transformer_model.bert.trainable = False 
  model = create_model(transformer_model)
model.summary()
tf.keras.utils.plot_model(model, show_shapes=True,show_dtype=True)

# + papermill={"duration": 32.095706, "end_time": "2022-02-18T20:00:55.705306", "exception": false, "start_time": "2022-02-18T20:00:23.609600", "status": "completed"} tags=[]
from sklearn.metrics import accuracy_score
test_embeddings = {'input_ids': df_tokenized_test['input_ids'].tolist() ,'token_type_ids': df_tokenized_test['token_type_ids'].tolist(),"attention_mask":df_tokenized_test['attention_mask'].tolist()}
#y_train = df_tokenized_train['label']
#y_test = df_tokenized_val['label']
#train_dataset = tf.data.Dataset.from_tensor_slices((train_embeddings , y_train))
test_dataset = tf.data.Dataset.from_tensor_slices((test_embeddings))
test_dataset = (
      test_dataset
      .map(test_prep_function , num_parallel_calls = config.AUTOTUNE)
      .batch(config.batch_size)
      .prefetch(config.AUTOTUNE)
  )
test_steps = len(test_embeddings['input_ids'])//config.batch_size
model.load_weights(f'{config.model_weights}/muril_fold_trained.h5')
y_predict = model.predict(test_dataset , verbose = 1)
predictions = y_predict
preds_classes = np.argmax(predictions, axis=-1)


# + papermill={"duration": 0.042108, "end_time": "2022-02-18T20:00:55.777918", "exception": false, "start_time": "2022-02-18T20:00:55.735810", "status": "completed"} tags=[]
preds_classes

# + papermill={"duration": 0.0389, "end_time": "2022-02-18T20:00:55.847720", "exception": false, "start_time": "2022-02-18T20:00:55.808820", "status": "completed"} tags=[]
df_pred = pd.DataFrame(preds_classes,columns=['label'])

# + papermill={"duration": 0.03987, "end_time": "2022-02-18T20:00:55.918401", "exception": false, "start_time": "2022-02-18T20:00:55.878531", "status": "completed"} tags=[]
df_pred.replace({0:'Counter-speech',
                 1:'Homophobia', 
                 2:'Hope-Speech', 
                 3:'Misandry', 
                 4:'Misogyny',
                 5:'None-of-the-above',
                 6:'Transphobic',
                 7:'Xenophobia'},inplace = True)

# + papermill={"duration": 0.04833, "end_time": "2022-02-18T20:00:55.997638", "exception": false, "start_time": "2022-02-18T20:00:55.949308", "status": "completed"} tags=[]
df_pred

# + papermill={"duration": 0.039655, "end_time": "2022-02-18T20:00:56.068705", "exception": false, "start_time": "2022-02-18T20:00:56.029050", "status": "completed"} tags=[]
df_test[list(df_pred.columns)] = df_pred

# + papermill={"duration": 0.047372, "end_time": "2022-02-18T20:00:56.147316", "exception": false, "start_time": "2022-02-18T20:00:56.099944", "status": "completed"} tags=[]
df_test

# + papermill={"duration": 0.04811, "end_time": "2022-02-18T20:00:56.227769", "exception": false, "start_time": "2022-02-18T20:00:56.179659", "status": "completed"} tags=[]
df_test.to_csv('BpHigh_tamil.tsv',sep="\t")

# + papermill={"duration": 0.042532, "end_time": "2022-02-18T20:00:56.302594", "exception": false, "start_time": "2022-02-18T20:00:56.260062", "status": "completed"} tags=[]
df_test.label.value_counts()

# + papermill={"duration": 0.031737, "end_time": "2022-02-18T20:00:56.366397", "exception": false, "start_time": "2022-02-18T20:00:56.334660", "status": "completed"} tags=[]


StarcoderdataJupyterScriptsDedupFiltered