class MeanRibosomeLoadPipeline(Pipeline):
"""
Mean-ribosome-load pipeline for 5'UTR models.
The pipeline accepts raw nucleotide sequences and returns a sequence-level mean ribosome load (MRL) score.
"""
def preprocess(
self,
inputs: str | Mapping[str, Any],
return_tensors: str | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
**preprocess_parameters,
) -> dict[str, GenericTensor]:
if return_tensors is None:
return_tensors = "pt"
sequence, extra_inputs = _resolve_inputs(inputs)
tokenizer_kwargs = _tokenizer_kwargs(tokenizer_kwargs)
model_inputs = self.tokenizer(sequence, return_tensors=return_tensors, **tokenizer_kwargs)
model_inputs.update(extra_inputs)
return model_inputs
def _forward(self, model_inputs):
outputs = _call_model(self.model, model_inputs)
outputs["input_ids"] = model_inputs["input_ids"]
outputs["attention_mask"] = model_inputs.get("attention_mask")
return outputs
def postprocess(self, model_outputs):
input_ids = model_outputs["input_ids"]
attention_mask = model_outputs.get("attention_mask")
sequences = _decode_sequences(self.tokenizer, input_ids, attention_mask)
scores = _get_scores(model_outputs)
channels = _output_channels(self.model, scores)
results = [
_mean_ribosome_load_result(sequence, sample_scores, channels)
for sequence, sample_scores in zip(sequences, _sample_tensors(scores, len(sequences)))
]
if len(results) == 1:
return results[0]
return results
def _sanitize_parameters(
self,
tokenizer_kwargs: dict[str, Any] | None = None,
):
preprocess_params: dict[str, Any] = {}
if tokenizer_kwargs is not None:
preprocess_params["tokenizer_kwargs"] = tokenizer_kwargs
return preprocess_params, {}, {}
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if not isinstance(self.model, torch.nn.Module):
raise NotImplementedError("Only PyTorch is supported for mean-ribosome-load.")