class MethylationPipeline(Pipeline):
"""
DNA methylation pipeline.
The pipeline accepts raw DNA sequences and returns sequence-level methylation scores, such as per-cell CpG
methylation probabilities.
"""
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 = _resolve_sequence(inputs)
tokenizer_kwargs = _tokenizer_kwargs(tokenizer_kwargs)
return self.tokenizer(sequence, return_tensors=return_tensors, **tokenizer_kwargs)
def _forward(self, model_inputs):
outputs = 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 = _processed_scores(model_outputs, model=self.model, task=self.task)
channels = _output_channels(self.model, scores)
results = [
_methylation_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 methylation.")