Source code for autogluon.tabular.models.text_prediction.text_prediction_v1_model
"""Wrapper of the Text Prediction in AutoGluon Text."""
from typing import Optional
import logging
import time
import copy
import os
import numpy as np
import pandas as pd
from autogluon.core.constants import REGRESSION, BINARY
from autogluon.core.features.types import R_OBJECT, R_INT, R_FLOAT, R_CATEGORY, \
S_TEXT_NGRAM, S_TEXT_AS_CATEGORY, S_TEXT_SPECIAL, S_IMAGE_PATH
from autogluon.core.utils import get_cpu_count, get_gpu_count_mxnet, try_import_mxnet, try_import_autogluon_text
from autogluon.core.models import AbstractModel
logger = logging.getLogger(__name__)
[docs]class TextPredictorModel(AbstractModel):
nn_model_name = 'text_nn'
def __init__(self, **kwargs):
"""Wrapper of autogluon.text.TextPredictor.
The features can be a mix of
- text column
- categorical column
- numerical column
The labels can be categorical or numerical.
Parameters
----------
path
The directory to store the modeling outputs.
name
Name of subdirectory inside path where model will be saved.
problem_type
Type of problem that this model will handle.
Valid options: ['binary', 'multiclass', 'regression'].
eval_metric
The evaluation metric.
num_classes
The number of classes.
stopping_metric
The stopping metric.
model
The internal model object.
hyperparameters
The hyperparameters of the model
features
Names of the features.
feature_metadata
The feature metadata.
"""
super().__init__(**kwargs)
self._label_column_name = None
self._load_model = None # Whether to load inner model when loading.
def _get_default_auxiliary_params(self) -> dict:
default_auxiliary_params = super()._get_default_auxiliary_params()
extra_auxiliary_params = dict(
get_features_kwargs=dict(
valid_raw_types=[R_INT, R_FLOAT, R_CATEGORY, R_OBJECT],
invalid_special_types=[S_TEXT_NGRAM, S_TEXT_AS_CATEGORY, S_TEXT_SPECIAL, S_IMAGE_PATH],
),
)
default_auxiliary_params.update(extra_auxiliary_params)
return default_auxiliary_params
@classmethod
def _get_default_ag_args(cls) -> dict:
default_ag_args = super()._get_default_ag_args()
extra_ag_args = {'valid_stacker': False}
default_ag_args.update(extra_ag_args)
return default_ag_args
def _set_default_params(self):
super()._set_default_params()
try_import_autogluon_text()
from autogluon.text import ag_text_presets
self.params = ag_text_presets.create('default')
def _fit(self,
X: pd.DataFrame,
y: pd.Series,
X_val: Optional[pd.DataFrame] = None,
y_val: Optional[pd.Series] = None,
time_limit: Optional[int] = None,
sample_weight=None,
**kwargs):
"""The internal fit function
Parameters
----------
X
Features of the training dataset
y
Labels of the training dataset
X_val
Features of the validation dataset
y_val
Labels of the validation dataset
time_limit
The time limits for the fit function
kwargs
Other keyword arguments
"""
try_import_mxnet()
try_import_autogluon_text()
from autogluon.text import TextPredictor
# Decide name of the label column
if 'label' in X.columns:
label_col_id = 0
while True:
self._label_column_name = 'label{}'.format(label_col_id)
if self._label_column_name not in X.columns:
break
label_col_id += 1
else:
self._label_column_name = 'label'
X_train = self.preprocess(X, fit=True)
if X_val is not None:
X_val = self.preprocess(X_val)
# Get arguments from kwargs
verbosity = kwargs.get('verbosity', 2)
num_cpus = kwargs.get('num_cpus', None)
num_gpus = kwargs.get('num_gpus', None)
if sample_weight is not None: # TODO: support
logger.log(15, "sample_weight not yet supported for TextPredictorModel, this model will ignore them in training.")
X_train.insert(len(X_train.columns), self._label_column_name, y)
if X_val is not None:
X_val.insert(len(X_val.columns), self._label_column_name, y_val)
assert self.params['tune_kwargs']['num_trials'] == 1 \
or self.params['tune_kwargs']['num_trials'] is None,\
'Currently, you cannot nest the hyperparameter search in text neural network ' \
'and the AutoGluon Tabular.'
verbosity_text = max(0, verbosity - 1)
root_logger = logging.getLogger()
root_log_level = root_logger.level
self.model = TextPredictor(label=self._label_column_name,
problem_type=self.problem_type,
path=self.path,
eval_metric=self.eval_metric,
verbosity=verbosity_text)
self.model.fit(train_data=X_train,
tuning_data=X_val,
time_limit=time_limit,
num_gpus=num_gpus,
num_cpus=num_cpus,
hyperparameters=self.params,
seed=self.params.get('seed', 0))
self.model.set_verbosity(verbosity)
root_logger.setLevel(root_log_level) # Reset log level
def save(self, path: str = None, verbose=True) -> str:
self._load_model = self.model is not None
__model = self.model
self.model = None
# save this AbstractModel object without NN weights
path = super().save(path=path, verbose=verbose)
self.model = __model
if self._load_model:
text_nn_path = os.path.join(path, self.nn_model_name)
self.model.save(text_nn_path)
logger.log(15, f"\tSaved Text NN weights and model hyperparameters to '{text_nn_path}'.")
self._load_model = None
return path
@classmethod
def load(cls, path: str, reset_paths=True, verbose=True):
model = super().load(path=path, reset_paths=reset_paths, verbose=verbose)
if model._load_model:
try_import_autogluon_text()
from autogluon.text import TextPredictor
model.model = TextPredictor.load(os.path.join(path, cls.nn_model_name))
model._load_model = None
return model
def get_memory_size(self) -> int:
"""Return the memory size by calculating the total number of parameters.
Returns
-------
memory_size
The total memory size in bytes.
"""
total_size = 0
for k, v in self.model._model.net.collect_params().items():
total_size += np.dtype(v.dtype).itemsize * np.prod(v.shape)
return total_size
def _get_default_resources(self):
num_cpus = get_cpu_count()
# TODO: use get_gpu_count_torch() or some better way once torch models are available.
num_gpus = get_gpu_count_mxnet()
return num_cpus, num_gpus
def _predict_proba(self, X, **kwargs):
X = self.preprocess(X, **kwargs)
if self.problem_type == REGRESSION:
return self.model.predict(X, as_pandas=False)
y_pred_proba = self.model.predict_proba(X, as_pandas=False)
return self._convert_proba_to_unified_form(y_pred_proba)