"""
This is an inference notebook with the `t5-large` model. You can achieve decent results with this model and choose to use the provided model in an ensemble of yours.
"""
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))
from typing import Dict, Any, Union

from pathlib import Path
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import StratifiedKFold

import torch
import torch.nn as nn

import torch.utils.data as D
from torch.utils.data.dataset import Dataset, IterableDataset
from torch.utils.data.dataloader import DataLoader

from transformers import AutoModel, AutoTokenizer, AutoConfig
from transformers import PreTrainedModel

from tqdm.notebook import tqdm

from sklearn.svm import SVR
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import mean_squared_error

import yaml, gc
"""
### Folders and Dataframes
"""
BASE_PATH = Path('/kaggle/input/commonlit-t5-large')
DATA_PATH = Path('/kaggle/input/commonlitreadabilityprize/')
assert DATA_PATH.exists()
MODELS_PATH = Path(BASE_PATH/'best_models')
assert MODELS_PATH.exists()
train_df = pd.read_csv(DATA_PATH/'train.csv')
test_df = pd.read_csv(DATA_PATH/'test.csv')
sample_df = pd.read_csv(DATA_PATH/'sample_submission.csv')

def remove_unnecessary(df):
    df.drop(df[df['target'] == 0].index, inplace=True)
    df.reset_index(drop=True, inplace=True)
    
remove_unnecessary(train_df)
"""
### Configuration
"""
class Config(): 
    NUM_FOLDS = 6
    NUM_EPOCHS = 3
    BATCH_SIZE = 16
    MAX_LEN = 248
    MODEL_PATH = BASE_PATH/'lm'
#     TOKENIZER_PATH = str(MODELS_PATH/'roberta-base-0')
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    SEED = 1000
    NUM_WORKERS = 2
    MODEL_FOLDER = MODELS_PATH
    model_name = 't5-large'
    svm_kernels = ['rbf']
    svm_c = 5

cfg = Config()
train_df['normalized_target'] = (train_df['target'] - train_df['target'].mean()) / train_df['target'].std()
"""
### Read Existing Models
"""
model_path = MODELS_PATH
assert model_path.exists()
!ls {MODELS_PATH}
class AttentionHead(nn.Module):
    
    def __init__(self, in_features, hidden_dim, num_targets):
        super().__init__()
        self.in_features = in_features
        
        self.hidden_layer = nn.Linear(in_features, hidden_dim)
        self.final_layer = nn.Linear(hidden_dim, num_targets)
        self.out_features = hidden_dim
        
    def forward(self, features):
        att = torch.tanh(self.hidden_layer(features))
        score = self.final_layer(att)
        attention_weights = torch.softmax(score, dim=1)
        return attention_weights
from transformers import T5EncoderModel

class CommonLitModel(nn.Module):
    def __init__(self):
        super(CommonLitModel, self).__init__()
        config = AutoConfig.from_pretrained(cfg.MODEL_PATH)
        config.update({
            "output_hidden_states": True,
            "hidden_dropout_prob": 0.0,
            "layer_norm_eps": 1e-7
        })
        self.transformer_model = T5EncoderModel.from_pretrained(cfg.MODEL_PATH, config=config)
        self.attention = AttentionHead(config.hidden_size, 512, 1)
        self.regressor = nn.Linear(config.hidden_size, 1)
    
    def forward(self, input_ids, attention_mask):
        last_layer_hidden_states = self.transformer_model(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
        weights = self.attention(last_layer_hidden_states)
        context_vector = torch.sum(weights * last_layer_hidden_states, dim=1) 
        return self.regressor(context_vector), context_vector
def load_model(i):
    inference_model = CommonLitModel()
    inference_model = inference_model.cuda()
    inference_model.load_state_dict(torch.load(str(model_path/f'{i + 1}_pytorch_model.bin')))
    inference_model.eval();
    return inference_model
"""
### DataSet and Tokenizers
"""
def convert_to_list(t):
    return t.flatten().long()

class CommonLitDataset(nn.Module):
    def __init__(self, text, test_id, tokenizer, max_len=128):
        self.excerpt = text
        self.test_id = test_id
        self.max_len = max_len
        self.tokenizer = tokenizer
    
    def __getitem__(self,idx):
        encode = self.tokenizer(self.excerpt[idx],
                                return_tensors='pt',
                                max_length=self.max_len,
                                padding='max_length',
                                truncation=True)
        return {'input_ids': convert_to_list(encode['input_ids']),
                'attention_mask': convert_to_list(encode['attention_mask']),
                'id': self.test_id[idx]}
    
    def __len__(self):
        return len(self.excerpt)
!ls {MODELS_PATH}/tokenizer-1
from transformers import T5Tokenizer

tokenizers = []
for i in range(1, cfg.NUM_FOLDS):
    tokenizer_path = MODELS_PATH/f"tokenizer-{i}"
    print(tokenizer_path)
    assert(Path(tokenizer_path).exists())
    tokenizer = T5Tokenizer.from_pretrained(str(tokenizer_path))
    tokenizers.append(tokenizer)
def create_dl(df, tokenizer):
    text = df['excerpt'].values
    ids = df['id'].values
    ds = CommonLitDataset(text, ids, tokenizer, max_len=cfg.MAX_LEN)
    return DataLoader(ds, 
                      batch_size = cfg.BATCH_SIZE,
                      shuffle=False,
                      num_workers = 1,
                      pin_memory=True,
                      drop_last=False
                     )
"""
#### Extract Embeddings
"""
def get_cls_embeddings(dl, transformer_model):
    cls_embeddings = []
    with torch.no_grad():
        for input_features in tqdm(dl, total=len(dl)):
            _, context_vector = transformer_model(input_features['input_ids'].cuda(), input_features['attention_mask'].cuda())
#             cls_embeddings.extend(output['last_hidden_state'][:,0,:].detach().cpu().numpy())
            embedding_out = context_vector.detach().cpu().numpy()
            cls_embeddings.extend(embedding_out)
    return np.array(cls_embeddings)
"""
#### Extract Number of Bins
"""
num_bins = int(np.ceil(np.log2(len(train_df))))
train_df['bins'] = pd.cut(train_df['target'], bins=num_bins, labels=False)
bins = train_df['bins'].values
"""
#### Training
"""
def rmse_score(X, y):
    return np.sqrt(mean_squared_error(X, y))
%%time

train_target = train_df['normalized_target'].values

def calc_mean(scores):
    return np.mean(np.array(scores), axis=0)

final_scores = []
final_rmse = []
for j, tokenizer in enumerate(tokenizers):
    print('Model', j)
    test_dl = create_dl(test_df, tokenizer)
    train_dl = create_dl(train_df, tokenizer)
    transformer_model = load_model(j)
    transformer_model.cuda()
    X = get_cls_embeddings(train_dl, transformer_model)
    y = train_target
    X_test = get_cls_embeddings(test_dl, transformer_model)
    kfold = StratifiedKFold(n_splits=cfg.NUM_FOLDS)
    scores = []
    rmse_scores = []
    for kernel in cfg.svm_kernels:
        print('Kernel', kernel)
        kernel_scores = []
        kernel_rmse_scores = []
        for k, (train_idx, valid_idx) in enumerate(kfold.split(X, bins)):

            print('Fold', k, train_idx.shape, valid_idx.shape)
            model = SVR(C=cfg.svm_c, kernel=kernel, gamma='auto')

            X_train, y_train = X[train_idx], y[train_idx]
            X_valid, y_valid = X[valid_idx], y[valid_idx]
            model.fit(X_train, y_train)
            prediction = model.predict(X_valid)
            kernel_rmse_scores.append(rmse_score(prediction, y_valid))
            print('rmse_score', kernel_rmse_scores[k])
            kernel_scores.append(model.predict(X_test))
        scores.append(calc_mean(kernel_scores))
        rmse_scores.append(calc_mean(kernel_rmse_scores))
    final_scores.append(calc_mean(scores))
    final_rmse.append(calc_mean(rmse_scores))
    del transformer_model
    torch.cuda.empty_cache()
    del tokenizer
    gc.collect()
print('FINAL RMSE score', np.mean(np.array(final_rmse)))
final_scores_bck = final_scores
final_scores_bck
# (train_df['target'] - cfg.train_target_mean) / cfg.train_target_std
final_scores = np.array(final_scores) * train_df['target'].std() + train_df['target'].mean()
final_scores
"""
#### Ensure the mean of the prediction equals the mean of the training data
"""
def calc_mean(scores):
    return np.mean(np.array(scores), axis=0)
target_mean = train_df['target'].mean()
final_scores_flat = calc_mean(final_scores).flatten()
final_scores_mean = final_scores_flat.mean()
target_mean, np.array(final_scores).mean()
mean_diff = target_mean - final_scores_mean
mean_diff, mean_diff / len(final_scores)
sample_df['target'] = final_scores_flat
# sample_df['target'] = len(final_scores) / np.sum(1 / np.array(final_scores), axis=0) # harmonic mean
sample_df
pd.DataFrame(sample_df).to_csv('submission.csv', index=False)
!cat submission.csv