# ---
# 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
# ---

# # Notebook to test a ranker based on scikit-learn models

# %load_ext autoreload
# %autoreload 2

# ## Load data

# +
from data import load_data

train, val, test = load_data()
# -

# Load proposals

from postprocessing import read_proposals_files

files = [
    "BART-base-submission-23.csv", 
    "BART-large-submission-21.csv",
    "pegasus-large-submission-30.csv",
    "t5-base-submission-40.csv",
    "t5-small-submission-36.csv",
]

train["name_proposals"] = read_proposals_files([f"train_{file}" for file in files])
train

val["name_proposals"] = read_proposals_files([f"val_{file}" for file in files])
val

# ## Data preprocessing

# +
from itertools import chain

all_train_texts = list(train["name"].values) + list(train["description"].values) + list(chain(*train["name_proposals"].values))
all_train_texts = set(all_train_texts)
# -

# ## Unroll training and validation data

from postprocessing import unroll_data, unroll_test_data

unrolled_train = unroll_data(train["name"], train["description"], train["name_proposals"])
unrolled_train = unrolled_train[unrolled_train["label"] == 1].append(unrolled_train[unrolled_train["label"] == 0].sample(frac=0.1)).sample(frac=1)
unrolled_train

unrolled_val = unroll_data(val["name"], val["description"], val["name_proposals"])
unrolled_val

# ## Model

# +
from sklearn.base import TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import HashingVectorizer
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import make_scorer, roc_auc_score
from sklearn.model_selection import PredefinedSplit, RandomizedSearchCV
from sklearn.pipeline import Pipeline

from xgboost import XGBClassifier

# +
from scipy.sparse import hstack

class NgramJoinTransformer(TransformerMixin):
    def fit(self, X, y=None):
        return self
    
    def transform(self, X, y=None):
        assert X.shape[0] > 0
        assert X.shape[1] % 2 == 0
        n_features = X.shape[1] // 2
        new_features = X[:, :n_features].multiply(X[:, n_features:])
        #return hstack([X, new_features])
        return new_features


# +
vectorizer = HashingVectorizer(analyzer="char", ngram_range=(1, 7), binary=True, alternate_sign=False)

feature_generator = Pipeline([
    ('vectorizer', ColumnTransformer([
        ("name_vectorizer", vectorizer, "name"),
        ("description_vectorizer", vectorizer, "description")
        ])
    ),
    ('ngram_join', NgramJoinTransformer())
])
# -

# %%time
XF_train = feature_generator.fit_transform(unrolled_train[["name", "description"]])

# %%time
XF_val = feature_generator.transform(unrolled_val[["name", "description"]])

# +
# %%time
from scipy.sparse import vstack

base_estimator = XGBClassifier(n_jobs=4, use_label_encoder=False)
param_grid = {
    'n_estimators': [10, 20, 50, 100, 200, 500, 1000],
    'learning_rate': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
    'gamma': [0, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3],
    'lambda': [0, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3],
    'alpha': [0, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2, 1e3],
    'max_depth': [3, 4, 5, 6, 7, 8, 9, 10],
    'min_child_weight': [1, 5, 10, 20, 50, 100],
    'subsample': [0.6, 0.8, 1.0],
    'colsample_bytree': [0.6, 0.8, 1.0],
}
model = RandomizedSearchCV(
    base_estimator, 
    param_grid, 
    n_iter=50, 
    scoring=make_scorer(roc_auc_score), 
    verbose=10, 
    n_jobs=5,
    cv=PredefinedSplit([-1] * len(unrolled_train) + [0] * len(unrolled_val)),
    refit=False
)
model.fit(vstack([XF_train, XF_val]), list(unrolled_train["label"].values) + list(unrolled_val["label"].values))
# -

print(model.best_params_)

best_estimator = XGBClassifier(n_jobs=20, use_label_encoder=False, **model.best_params_).fit(XF_train, unrolled_train["label"].values)

# +
import pickle

with open("ranker.pkl", "wb") as f:
    pickle.dump(best_estimator, f)
# -

preds = best_estimator.predict_proba(XF_val)
print(roc_auc_score(unrolled_val["label"], preds[:, 1]))

preds = best_estimator.predict_proba(XF_train)
print(roc_auc_score(unrolled_train["label"], preds[:, 1]))

# ## Sort predictions

submissions = ["submission_21.csv", "submission_23.csv", "submission_30.csv", "submission_40.csv", "submission_36.csv"]
test["name_proposals"] = read_proposals_files(submissions, remove_duplicates=True)
test

unrolled_test = unroll_test_data(test["description"], test["name_proposals"])
unrolled_test

# %%time
XF_test = feature_generator.transform(unrolled_test[["name", "description"]])

# %%time
unrolled_test["score"] = best_estimator.predict_proba(XF_test)[:, 1]
unrolled_test

# +
from postprocessing import reroll_score_test_data

proposals = reroll_score_test_data(unrolled_test)
proposals

# +
from postprocessing import save_submission

save_submission(proposals.values, "submission_44")
save_submission(proposals.values, "submission_44", zip=False)
# -

# ## Results

# |Model|Train AUC|Val AUC|
# |-----|---------|-------|
# |XGB 10 trees|0.576601653865533|0.5696578534991624|
# |XGB 100 trees|0.6027082266547432|0.6885317734719824|
# |XGB 100 trees|0.7342786432389221|0.605825065246392|



StarcoderdataJupyterScriptsDedupFiltered