File size: 3,590 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | from typing import Dict, Optional, Union
import torch
import transformers
from lightning import LightningModule
from torch import nn
from src.data.loading.components.interfaces import ItemData
from src.models.components.interfaces import OneKeyPerPredictionOutput
class SemanticEmbeddingInferenceModule(LightningModule):
def __init__(
self,
semantic_embedding_model: Union[nn.Module, transformers.PreTrainedModel],
semantic_embedding_model_input_map: Optional[Dict[str, str]] = None,
**kwargs,
) -> None:
"""
Initialize the SemanticEmbeddingInferenceModule.
This module is used to compute semantic embeddings from input data using a
pre-trained, frozen semantic embedding model. It is intended to be used only for
inference.
Args:
semantic_embedding_model: The model to use for computing semantic embeddings.
semantic_embedding_model_input_map: The mapping from feature names to input names
expected by the semantic embedding model.
"""
super().__init__()
self.semantic_embedding_model = semantic_embedding_model
# We use a frozen embedding module to compute the input embeddings
for param in self.semantic_embedding_model.parameters():
param.requires_grad = False
self.semantic_embedding_model_input_map = semantic_embedding_model_input_map
def forward(self, model_input: ItemData) -> torch.Tensor:
"""
Get the semantic embeddings from the input data.
Args:
model_input: ItemData consisting of the batch of input features.
Returns:
semantic_embeddings: The semantic embeddings.
Shape (batch_size, n_features)
"""
semantic_embedding_model_input_name_to_feature = {
input_embedding_model_input_name: model_input.transformed_features[
feature_name
]
for input_embedding_model_input_name, feature_name in self.semantic_embedding_model_input_map.items()
}
with torch.no_grad():
semantic_embeddings = self.semantic_embedding_model(
**semantic_embedding_model_input_name_to_feature
)
return semantic_embeddings
def model_step(self, model_input: ItemData) -> torch.Tensor:
semantic_embeddings = self.forward(model_input)
return semantic_embeddings
def predict_step(self, batch: ItemData) -> OneKeyPerPredictionOutput:
"""
Perform a single prediction step on a batch of data.
Save the semantic embeddings of the input items and the corresponding item ids
in a OneKeyAcrossPredictionsOutput object.
Args:
batch: A batch of data of ItemData type.
batch_idx: The index of the batch.
Returns:
model_output: A SharedKeyAcrossPredictionsOutput object containing the item
ids as keys and the semantic embeddings as predictions.
"""
semantic_embeddings = self.model_step(batch)
item_ids = [
item_id.item() if isinstance(item_id, torch.Tensor) else item_id
for item_id in batch.item_ids
]
model_output = OneKeyPerPredictionOutput(
keys=item_ids,
predictions=semantic_embeddings,
key_name="item_id",
prediction_name="embedding",
)
return model_output
|