"""
# About this notebook

This uses the implementation from @maunish as a basis and makes an inference using Electra-large model from google, scoring 0.504 on LB.

If you like it, please upvote maunish's notebooks mentioned below.

This uses a dataset that contains Electra pre-trained for 3 epochs, trained for 8 epochs, 3 folds
"""
"""
# Original context from @maunish

This is inference notebooks that is trained using below notebooks.

This notebook uses the model created in pretrain any model notebook.

1. Pretrain Roberta Model: https://www.kaggle.com/maunish/clrp-pytorch-roberta-pretrain
2. Finetune Roberta Model: https://www.kaggle.com/maunish/clrp-pytorch-roberta-finetune <br/>
   Finetune Roberta Model TPU: https://www.kaggle.com/maunish/clrp-pytorch-roberta-finetune-tpu
3. Inference Notebook: this notebook
4. Roberta + SVM: https://www.kaggle.com/maunish/clrp-roberta-svm

"""
"""
# Imports
"""
import os
import gc
import sys
import cv2
import math
import time
import tqdm
import random
import numpy as np
import pandas as pd
import seaborn as sns
from tqdm import tqdm
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

from sklearn.model_selection import KFold,StratifiedKFold

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader,Sampler

import multiprocessing
import more_itertools


from transformers import (AutoModel, AutoTokenizer, AutoModelForSequenceClassification)

from colorama import Fore, Back, Style
y_ = Fore.YELLOW
r_ = Fore.RED
g_ = Fore.GREEN
b_ = Fore.BLUE
m_ = Fore.MAGENTA
c_ = Fore.CYAN
sr_ = Style.RESET_ALL
"""
# Load data
"""
train_data = pd.read_csv('../input/commonlitreadabilityprize/train.csv')
test_data = pd.read_csv('../input/commonlitreadabilityprize/test.csv')
sample = pd.read_csv('../input/commonlitreadabilityprize/sample_submission.csv')
"""
# Config
"""
config = {
    'learning_rate':2e-5,
    'batch_size':16,
    'epochs':3,
    'nfolds':5,
    'seed':1, # 42
    'max_len':256,
}

def seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONASSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

seed_everything(seed=config['seed'])
"""
# Dataset class
"""
class CLRPDataset(Dataset):
    def __init__(self,df,tokenizer):
        self.excerpt = df['excerpt'].to_numpy()
        self.tokenizer = tokenizer
    
    def __getitem__(self,idx):
        encode = self.tokenizer(self.excerpt[idx],return_tensors='pt',
                                max_length=config['max_len'],
                                padding='max_length',truncation=True)
        return encode
    
    def __len__(self):
        return len(self.excerpt)
"""
# Electra Model
"""
class AttentionHead(nn.Module):
    def __init__(self, in_features, hidden_dim, num_targets):
        super().__init__()
        self.in_features = in_features
        self.middle_features = hidden_dim

        self.W = nn.Linear(in_features, hidden_dim)
        self.V = nn.Linear(hidden_dim, 1)
        self.out_features = hidden_dim

    def forward(self, features):
        att = torch.tanh(self.W(features))

        score = self.V(att)

        attention_weights = torch.softmax(score, dim=1)

        context_vector = attention_weights * features
        context_vector = torch.sum(context_vector, dim=1)

        return context_vector
class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.electra = AutoModel.from_pretrained('../input/clrp-electra-large/clrp_electra_large')
        self.head = AttentionHead(256,256,1) 
        self.dropout = nn.Dropout(0.1)
        self.linear = nn.Linear(256,1)

    def forward(self,**xb):
        x = self.electra(**xb)[0]
        x = self.head(x)
        x = self.dropout(x)
        x = self.linear(x)
        return x
"""
# Predict
"""
def get_prediction(df,path,device='cuda'):        
    model = Model()
    model.load_state_dict(torch.load(path,map_location=device))
    model.to(device)
    model.eval()
    
    tokenizer = AutoTokenizer.from_pretrained('../input/electra-tokenizer-files')
    
    test_ds = CLRPDataset(df,tokenizer)
    test_dl = DataLoader(test_ds,
                        batch_size = config["batch_size"],
                        shuffle=False,
                        num_workers = 4,
                        pin_memory=True)
    
    predictions = list()
    for i, (inputs) in tqdm(enumerate(test_dl)):
        inputs = {key:val.reshape(val.shape[0],-1).to(device) for key,val in inputs.items()}
        outputs = model(**inputs)
        outputs = outputs.cpu().detach().numpy().ravel().tolist()
        predictions.extend(outputs)
        
    torch.cuda.empty_cache()
    return np.array(predictions)
"""
# Prediction
"""
pred1 = get_prediction(test_data,'../input/electra-large-fit8/model0/model0.bin')
pred2 = get_prediction(test_data,'../input/electra-large-fit8/model1/model1.bin')
pred3 = get_prediction(test_data,'../input/electra-large-fit8/model2/model2.bin')

predictions1 = (pred1 + pred2 + pred3)/3
"""
# Submission
"""
sample['target'] = predictions1
sample.to_csv('submission.csv',index=False)
sample