Custom Model

A custom classifier based on the RobertaModel class from transformers.

Created on February 2022

@author: José Antonio Espinosa

class src.domain_classifier.custom_model.CustomClassificationHead(hidden_dropout_prob=0.1, hidden_size=768)

Bases: Module

Copy of class RobertaClassificationHead (from transformers.models.roberta.modeling_roberta) Head for sentence-level classification tasks.

__init__(hidden_dropout_prob=0.1, hidden_size=768)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(features: Tensor)
Parameters:

features (Tensor) – The encoded text

class src.domain_classifier.custom_model.CustomDataset(df)

Bases: Dataset

__init__(df)
class src.domain_classifier.custom_model.CustomEncoderLayer(hidden_act='gelu', hidden_size=768, intermediate_size=3072, layer_norm_eps=1e-05, num_attention_heads=12, num_hidden_layers=1)

Bases: Module

Custom encoder layer of transformer for classification.

__init__(hidden_act='gelu', hidden_size=768, intermediate_size=3072, layer_norm_eps=1e-05, num_attention_heads=12, num_hidden_layers=1)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(features: Tensor, mask: Tensor)
Parameters:
  • features (Tensor) – The sequence to the encoder

  • mask (Tensor) – The mask for the src keys per batch

class src.domain_classifier.custom_model.CustomModel(config, path_model, model_type, model_name)

Bases: Module

Copy of class RobertaClassificationHead (from transformers.models.roberta.modeling_roberta) Head for sentence-level classification tasks.

__init__(config, path_model, model_type, model_name)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

create_data_loader(df, batch_size=8, shuffle=True)

Creates a DataLoader from a DataFrame to train/eval model

eval_model(df_eval, device='cuda', batch_size=8)

Evaluate trained model

Parameters:
  • df_train (DataFrame) – Training dataframe

  • epochs (int) – Number of epochs to train model

  • device (str, optional (default=”cuda”)) – If “cuda”, a GPU is used if available

  • batch_size (int, optiona (default=8)) – Batch size

forward(features: Tensor, mask: Tensor)
Parameters:
  • features (Tensor) – The sequence to the encoder

  • mask (Tensor) – The mask for the src keys per batch

load_embeddings()

Load embeddings layer.

If there is no previous configuration, copy it from simpletransformers ClassificationModel and save it.

train_model(df_train, device='cuda', batch_size=8)

Train the model

Parameters:
  • df_train (DataFrame) – Training dataframe

  • epochs (int) – Number of epochs to train model

  • device (str, optional (default=”cuda”)) – If “cuda”, a GPU is used if available

  • batch_size (int, optiona (default=8)) – Batch size