!nvidia-smi

%%capture
!echo "deb http://packages.cloud.google.com/apt gcsfuse-`lsb_release -c -s` main" | sudo tee /etc/apt/sources.list.d/gcsfuse.list
!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
!apt-get update
!apt-get install gcsfuse
!mkdir -p /content/gcs
!gcsfuse bucket-paper-acception /content/gcs

%%capture
!pip install polars
! pip install transformers
! pip install sentencepiece
!pip install fontstyle

# ----------
# ライブラリ
# ----------
import os
import random
import numpy as np
import torch
from psutil import virtual_memory

import polars as pl
import pandas as pd
from sklearn.model_selection import StratifiedKFold

from transformers import AutoModel, AutoTokenizer
import torch
import torch.nn.functional as F
from tqdm.auto import tqdm

import xgboost as xgb
import scipy.stats as stats
import lightgbm as lgbm
from sklearn.metrics import accuracy_score
import fontstyle
import warnings
warnings.simplefilter('ignore')

# ----------
# 設定
# ----------
num_fold = 5
seed = 0

DEVICE = "cuda" # "cpu" or "cuda"
tokenizer = None
BATCH_SIZE = 16
MAX_LEN = 768

# テキスト特徴として連結するカラム
txt_columns = ['title', 'keywords', 'abstract']

# ----------
# 関数
# ----------
def set_seed(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        current_device = torch.cuda.current_device()
        print("Device:", torch.cuda.get_device_name(current_device))
        ram_gb = virtual_memory().total / 1e9
        print('Your runtime has {:.1f} gigabytes of available RAM\n'.format(ram_gb))

def get_stratifiedkfold(train, target_col, n_splits, seed):
    kf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    generator = kf.split(train.to_numpy(), train.get_column(target_col).to_numpy())
    fold_array = np.zeros(len(train))
    for fold, (_, idx_valid) in enumerate(generator):
        fold_array[idx_valid] = fold
    return fold_array

# ----------
# データ
# ----------
train = pl.read_csv('/content/gcs/train.csv')
test = pl.read_csv('/content/gcs/test.csv')
sub = pl.read_csv('/content/gcs/submission.csv')

# ----------
# 前処理・特徴生成
# ----------
set_seed(seed)

# テキスト特徴の作成
# グループごとにFold数を設定
train =\
train.with_columns(
    pl.concat_str(txt_columns, sep='. ').alias('txt_feat'),
    # title
    pl.when(pl.col('title').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('title').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_title'),
    # abstract
    pl.when(pl.col('abstract').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('abstract').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_abstract'),
    # keywords
    pl.when(pl.col('keywords').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('keywords').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_keywords'),
    # group
    pl.concat_str(['year', 'y'], sep='-').alias('group'),
    )

test = \
test.with_columns(
    pl.concat_str(txt_columns, sep='. ').alias('txt_feat'),
    # title
    pl.when(pl.col('title').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('title').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_title'),
    # abstract
    pl.when(pl.col('abstract').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('abstract').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_abstract'),
    # keywords
    pl.when(pl.col('keywords').str.to_lowercase().is_in(['', 'nan', '0', 'blank']))
    .then(0)
    .otherwise(pl.col('keywords').str.to_lowercase().str.count_match(' ') + 1)
    .alias('num_keywords'),
    )

display(train.head(3))
display(test.head(3))

# ----------
# BERT
# ----------

# Dataset
class EmbedDataset(torch.utils.data.Dataset):
    def __init__(self, df):
        self.df = df
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        text = self.df[idx, 'txt_feat']
        tokens = tokenizer(
                text,
                None,
                add_special_tokens=True,
                padding='max_length',
                truncation=True,
                max_length=MAX_LEN,
                return_tensors="pt")
        tokens = {k:v.squeeze(0) for k,v in tokens.items()}
        return tokens

# Pooling
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output.last_hidden_state.detach().cpu()
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    )
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )

def get_embeddings(MODEL_NM='', MAX_LEN=512, BATCH_SIZE=4, verbose=True):
    global tokenizer, DEVICE

    model = AutoModel.from_pretrained(MODEL_NM)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NM)
    
    model = model.to(DEVICE)
    model.eval()
        
    # train
    all_train_text_feats = []
    for batch in tqdm(embed_dataloader_tr,total=len(embed_dataloader_tr)):
        input_ids = batch["input_ids"].to(DEVICE)
        attention_mask = batch["attention_mask"].to(DEVICE)
        with torch.no_grad():
            model_output = model(input_ids=input_ids,attention_mask=attention_mask)
        sentence_embeddings = mean_pooling(model_output, attention_mask.detach().cpu())
        # Normalize the embeddings
        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
        sentence_embeddings =  sentence_embeddings.squeeze(0).detach().cpu().numpy()
        all_train_text_feats.extend(sentence_embeddings)
    all_train_text_feats = np.array(all_train_text_feats)
    if verbose:
        print('Train embeddings shape',all_train_text_feats.shape)
    
    # test
    te_text_feats = []
    for batch in tqdm(embed_dataloader_te,total=len(embed_dataloader_te)):
        input_ids = batch["input_ids"].to(DEVICE)
        attention_mask = batch["attention_mask"].to(DEVICE)
        with torch.no_grad():
            model_output = model(input_ids=input_ids,attention_mask=attention_mask)
        sentence_embeddings = mean_pooling(model_output, attention_mask.detach().cpu())
        # Normalize the embeddings
        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
        sentence_embeddings =  sentence_embeddings.squeeze(0).detach().cpu().numpy()
        te_text_feats.extend(sentence_embeddings)
    te_text_feats = np.array(te_text_feats)
    if verbose:
        print('Test embeddings shape',te_text_feats.shape)
        
    # save feat
    np.save(f"{MODEL_NM.split('/')[-1]}_train", all_train_text_feats)
    np.save(f"{MODEL_NM.split('/')[-1]}_test", te_text_feats)

    return all_train_text_feats, te_text_feats

ds_tr = EmbedDataset(train)
embed_dataloader_tr = torch.utils.data.DataLoader(ds_tr,\
                        batch_size=BATCH_SIZE,\
                        shuffle=False)
ds_te = EmbedDataset(test)
embed_dataloader_te = torch.utils.data.DataLoader(ds_te,\
                        batch_size=BATCH_SIZE,\
                        shuffle=False)

%%time
MODEL_NM = 'microsoft/deberta-v3-base'
train_emb, test_emb = get_embeddings(MODEL_NM, MAX_LEN, BATCH_SIZE)

emb_col = [f'emb{i}' for i in range(train_emb.shape[1])]
train_emb_df = pl.DataFrame(train_emb, schema=emb_col)
train = pl.concat([train, train_emb_df], how='horizontal')
test_emb_df = pl.DataFrame(test_emb, schema=emb_col)
test = pl.concat([test, test_emb_df], how='horizontal')

# Run XGBoost
use_col = ['num_title', 'num_abstract', 'num_keywords'] + emb_col
test_x = test.select(use_col).to_numpy()

whole_va_preds = []
whole_test_preds = []
for seed in range(3):
    print(fontstyle.apply(f'< Seed : {seed} >', 'BLACK/BOLD'))
    set_seed(seed)
    train = train.with_columns(
        pl.Series(get_stratifiedkfold(train, 'group', num_fold, seed))
        .alias('folds')
        )
    
    oof_preds = np.zeros((len(train), ), dtype=np.float32)
    preds = []
    for fold in range(num_fold):
        tr_x = train.filter(pl.col('folds')!=fold).select(use_col).to_numpy()
        tr_y = train.filter(pl.col('folds')!=fold).select('y').to_numpy()
        va_x = train.filter(pl.col('folds')==fold).select(use_col).to_numpy()
        va_y = train.filter(pl.col('folds')==fold).select('y').to_numpy()

        params = {
        'objective': 'binary:logistic',
        'n_estimators': 10000,
        'random_state': 0, 
        'learning_rate': 0.01,
        'max_depth': 8,
        'colsample_bytree': 1.0,
        'colsample_bylevel': 0.5,
        'subsample': 0.9,
        'gamma': 0,
        'lambda': 1,
        'alpha': 0,
        'min_child_weight': 1,
        'tree_method': 'gpu_hist',
        }

        clf = xgb.XGBClassifier(**params)
        clf.fit(
            tr_x, tr_y,
            eval_set=[(va_x, va_y)],
            early_stopping_rounds=100,
            verbose=100)

        va_preds_p = clf.predict_proba(va_x)[:, 1]
        oof_preds[
            train.select(
                pl.when(pl.col('folds')==fold).then(True).otherwise(False)
                ).to_numpy().reshape(-1)
                ] = va_preds_p
        va_preds = (va_preds_p > 0.5).astype(int)
        score = accuracy_score(va_y, va_preds)
        print(f'Fold : {fold+1} Accuracy score: {score}')
        print()
        test_preds_p = clf.predict_proba(test_x)[:, 1]
        preds.append(test_preds_p)

    score_s = accuracy_score(train.select('y').to_numpy(), oof_preds > 0.5)
    print(fontstyle.apply(f'Seed{seed} Accuracy score : {score_s}', 'BLACK/BOLD'))
    print()
    whole_va_preds.append(oof_preds)
    whole_test_preds.append(preds)

# preds_va_p = np.mean(whole_va_preds, axis=0)
# whole_score = accuracy_score(train.select('y').to_numpy(), preds_va_p > 0.5)
# preds_test = (np.mean(np.mean(whole_test_preds, axis=0), axis=0) > 0.5).astype(int)
preds_va = np.array([np.where(preds > 0.5, 1, 0) for preds in whole_va_preds])
whole_score = accuracy_score(train.select('y').to_numpy(), stats.mode(preds_va, axis=0).mode.flatten())
test_preds_array = np.array(whole_test_preds)
test_preds_array = test_preds_array.reshape(test_preds_array.shape[0]*test_preds_array.shape[1], -1)
preds_test = np.array([np.where(preds > 0.5, 1, 0) for preds in test_preds_array])
preds_test = stats.mode(preds_test, axis=0).mode.flatten()
print()  
print(fontstyle.apply(f'whole Accuracy score: {whole_score}', 'BLACK/BOLD'))
print()

display(pl.Series(preds_test).value_counts())

# # Run LightGBM
# use_col = ['num_title', 'num_abstract', 'num_keywords'] + emb_col
# test_x = test.select(use_col).to_numpy()

# whole_va_preds = []
# whole_test_preds = []
# for seed in range(3):
#     print(fontstyle.apply(f'< Seed : {seed} >', 'BLACK/BOLD'))
#     set_seed(seed)
#     train = train.with_columns(
#         pl.Series(get_stratifiedkfold(train, 'group', num_fold, seed))
#         .alias('folds')
#         )
    
#     oof_preds = np.zeros((len(train), ), dtype=np.float32)
#     preds = []
#     for fold in range(num_fold):
#         tr_x = train.filter(pl.col('folds')!=fold).select(use_col).to_numpy()
#         tr_y = train.filter(pl.col('folds')!=fold).select('y').to_numpy()
#         va_x = train.filter(pl.col('folds')==fold).select(use_col).to_numpy()
#         va_y = train.filter(pl.col('folds')==fold).select('y').to_numpy()

#         params = {
#             'n_estimators' : 10000,
#             'learning_rate': 0.01,
#             'random_seed': seed,
#         }
#         clf = lgbm.LGBMClassifier(**params)
#         clf.fit(
#             tr_x, tr_y,
#             eval_set=[(va_x, va_y)],
#             early_stopping_rounds=100,
#             verbose=100)

#         va_preds_p = clf.predict_proba(va_x)[:, 1]
#         oof_preds[
#             train.select(
#                 pl.when(pl.col('folds')==fold).then(True).otherwise(False)
#                 ).to_numpy().reshape(-1)
#                 ] = va_preds_p
#         va_preds = (va_preds_p > 0.5).astype(int)
#         score = accuracy_score(va_y, va_preds)
#         print(f'Fold : {fold+1} Accuracy score: {score}')
#         print()
#         test_preds_p = clf.predict_proba(test_x)[:, 1]
#         preds.append(test_preds_p)

#     score_s = accuracy_score(train.select('y').to_numpy(), oof_preds > 0.5)
#     print(fontstyle.apply(f'Seed{seed} Accuracy score : {score_s}', 'BLACK/BOLD'))
#     print()
#     whole_va_preds.append(oof_preds)
#     whole_test_preds.append(preds)

# # preds_va_p = np.mean(whole_va_preds, axis=0)
# # whole_score = accuracy_score(train.select('y').to_numpy(), preds_va_p > 0.5)
# # preds_test = (np.mean(np.mean(whole_test_preds, axis=0), axis=0) > 0.5).astype(int)
# preds_va = np.array([np.where(preds > 0.5, 1, 0) for preds in whole_va_preds])
# whole_score = accuracy_score(train.select('y').to_numpy(), stats.mode(preds_va, axis=0).mode.flatten())
# test_preds_array = np.array(whole_test_preds)
# test_preds_array = test_preds_array.reshape(test_preds_array.shape[0]*test_preds_array.shape[1], -1)
# preds_test = np.array([np.where(preds > 0.5, 1, 0) for preds in test_preds_array])
# preds_test = stats.mode(preds_test, axis=0).mode.flatten()
# print()  
# print(fontstyle.apply(f'whole Accuracy score: {whole_score}', 'BLACK/BOLD'))
# print()

# display(pl.Series(preds_test).value_counts())