-#
-# Copyright (c) 2021, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import copy
-import inspect
-import os
-import pathlib
-from collections import defaultdict
-from types import SimpleNamespace
-from typing import Callable, Dict, Iterable, List, Optional, Type, Union, cast
-
-import numpy as np
-import torch
-import torchmetrics as tm
-from merlin.models.utils.registry import camelcase_to_snakecase
-from merlin.schema import ColumnSchema
-from merlin.schema import Schema as Core_Schema
-from merlin.schema import Tags
-from tqdm import tqdm
-from transformers.modeling_utils import SequenceSummary
-
-from merlin_standard_lib import Schema
-
-from ..block.base import BlockBase, BlockOrModule, BlockType
-from ..features.base import InputBlock
-from ..features.sequence import TabularFeaturesType
-from ..typing import TabularData
-from ..utils.padding import pad_inputs
-from ..utils.torch_utils import LossMixin, MetricsMixin
-
-
-def name_fn(name, inp):
- return "/".join([name, inp]) if name else None
-
-
-[docs]class PredictionTask(torch.nn.Module, LossMixin, MetricsMixin):
-
"""Individual prediction-task of a model.
-
Parameters
-
----------
-
loss: torch.nn.Module
-
The loss to use during training of this task.
-
metrics: torch.nn.Module
-
The metrics to calculate during training & evaluation.
-
target_name: str, optional
-
Name of the target, this is needed when there are multiple targets.
-
task_name: str, optional
-
Name of the prediction task, if not provided a name will be automatically constructed based
-
on the target-name & class-name.
-
forward_to_prediction_fn: Callable[[torch.Tensor], torch.Tensor]
-
Function to apply before the prediction
-
task_block: BlockType
-
Module to transform input tensor before computing predictions.
-
pre: BlockType
-
Module to compute the predictions probabilities.
-
summary_type: str
-
This is used to summarize a sequence into a single tensor. Accepted values are:
-
- `"last"` -- Take the last token hidden state (like XLNet)
-
- `"first"` -- Take the first token hidden state (like Bert)
-
- `"mean"` -- Take the mean of all tokens hidden states
-
- `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
-
- `"attn"` -- Not implemented now, use multi-head attention
-
"""
-
-
def __init__(
-
self,
-
loss: torch.nn.Module,
-
metrics: Iterable[tm.Metric] = None,
-
target_name: Optional[str] = None,
-
task_name: Optional[str] = None,
-
forward_to_prediction_fn: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
-
task_block: Optional[BlockType] = None,
-
pre: Optional[BlockType] = None,
-
summary_type: str = "last",
-
):
-
super().__init__()
-
self.summary_type = summary_type
-
self.sequence_summary = SequenceSummary(
-
SimpleNamespace(summary_type=self.summary_type) # type: ignore
-
) # noqa
-
self.target_name = target_name
-
self.forward_to_prediction_fn = forward_to_prediction_fn
-
self.set_metrics(metrics)
-
self.loss = loss
-
self.pre = pre
-
self.task_block = task_block
-
self._task_name = task_name
-
-
[docs] def build(
-
self,
-
body: BlockType,
-
input_size,
-
inputs: Optional[InputBlock] = None,
-
device=None,
-
task_block: Optional[BlockType] = None,
-
pre=None,
-
):
-
"""
-
The method will be called when block is converted to a model,
-
i.e when linked to prediction head.
-
Parameters
-
----------
-
block:
-
the model block to link with head
-
device:
-
set the device for the metrics and layers of the task
-
"""
-
-
if task_block:
-
# TODO: What to do when `self.task_block is not None`?
-
self.task_block = task_block
-
if pre:
-
# TODO: What to do when `self.pre is not None`?
-
self.pre = pre
-
-
# Build task block
-
pre_input_size = input_size
-
if self.task_block:
-
if isinstance(self.task_block, torch.nn.Module):
-
self.task_block = copy.deepcopy(self.task_block)
-
else:
-
self.task_block = self.task_block.build(input_size)
-
pre_input_size = self.task_block.output_size() # type: ignore
-
-
if self.pre:
-
if isinstance(self.pre, torch.nn.Module):
-
self.pre = copy.deepcopy(self.pre)
-
else:
-
self.pre = self.pre.build(pre_input_size)
-
-
if device:
-
self.to(device)
-
for metric in self.metrics:
-
metric.to(device)
-
self.built = True
-
-
[docs] def forward(
-
self,
-
inputs: torch.Tensor,
-
targets: torch.Tensor = None,
-
training: bool = False,
-
testing: bool = False,
-
):
-
x = inputs
-
-
if len(x.size()) == 3 and self.summary_type:
-
x = self.sequence_summary(x)
-
-
if self.task_block:
-
x = self.task_block(x) # type: ignore
-
-
if self.pre:
-
x = self.pre(x) # type: ignore
-
-
if training or testing:
-
# add support of computing the loss inside the forward
-
# and return a dictionary as standard output
-
if self.summary_type is None:
-
if targets.dim() != 2:
-
raise ValueError(
-
"If `summary_type==None`, targets are expected to be a 2D tensor, "
-
f"but got a tensor with shape {targets.shape}"
-
)
-
-
loss = self.loss(x, target=targets)
-
return {"loss": loss, "labels": targets, "predictions": x}
-
-
return x
-
-
@property
-
def task_name(self):
-
if self._task_name:
-
return self._task_name
-
-
base_name = camelcase_to_snakecase(self.__class__.__name__)
-
-
return name_fn(self.target_name, base_name) if self.target_name else base_name
-
-
[docs] def child_name(self, name):
-
return name_fn(self.task_name, name)
-
-
[docs] def set_metrics(self, metrics):
-
self.metrics = torch.nn.ModuleList(metrics)
-
-
[docs] def calculate_metrics( # type: ignore
-
self,
-
predictions: torch.Tensor,
-
targets: torch.Tensor,
-
) -> Dict[str, torch.Tensor]:
-
outputs = {}
-
-
predictions = self.forward_to_prediction_fn(cast(torch.Tensor, predictions))
-
-
from .prediction_task import BinaryClassificationTask
-
-
for metric in self.metrics:
-
if isinstance(metric, tuple(type(x) for x in BinaryClassificationTask.DEFAULT_METRICS)):
-
targets = cast(torch.Tensor, targets).int()
-
outputs[self.metric_name(metric)] = metric(predictions, targets)
-
-
return outputs
-
-
[docs] def compute_metrics(self, **kwargs):
-
return {self.metric_name(metric): metric.compute() for metric in self.metrics}
-
-
[docs] def metric_name(self, metric: tm.Metric) -> str:
-
return self.child_name(camelcase_to_snakecase(metric.__class__.__name__))
-
-
[docs] def reset_metrics(self):
-
for metric in self.metrics:
-
metric.reset()
-
-
[docs] def to_head(self, body, inputs=None, **kwargs) -> "Head":
-
return Head(body, self, inputs=inputs, **kwargs)
-
-
[docs] def to_model(self, body, inputs=None, **kwargs) -> "Model":
-
return Model(Head(body, self, inputs=inputs, **kwargs), **kwargs)
-
-
-[docs]class Head(torch.nn.Module, LossMixin, MetricsMixin):
-
"""Head of a Model, a head has a single body but could have multiple prediction-tasks.
-
Parameters
-
----------
-
body: Block
-
TODO
-
prediction_tasks: Union[List[PredictionTask], PredictionTask], optional
-
TODO
-
task_blocks
-
TODO
-
task_weights: List[float], optional
-
TODO
-
loss_reduction: str, default="mean"
-
TODO
-
inputs: TabularFeaturesType, optional
-
TODO
-
"""
-
-
def __init__(
-
self,
-
body: BlockBase,
-
prediction_tasks: Union[List[PredictionTask], PredictionTask],
-
task_blocks: Optional[Union[BlockType, Dict[str, BlockType]]] = None,
-
task_weights: Optional[List[float]] = None,
-
loss_reduction: str = "mean",
-
inputs: Optional[TabularFeaturesType] = None,
-
):
-
super().__init__()
-
self.body = body
-
self.loss_reduction = loss_reduction
-
self.prediction_task_dict = torch.nn.ModuleDict()
-
if prediction_tasks:
-
if not isinstance(prediction_tasks, list):
-
prediction_tasks = [prediction_tasks]
-
for i, task in enumerate(prediction_tasks):
-
self.prediction_task_dict[task.task_name] = task
-
-
self._task_weights = defaultdict(lambda: 1.0)
-
if task_weights:
-
for task, val in zip(cast(List[PredictionTask], prediction_tasks), task_weights):
-
self._task_weights[task.task_name] = val
-
-
self.build(inputs=inputs, task_blocks=task_blocks)
-
-
[docs] def build(self, inputs=None, device=None, task_blocks=None):
-
"""Build each prediction task that's part of the head.
-
Parameters
-
----------
-
body
-
inputs
-
device
-
task_blocks
-
"""
-
if not getattr(self.body, "output_size", lambda: None)():
-
raise ValueError(
-
"Can't infer output-size of the body, please provide "
-
"a `Block` with a output-size. You can wrap any torch.Module in a Block."
-
)
-
-
input_size = self.body.output_size()
-
-
if device:
-
self.to(device)
-
-
for name, task in self.prediction_task_dict.items():
-
task_block = task_blocks
-
if task_blocks and isinstance(task_blocks, dict) and name in task_blocks:
-
task_block = task_blocks[name]
-
task.build(self.body, input_size, inputs=inputs, device=device, task_block=task_block)
-
self.input_size = input_size
-
-
[docs] @classmethod
-
def from_schema(
-
cls,
-
schema: Schema,
-
body: BlockBase,
-
task_blocks: Optional[Union[BlockType, Dict[str, BlockType]]] = None,
-
task_weight_dict: Optional[Dict[str, float]] = None,
-
loss_reduction: str = "mean",
-
inputs: Optional[TabularFeaturesType] = None,
-
) -> "Head":
-
"""Instantiate a Head from a Schema through tagged targets.
-
Parameters
-
----------
-
schema: DatasetSchema
-
Schema to use for inferring all targets based on the tags.
-
body
-
task_blocks
-
task_weight_dict
-
loss_reduction
-
inputs
-
Returns
-
-------
-
Head
-
"""
-
task_weight_dict = task_weight_dict or {}
-
tasks: List[PredictionTask] = []
-
task_weights = []
-
-
from .prediction_task import BinaryClassificationTask, RegressionTask
-
-
for binary_target in schema.select_by_tag([Tags.BINARY, Tags.CLASSIFICATION]).column_names:
-
tasks.append(BinaryClassificationTask(binary_target))
-
task_weights.append(task_weight_dict.get(binary_target, 1.0))
-
-
for regression_target in schema.select_by_tag(Tags.REGRESSION).column_names:
-
tasks.append(RegressionTask(regression_target))
-
task_weights.append(task_weight_dict.get(regression_target, 1.0))
-
-
# TODO: Add multi-class classification here. Figure out how to get number of classes
-
-
return cls(
-
body,
-
tasks,
-
task_blocks=task_blocks,
-
task_weights=task_weights,
-
loss_reduction=loss_reduction,
-
inputs=inputs,
-
)
-
-
[docs] def pop_labels(self, inputs: TabularData) -> TabularData:
-
"""Pop the labels from the different prediction_tasks from the inputs.
-
Parameters
-
----------
-
inputs: TabularData
-
Input dictionary containing all targets.
-
Returns
-
-------
-
TabularData
-
"""
-
outputs = {}
-
for name in self.prediction_task_dict.keys():
-
outputs[name] = inputs.pop(name)
-
-
return outputs
-
-
[docs] def forward(
-
self,
-
body_outputs: Union[torch.Tensor, TabularData],
-
training: bool = False,
-
testing: bool = False,
-
targets: Union[torch.Tensor, TabularData] = None,
-
call_body: bool = False,
-
top_k: Optional[int] = None,
-
**kwargs,
-
) -> Union[torch.Tensor, TabularData]:
-
outputs = {}
-
-
from transformers4rec.torch.model.prediction_task import NextItemPredictionTask
-
-
if call_body:
-
body_outputs = self.body(body_outputs, training=training, testing=testing, **kwargs)
-
-
if training or testing:
-
losses = []
-
labels = {}
-
predictions = {}
-
for name, task in self.prediction_task_dict.items():
-
if isinstance(targets, dict):
-
label = targets.get(task.target_name, None)
-
else:
-
label = targets
-
if label is not None:
-
label = label.float()
-
task_output = task(
-
body_outputs, targets=label, training=training, testing=testing, **kwargs
-
)
-
labels[name] = task_output["labels"]
-
predictions[name] = task_output["predictions"]
-
losses.append(task_output["loss"] * self._task_weights[name])
-
loss_tensor = torch.stack(losses)
-
loss = getattr(loss_tensor, self.loss_reduction)()
-
outputs = {"loss": loss, "labels": labels, "predictions": predictions}
-
else:
-
for name, task in self.prediction_task_dict.items():
-
if isinstance(task, NextItemPredictionTask):
-
outputs[name] = task(
-
body_outputs,
-
targets=targets,
-
training=training,
-
testing=testing,
-
top_k=top_k,
-
**kwargs,
-
)
-
-
else:
-
outputs[name] = task(
-
body_outputs, targets=targets, training=training, testing=testing, **kwargs
-
)
-
-
return outputs
-
-
[docs] def calculate_metrics( # type: ignore
-
self,
-
predictions: Union[torch.Tensor, TabularData],
-
targets: Union[torch.Tensor, TabularData],
-
) -> Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]]:
-
"""Calculate metrics of the task(s) set in the Head instance.
-
Parameters
-
----------
-
predictions: Union[torch.Tensor, TabularData]
-
The predictions tensors to use for calculate metrics.
-
They can be either a torch.Tensor if a single task is used or
-
a dictionary of torch.Tensor if multiple tasks are used. In the
-
second case, the dictionary is indexed by the tasks names.
-
targets:
-
The tensor or dictionary of targets to use for computing the metrics of
-
one or multiple tasks.
-
"""
-
metrics = {}
-
-
for name, task in self.prediction_task_dict.items():
-
label = targets
-
output = predictions
-
if isinstance(targets, dict):
-
# The labels are retrieved from the task's output
-
# and indexed by the task name.
-
label = targets[name]
-
if isinstance(predictions, dict):
-
output = predictions[name]
-
-
metrics.update(
-
task.calculate_metrics(
-
predictions=output,
-
targets=label,
-
)
-
)
-
-
return _output_metrics(metrics)
-
-
[docs] def compute_metrics(self, mode: str = None) -> Dict[str, Union[float, torch.Tensor]]:
-
def name_fn(x):
-
return "_".join([mode, x]) if mode else x
-
-
metrics = {
-
name_fn(name): task.compute_metrics()
-
for name, task in self.prediction_task_dict.items()
-
}
-
-
return _output_metrics(metrics)
-
-
[docs] def reset_metrics(self):
-
""""""
-
for task in self.prediction_task_dict.values():
-
task.reset_metrics()
-
-
@property
-
def task_blocks(self) -> Dict[str, Optional[BlockOrModule]]:
-
return {name: task.task_block for name, task in self.prediction_task_dict.items()}
-
-
[docs] def to_model(self, **kwargs) -> "Model":
-
"""Convert the head to a Model.
-
Returns
-
-------
-
Model
-
"""
-
-
return Model(self, **kwargs)
-
-
-[docs]class Model(torch.nn.Module, LossMixin, MetricsMixin):
-
def __init__(
-
self,
-
*head: Head,
-
head_weights: Optional[List[float]] = None,
-
head_reduction: str = "mean",
-
optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
-
name: str = None,
-
max_sequence_length: Optional[int] = None,
-
top_k: Optional[int] = None,
-
):
-
"""Model class that can aggregate one or multiple heads.
-
Parameters
-
----------
-
head: Head
-
One or more heads of the model.
-
head_weights: List[float], optional
-
Weight-value to use for each head.
-
head_reduction: str, optional
-
How to reduce the losses into a single tensor when multiple heads are used.
-
optimizer: Type[torch.optim.Optimizer]
-
Optimizer-class to use during fitting
-
name: str, optional
-
Name of the model.
-
max_sequence_length: int, optional
-
The maximum sequence length supported by the model.
-
Used to truncate sequence inputs longer than this value.
-
top_k: int, optional
-
The number of items to return at the inference step once the model is deployed.
-
Default is None, which will return all items.
-
"""
-
if head_weights:
-
if not isinstance(head_weights, list):
-
raise ValueError("`head_weights` must be a list")
-
if not len(head_weights) == len(head):
-
raise ValueError(
-
"`head_weights` needs to have the same length " "as the number of heads"
-
)
-
-
super().__init__()
-
-
self.name = name
-
self.heads = torch.nn.ModuleList(head)
-
self.head_weights = head_weights or [1.0] * len(head)
-
self.head_reduction = head_reduction
-
self.optimizer = optimizer
-
self.max_sequence_length = max_sequence_length
-
self.top_k = top_k
-
-
[docs] def forward(self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs):
-
# Convert inputs to float32 which is the default type, expected by PyTorch
-
for name, val in inputs.items():
-
if torch.is_floating_point(val):
-
inputs[name] = val.to(torch.float32)
-
-
# pad ragged inputs
-
inputs = pad_inputs(inputs, self.max_sequence_length)
-
-
if isinstance(targets, dict) and len(targets) == 0:
-
# `pyarrow`` dataloader is returning {} instead of None
-
# TODO remove this code when `PyarraowDataLoader` is dropped
-
targets = None
-
-
# TODO: Optimize this
-
if training or testing:
-
losses = []
-
labels = {}
-
predictions = {}
-
for i, head in enumerate(self.heads):
-
head_output = head(
-
inputs,
-
call_body=True,
-
targets=targets,
-
training=training,
-
testing=testing,
-
**kwargs,
-
)
-
labels.update(head_output["labels"])
-
predictions.update(head_output["predictions"])
-
losses.append(head_output["loss"] * self.head_weights[i])
-
loss_tensor = torch.stack(losses)
-
loss = getattr(loss_tensor, self.head_reduction)()
-
if len(labels) == 1:
-
labels = list(labels.values())[0]
-
predictions = list(predictions.values())[0]
-
return {"loss": loss, "labels": labels, "predictions": predictions}
-
else:
-
outputs = {}
-
for head in self.heads:
-
outputs.update(
-
head(
-
inputs,
-
call_body=True,
-
targets=targets,
-
training=training,
-
testing=testing,
-
top_k=self.top_k,
-
**kwargs,
-
)
-
)
-
if len(outputs) == 1:
-
return list(outputs.values())[0]
-
-
return outputs
-
-
[docs] def calculate_metrics( # type: ignore
-
self,
-
predictions: Union[torch.Tensor, TabularData],
-
targets: Union[torch.Tensor, TabularData],
-
) -> Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]]:
-
"""Calculate metrics of the task(s) set in the Head instance.
-
Parameters
-
----------
-
predictions: Union[torch.Tensor, TabularData]
-
The predictions tensors returned by the model.
-
They can be either a torch.Tensor if a single task is used or
-
a dictionary of torch.Tensor if multiple heads/tasks are used. In the
-
second case, the dictionary is indexed by the tasks names.
-
targets:
-
The tensor or dictionary of targets returned by the model.
-
They are used for computing the metrics of one or multiple tasks.
-
"""
-
outputs = {}
-
for head in self.heads:
-
outputs.update(
-
head.calculate_metrics(
-
predictions,
-
targets,
-
)
-
)
-
-
return outputs
-
-
[docs] def compute_metrics(self, mode=None) -> Dict[str, Union[float, torch.Tensor]]:
-
metrics = {}
-
for head in self.heads:
-
metrics.update(head.compute_metrics(mode=mode))
-
-
return metrics
-
-
[docs] def reset_metrics(self):
-
for head in self.heads:
-
head.reset_metrics()
-
-
[docs] def to_lightning(self):
-
import pytorch_lightning as pl
-
-
parent_self = self
-
-
class BlockWithHeadLightning(pl.LightningModule):
-
def __init__(self):
-
super(BlockWithHeadLightning, self).__init__()
-
self.parent = parent_self
-
-
def forward(self, inputs, targets=None, training=False, testing=False, *args, **kwargs):
-
return self.parent(
-
inputs, targets=targets, training=training, testing=testing, *args, **kwargs
-
)
-
-
def training_step(self, batch, batch_idx, targets=None, training=True, testing=False):
-
loss = self.parent(*batch, targets=targets, training=training, testing=testing)[
-
"loss"
-
]
-
self.log("train_loss", loss)
-
-
return loss
-
-
def configure_optimizers(self):
-
optimizer = self.parent.optimizer(self.parent.parameters(), lr=1e-3)
-
-
return optimizer
-
-
return BlockWithHeadLightning()
-
-
[docs] def fit(
-
self,
-
dataloader,
-
optimizer=torch.optim.Adam,
-
eval_dataloader=None,
-
num_epochs=1,
-
amp=False,
-
train=True,
-
verbose=True,
-
compute_metric=True,
-
):
-
if isinstance(dataloader, torch.utils.data.DataLoader):
-
dataset = dataloader.dataset
-
else:
-
dataset = dataloader
-
-
if inspect.isclass(optimizer):
-
optimizer = optimizer(self.parameters())
-
-
self.train(mode=train)
-
epoch_losses = []
-
with torch.set_grad_enabled(mode=train):
-
for epoch in range(num_epochs):
-
losses = []
-
batch_iterator = enumerate(iter(dataset))
-
if verbose:
-
batch_iterator = tqdm(batch_iterator)
-
for batch_idx, (x, y) in batch_iterator:
-
if amp:
-
with torch.cuda.amp.autocast():
-
output = self(x, targets=y, training=True)
-
else:
-
output = self(x, targets=y, training=True)
-
losses.append(float(output["loss"]))
-
if compute_metric:
-
self.calculate_metrics(
-
output["predictions"],
-
targets=output["labels"],
-
)
-
if train:
-
optimizer.zero_grad()
-
output["loss"].backward()
-
optimizer.step()
-
if verbose:
-
print(self.compute_metrics(mode="train"))
-
if eval_dataloader:
-
print(self.evaluate(eval_dataloader, verbose=False))
-
epoch_losses.append(np.mean(losses))
-
-
return np.array(epoch_losses)
-
-
[docs] def evaluate(
-
self, dataloader, targets=None, training=False, testing=True, verbose=True, mode="eval"
-
):
-
if isinstance(dataloader, torch.utils.data.DataLoader):
-
dataset = dataloader.dataset
-
else:
-
dataset = dataloader
-
-
batch_iterator = enumerate(iter(dataset))
-
if verbose:
-
batch_iterator = tqdm(batch_iterator)
-
self.reset_metrics()
-
for batch_idx, (x, y) in batch_iterator:
-
output = self(x, targets=y, training=training, testing=testing)
-
self.calculate_metrics(
-
output["predictions"],
-
targets=output["labels"],
-
)
-
-
return self.compute_metrics(mode=mode)
-
-
def _get_name(self):
-
if self.name:
-
return self.name
-
-
return super(Model, self)._get_name()
-
-
@property
-
def input_schema(self):
-
# return the input schema given by the model
-
# loop over the heads to get input schemas
-
schemas = []
-
for head in self.heads:
-
schemas.append(head.body.inputs.schema)
-
if all(isinstance(s, Core_Schema) for s in schemas):
-
return sum(schemas, Core_Schema())
-
-
model_schema = sum(schemas, Schema())
-
-
# TODO: rework T4R to use Merlin Schemas.
-
# In the meantime, we convert model_schema to merlin core schema
-
core_schema = Core_Schema()
-
for column in model_schema:
-
name = column.name
-
-
dtype = {0: np.float32, 2: np.int64, 3: np.float32}[column.type]
-
tags = column.tags
-
dims = None
-
if column.value_count.max > 0:
-
dims = (None, (column.value_count.min, column.value_count.max))
-
int_domain = {"min": column.int_domain.min, "max": column.int_domain.max}
-
properties = {
-
"int_domain": int_domain,
-
}
-
-
col_schema = ColumnSchema(
-
name, dtype=dtype, tags=tags, properties=properties, dims=dims
-
)
-
core_schema[name] = col_schema
-
return core_schema
-
-
@property
-
def output_schema(self):
-
from merlin.schema import Tags
-
-
from .prediction_task import BinaryClassificationTask, RegressionTask
-
-
# if the model has one head with one task, the output is a tensor
-
# if multiple heads and/or multiple prediction task, the output is a dictionary
-
output_cols = []
-
for head in self.heads:
-
dims = None
-
for name, task in head.prediction_task_dict.items():
-
target_dim = task.target_dim
-
int_domain = {"min": target_dim, "max": target_dim}
-
if (
-
isinstance(task, (BinaryClassificationTask, RegressionTask))
-
and not task.summary_type
-
):
-
dims = (None, (1, None))
-
elif (
-
isinstance(task, (BinaryClassificationTask, RegressionTask))
-
and task.summary_type
-
):
-
dims = (None,)
-
else:
-
dims = (None, task.target_dim)
-
properties = {
-
"int_domain": int_domain,
-
}
-
# in case one sets top_k at the inference step we return two outputs
-
if self.top_k:
-
# be sure categ item-id dtype in model.input schema and output schema matches
-
col_name = self.input_schema.select_by_tag(Tags.ITEM_ID).column_names[0]
-
col_dtype = (
-
self.input_schema.select_by_tag(Tags.ITEM_ID)
-
.column_schemas[col_name]
-
.dtype.name
-
)
-
col_schema_scores = ColumnSchema(
-
"item_id_scores", dtype=np.float32, properties=properties, dims=dims
-
)
-
col_schema_ids = ColumnSchema(
-
"item_ids", dtype=np.dtype(col_dtype), properties=properties, dims=dims
-
)
-
output_cols.append(col_schema_scores)
-
output_cols.append(col_schema_ids)
-
else:
-
col_schema = ColumnSchema(
-
name, dtype=np.float32, properties=properties, dims=dims
-
)
-
output_cols.append(col_schema)
-
-
return Core_Schema(output_cols)
-
-
@property
-
def prediction_tasks(self):
-
return [task for head in self.heads for task in list(head.prediction_task_dict.values())]
-
-
[docs] def save(self, path: Union[str, os.PathLike], model_name="t4rec_model_class"):
-
"""Saves the model to f"{export_path}/{model_name}.pkl" using `cloudpickle`
-
Parameters
-
----------
-
path : Union[str, os.PathLike]
-
Path to the directory where the T4Rec model should be saved.
-
model_name : str, optional
-
the name given to the pickle file storing the T4Rec model,
-
by default 't4rec_model_class'
-
"""
-
try:
-
import cloudpickle
-
except ImportError:
-
raise ValueError("cloudpickle is required to save model class")
-
-
export_path = pathlib.Path(path)
-
export_path.mkdir(exist_ok=True)
-
-
model_name = model_name + ".pkl"
-
export_path = export_path / model_name
-
with open(export_path, "wb") as out:
-
cloudpickle.dump(self, out)
-
-
[docs] @classmethod
-
def load(cls, path: Union[str, os.PathLike], model_name="t4rec_model_class") -> "Model":
-
"""Loads a T4Rec model that was saved with `model.save()`.
-
Parameters
-
----------
-
path : Union[str, os.PathLike]
-
Path to the directory where the T4Rec model is saved.
-
model_name : str, optional
-
the name given to the pickle file storing the T4Rec model,
-
by default 't4rec_model_class'.
-
"""
-
try:
-
import cloudpickle
-
except ImportError:
-
raise ValueError("cloudpickle is required to load T4Rec model")
-
-
export_path = pathlib.Path(path)
-
model_name = model_name + ".pkl"
-
export_path = export_path / model_name
-
return cloudpickle.load(open(export_path, "rb"))
-
-
-def _output_metrics(metrics):
- # If there is only a single head with metrics, returns just those metrics
- if len(metrics) == 1 and isinstance(metrics[list(metrics.keys())[0]], dict):
- return metrics[list(metrics.keys())[0]]
-
- return metrics
-
