|
import os |
|
from pathlib import Path |
|
from typing import Union |
|
|
|
import pandas as pd |
|
import torch.cuda |
|
from catboost import CatBoostRegressor |
|
from catboost.utils import get_gpu_device_count |
|
from easydict import EasyDict as edict |
|
|
|
from src.cross_validate import CrossValidation |
|
from src.feature_extractors.bert_pretrain_extractor import \ |
|
ManyBertPretrainFeatureExtractor |
|
from src.feature_extractors.text_statistics_extractor import \ |
|
HandcraftedTextFeatureExtractor |
|
from src.solutions.base_solution import BaseSolution |
|
from src.solutions.constant_predictor import load_train_test_df |
|
from src.spell_checker import SmartSpellChecker |
|
from src.text_preprocessings.spellcheck_preprocessing import \ |
|
SpellcheckTextPreprocessor |
|
from src.utils import get_x_columns, seed_everything, validate_x, validate_y |
|
|
|
seed_everything() |
|
|
|
spellcheck = SmartSpellChecker() |
|
|
|
|
|
class ManyBertWithHandcraftedFeaturePredictor(BaseSolution): |
|
|
|
def __init__( |
|
self, |
|
model_names: list, |
|
catboost_iter: int, |
|
saving_dir: str, |
|
): |
|
super(ManyBertWithHandcraftedFeaturePredictor, self).__init__() |
|
|
|
self.feature_extractor = HandcraftedTextFeatureExtractor(spellcheck) |
|
self.text_preprocessing = SpellcheckTextPreprocessor(spellcheck) |
|
self.berts = ManyBertPretrainFeatureExtractor(model_names=model_names) |
|
|
|
self.device = 'GPU' if torch.cuda.is_available() else None |
|
self.task_type = 'GPU' if get_gpu_device_count() > 0 else 'CPU' |
|
|
|
|
|
self.columns = ['cohesion', 'syntax', 'vocabulary', 'phraseology', 'grammar', 'conventions'] |
|
self.models = [ |
|
CatBoostRegressor( |
|
iterations=catboost_iter, |
|
task_type=self.task_type, |
|
verbose=True, |
|
) for _ in range(len(self.columns)) |
|
] |
|
|
|
def transform_data(self, X: pd.Series) -> pd.DataFrame: |
|
cleaned_text = self.text_preprocessing.preprocess_data(X) |
|
bert_features = self.berts.generate_features(cleaned_text) |
|
handcrafted_features = self.feature_extractor.generate_features(X) |
|
features_df = pd.concat([bert_features, handcrafted_features], axis='columns') |
|
return features_df |
|
|
|
def fit(self, X: pd.DataFrame, y: pd.DataFrame, **kwargs) -> None: |
|
validate_x(X) |
|
validate_y(y) |
|
|
|
features_df = self.transform_data(X.full_text) |
|
|
|
for ii, column in enumerate(self.columns): |
|
print(f"-> Training model on: {column}...") |
|
model = self.models[ii] |
|
target = y[column] |
|
torch.cuda.empty_cache() |
|
model.fit(X=features_df, y=target) |
|
torch.cuda.empty_cache() |
|
|
|
def predict(self, X: pd.DataFrame) -> pd.DataFrame: |
|
validate_x(X) |
|
|
|
features_df = self.transform_data(X.full_text) |
|
|
|
prediction = {} |
|
for ii, column in enumerate(self.columns): |
|
print(f"-> Predicting model on: {column}") |
|
model = self.models[ii] |
|
prediction[column] = model.predict(features_df) |
|
|
|
y_pred = pd.DataFrame(prediction, index=X.index) |
|
y_pred['text_id'] = X.text_id |
|
|
|
return y_pred |
|
|
|
def save(self, directory: Union[str, Path]) -> None: |
|
directory = Path(directory) |
|
if not directory.exists(): |
|
directory.mkdir(parents=True) |
|
|
|
for ii, model in enumerate(self.models): |
|
column = self.columns[ii] |
|
path = directory / f'catboost_{column}.cbm' |
|
model.save_model(str(path)) |
|
|
|
def load(self, directory: Union[str, Path]) -> None: |
|
directory = Path(directory) |
|
if not directory.is_dir(): |
|
raise OSError(f"Dir. {directory.absolute()} does not exist") |
|
|
|
for ii, model in enumerate(self.models): |
|
column = self.columns[ii] |
|
path = directory / f'catboost_{column}.cbm' |
|
model.load_model(str(path)) |
|
|
|
|
|
def main(): |
|
config = edict( |
|
dict( |
|
model_names=[ |
|
'bert-base-uncased', |
|
'bert-base-cased', |
|
'vblagoje/bert-english-uncased-finetuned-pos', |
|
'bert-base-multilingual-cased', |
|
'unitary/toxic-bert', |
|
'bert-large-uncased' |
|
], |
|
catboost_iter=5000, |
|
n_splits=5, |
|
saving_dir='checkpoints/ManyBertWithHandcraftedFeaturePredictor', |
|
) |
|
) |
|
|
|
train_df, test_df = load_train_test_df() |
|
|
|
x_columns = get_x_columns() |
|
train_x, train_y = train_df[x_columns], train_df.drop(columns=['full_text']) |
|
|
|
predictor = ManyBertWithHandcraftedFeaturePredictor( |
|
model_names=config.model_names, |
|
catboost_iter=config.catboost_iter, |
|
saving_dir=config.saving_dir, |
|
) |
|
cv = CrossValidation(saving_dir=config.saving_dir, n_splits=config.n_splits) |
|
|
|
results = cv.fit(predictor, train_x, train_y) |
|
print("CV results") |
|
print(results) |
|
|
|
print(f"CV mean: {results.iloc[len(results) - 1].mean()}") |
|
|
|
cv.save(config.saving_dir) |
|
|
|
submission_df = cv.predict(test_df) |
|
submission_path = os.path.join(config.saving_dir, "submission.csv") |
|
submission_df.to_csv(submission_path, index=False) |
|
|
|
cv_results_path = os.path.join(config.saving_dir, "cv_results.csv") |
|
results.to_csv(cv_results_path) |
|
|
|
print("Finished training!") |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
