inference/inference API

dnallm.inference.inference

DNA Language Model Inference Module.

This module implements core model inference functionality, including:

1. DNAInference class
2. Model loading and initialization
3. Batch sequence inference
4. Result post-processing
5. Device management

6. Half-precision inference support

7. Core features:

8. Model state management
9. Batch inference
10. Result merging
11. Inference result saving

12. Memory optimization

13. Inference optimization:

14. Batch parallelization
15. GPU acceleration
16. Half-precision computation
17. Memory efficiency optimization

Example

inference_engine = DNAInference(
    model=model,
    tokenizer=tokenizer,
    config=config
)
results = inference_engine.infer(sequences)

Classes

DNAInference

DNAInference(
    model, tokenizer, config, lora_adapter=None, **kwargs
)

DNA sequence inference engine using fine-tuned models.

This class provides comprehensive functionality for performing inference using DNA language models. It handles model loading, inference, result processing, and various output formats including hidden states and attention weights for model interpretability.

Attributes:

Name	Type	Description
model		Fine-tuned model instance for inference
tokenizer		Tokenizer for encoding DNA sequences
task_config		Configuration object containing task settings
pred_config		Configuration object containing inference parameters
device		Device (CPU/GPU/MPS) for model inference
sequences	list[str]	List of input sequences
labels	list[Any]	List of true labels (if available)
embeddings	list[Any]	Dictionary containing hidden states and attention weights

Initialize the inference engine.

Parameters:

Name	Type	Description	Default
model	Any	Fine-tuned model instance for inference	required
tokenizer	Any	Tokenizer for encoding DNA sequences	required
config	dict	Configuration dictionary containing task settings and inference parameters	required
lora_adapter	str | None	Optional path to LoRA adapter for model	None

Source code in dnallm/inference/inference.py

def __init__(
    self,
    model: Any,
    tokenizer: Any,
    config: dict,
    lora_adapter: str | None = None,
    **kwargs,
) -> None:
    """Initialize the inference engine.

    Args:
        model: Fine-tuned model instance for inference
        tokenizer: Tokenizer for encoding DNA sequences
        config: Configuration dictionary containing task settings\
                and inference parameters
        lora_adapter: Optional path to LoRA adapter for model
    """

    default_forward_args = {
        "input_ids",
        "attention_mask",
        "token_type_ids",
        "position_ids",
        "inputs_embeds",
        "labels",
        "output_attentions",
        "output_hidden_states",
        "return_dict",
        "past_key_values",
        "use_cache",
    }

    if lora_adapter:
        from peft import PeftModel
        from ..models.model import peft_forward_compatiable

        if os.path.isdir(lora_adapter):
            source = "local"
        else:
            source = model.source if hasattr(model, "source") else "huggingface"
        try:
            lora_adapter_path, _ = _get_model_path_and_imports(lora_adapter, source)
        except Exception as e:
            raise ValueError(f"Failed to load LoRA adapter from {lora_adapter}: {e}") from e

        if model is not None:
            self.accepted_args = self._get_accepted_forward_args(model)
        else:
            self.accepted_args = set(default_forward_args)

        model = peft_forward_compatiable(model)
        self.model = PeftModel.from_pretrained(model, lora_adapter_path)
        logger.info(f"Loaded LoRA adapter from {lora_adapter}")
    else:
        self.model = model
        if model is not None:
            if "CustomEvo" in str(type(self.model)):
                self.accepted_args = self._get_accepted_forward_args(model.model)
            else:
                self.accepted_args = self._get_accepted_forward_args(model)
        else:
            self.accepted_args = set(default_forward_args)
    self.tokenizer = tokenizer
    self.config = config
    self.pad_id = self._get_pad_id()
    self.task_config = config["task"]
    self.pred_config = config["inference"]
    self.device = self._get_device()
    if model:
        if "CustomEvo" in str(type(self.model)):
            self.model.model.to(self.device)
        else:
            self.model.to(self.device)
        # mamba only support cuda and cpu, and only allow fp32
        if "mamba" in str(type(self.model)).lower():
            if self.device.type != "cuda":
                self.device = torch.device("cpu")
            if self.pred_config.use_fp16:
                self.pred_config.use_fp16 = False
        logger.info(f"Using device: {self.device}")
    self.sequences: list[str] = []
    self.labels: list[Any] = []

Methods:

batch_infer

batch_infer(
    dataloader,
    do_pred=True,
    output_hidden_states=False,
    output_attentions=False,
    reduce_hidden_states=False,
    reduce_strategy="mean",
    return_dict=True,
)

Perform batch inference on sequences.

This method runs inference on batches of sequences and optionally extracts hidden states and attention weights for model interpretability.

Parameters:

Name	Type	Description	Default
dataloader	DataLoader	DataLoader object containing sequences for inference	required
do_pred	bool	Whether to convert logits to predictions	True
output_hidden_states	bool	Whether to output hidden states from all layers	False
output_attentions	bool	Whether to output attention weights from all layers	False
reduce_hidden_states	bool	Whether to average hidden states across layers	False

Returns:

Type	Description
tuple[Tensor, dict | None, dict]	Tuple containing: - torch.Tensor: All logits from the model - Optional[Dict]: Predictions dictionary if do_pred=True, otherwise None - Dict: Embeddings dictionary containing hidden states and/or attention weights

Note

Setting output_hidden_states or output_attentions to True will consume significant memory, especially for long sequences or large models.

Source code in dnallm/inference/inference.py

@torch.inference_mode()
def batch_infer(
    self,
    dataloader: DataLoader,
    do_pred: bool = True,
    output_hidden_states: bool = False,
    output_attentions: bool = False,
    reduce_hidden_states: bool = False,
    reduce_strategy: str | int = "mean",
    return_dict: bool = True,
) -> tuple[torch.Tensor, dict | None, dict]:
    """Perform batch inference on sequences.

    This method runs inference on batches of sequences and optionally
    extracts hidden states and attention weights for model
    interpretability.

    Args:
        dataloader: DataLoader object containing sequences for inference
        do_pred: Whether to convert logits to predictions
        output_hidden_states: Whether to output hidden states from
            all layers
        output_attentions: Whether to output attention weights from
            all layers
        reduce_hidden_states: Whether to average hidden states across
            layers

    Returns:
        Tuple containing:
            - torch.Tensor: All logits from the model
            - Optional[Dict]: Predictions dictionary if do_pred=True,
              otherwise None
            - Dict: Embeddings dictionary containing hidden states
              and/or attention weights

    Note:
        Setting output_hidden_states or output_attentions to True will
        consume significant memory, especially for long sequences or
        large models.
    """
    # Set model to evaluation mode
    self.model.eval()
    all_logits = []

    # Setup configurations for outputs
    output_hidden_states, hidden_embeddings, params = self._setup_hidden_states_config(
        output_hidden_states
    )
    output_attentions, attention_embeddings = self._setup_attentions_config(
        output_attentions, params
    )

    # Combine embeddings dictionaries
    embeddings = {**hidden_embeddings, **attention_embeddings}

    # Check model precision settings
    if self.pred_config.use_fp16:
        dtype = torch.float16
    elif self.pred_config.use_bf16:
        dtype = torch.bfloat16
    else:
        dtype = torch.float32

    # Iterate over batches
    for batch in tqdm(dataloader, desc="Inferring"):
        inputs = {k: v.to(self.device) if hasattr(v, "to") else v for k, v in batch.items()}
        # Add output flags if supported
        # In case model config does not recognize these args
        if output_attentions:
            if "output_attentions" in self.accepted_args:
                inputs["output_attentions"] = True
            elif "**kwargs" in self.accepted_args:
                inputs["output_attentions"] = True
        if output_hidden_states:
            if "output_hidden_states" in self.accepted_args:
                inputs["output_hidden_states"] = True
            elif "**kwargs" in self.accepted_args:
                inputs["output_hidden_states"] = True

        # Run model inference
        # check accepted forward method
        args = inputs.keys()
        accepted_inputs = {}
        for arg in args:
            if arg in self.accepted_args or "**kwargs" in self.accepted_args:
                accepted_inputs[arg] = inputs[arg]

        # Use autocast for mixed precision if enabled
        if self.pred_config.use_fp16 or self.pred_config.use_bf16:
            with torch.amp.autocast("cuda", dtype=dtype):
                outputs = self.model(**accepted_inputs)
        else:
            outputs = self.model(**accepted_inputs)

        # Process batch outputs
        logits = self._process_batch_outputs(
            outputs,
            inputs,
            output_hidden_states,
            output_attentions,
            embeddings,
            reduce_hidden_states,
            reduce_strategy,
        )
        all_logits.append(logits)

    # Concatenate all logits
    if all_logits and all_logits[0] is not None:
        all_logits = torch.cat(all_logits, dim=0)  # type: ignore

    # Finalize embeddings
    self._finalize_embeddings(embeddings, output_hidden_states, output_attentions)

    # Get predictions if requested
    predictions = None
    if do_pred and len(all_logits) > 0:
        predictions = self.logits_to_preds(all_logits)  # type: ignore
        if return_dict:
            predictions = self.format_output(predictions)  # type: ignore

    return all_logits, predictions, embeddings  # type: ignore

calculate_metrics

calculate_metrics(logits, labels, plot=False)

Calculate evaluation metrics for model predictions.

This method computes task-specific evaluation metrics using the configured metrics computation module.

Parameters:

Name	Type	Description	Default
logits	list | Tensor	Model predictions (logits or probabilities)	required
labels	list | Tensor	True labels for evaluation	required
plot	bool	Whether to generate metric plots	False

Returns:

Type	Description
dict[Any, Any]	Dictionary containing evaluation metrics for the task

Source code in dnallm/inference/inference.py

def calculate_metrics(
    self,
    logits: list | torch.Tensor,
    labels: list | torch.Tensor,
    plot: bool = False,
) -> dict[Any, Any]:
    """Calculate evaluation metrics for model predictions.

    This method computes task-specific evaluation metrics using the
    configured metrics computation module.

    Args:
        logits: Model predictions (logits or probabilities)
        labels: True labels for evaluation
        plot: Whether to generate metric plots

    Returns:
        Dictionary containing evaluation metrics for the task
    """
    # Calculate metrics based on task type
    compute_metrics_func = compute_metrics(self.task_config, plot=plot)
    metrics: dict[Any, Any] = compute_metrics_func((logits, labels))

    return metrics

estimate_memory_usage

estimate_memory_usage(batch_size=1, sequence_length=1000)

Estimate memory usage for inference.

Parameters:

Name	Type	Description	Default
batch_size	int	Batch size for inference	1
sequence_length	int	Maximum sequence length	1000

Returns:

Type	Description
dict[str, Any]	Dict containing memory usage estimates

Source code in dnallm/inference/inference.py

def estimate_memory_usage(
    self, batch_size: int = 1, sequence_length: int = 1000
) -> dict[str, Any]:
    """Estimate memory usage for inference.

    Args:
        batch_size: Batch size for inference
        sequence_length: Maximum sequence length

    Returns:
        Dict containing memory usage estimates
    """
    try:
        # Get model parameters
        total_params = sum(p.numel() for p in self.model.parameters())
        param_memory_mb = (total_params * 4) / (1024 * 1024)  # Assuming float32

        # Estimate activation memory (rough approximation)
        if hasattr(self.model, "config"):
            config = self.model.config
            hidden_size = getattr(config, "hidden_size", 768)
            num_layers = getattr(config, "num_hidden_layers", 12)
        else:
            hidden_size, num_layers = 768, 12

        # Rough estimate for activations
        activation_memory_mb = (batch_size * sequence_length * hidden_size * num_layers * 2) / (
            1024 * 1024
        )

        total_memory_mb = param_memory_mb + activation_memory_mb

        return {
            "total_estimated_mb": f"{total_memory_mb:.1f}",
            "parameter_memory_mb": f"{param_memory_mb:.1f}",
            "activation_memory_mb": f"{activation_memory_mb:.1f}",
            "note": "Estimates are approximate and may vary based on actual usage",
        }
    except Exception as e:
        return {"error": str(e)}

force_eager_attention

force_eager_attention()

Force the model to use eager attention implementation.

This method attempts to switch the model from SDPA to eager attention implementation to ensure compatibility with output_attentions=True.

Returns:

Name	Type	Description
bool	bool	True if successful, False otherwise

Source code in dnallm/inference/inference.py

def force_eager_attention(self) -> bool:
    """Force the model to use eager attention implementation.

    This method attempts to switch the model from SDPA to eager attention
    implementation to ensure compatibility with output_attentions=True.

    Returns:
        bool: True if successful, False otherwise
    """
    try:
        if hasattr(self.model, "config") and hasattr(self.model.config, "attn_implementation"):
            self.model.config.attn_implementation = "eager"
            logger.success("Switched to eager attention implementation")
            return True
    except Exception as e:
        logger.failure(f"Failed to switch to eager attention: {e}")
    return False

format_output

format_output(predictions)

Format output predictions into a structured dictionary.

This method converts raw predictions into a user-friendly format with sequences, labels, and confidence scores.

Parameters:

Name	Type	Description	Default
predictions	tuple[Tensor, list]	Tuple containing (probabilities, labels)	required

Returns:

Type	Description
dict	Dictionary containing formatted predictions with structure:
dict	{index: {'sequence': str, 'label': str/list, 'scores': dict/list}}

Source code in dnallm/inference/inference.py

def format_output(self, predictions: tuple[torch.Tensor, list]) -> dict:
    """Format output predictions into a structured dictionary.

    This method converts raw predictions into a user-friendly format with
    sequences, labels, and confidence scores.

    Args:
        predictions: Tuple containing (probabilities, labels)

    Returns:
        Dictionary containing formatted predictions with structure:
        {index: {'sequence': str, 'label': str/list, 'scores': dict/list}}
    """
    # Get task type from config
    task_type = self.task_config.task_type
    formatted_predictions = {}
    probs, labels = predictions
    probs = probs.numpy().tolist()
    keep_seqs = True if len(self.sequences) else False
    label_names = self.task_config.label_names
    for i, label in enumerate(labels):
        prob = probs[i]
        if task_type == "regression":
            scores = {label_names[0]: prob}
        elif task_type == "token":
            scores = [max(x) for x in prob]  # type: ignore
        else:
            scores = {label_names[j]: p for j, p in enumerate(prob)}
        formatted_predictions[i] = {
            "sequence": self.sequences[i] if keep_seqs else "",
            "label": label,
            "scores": scores,
        }
    return formatted_predictions

generate

generate(
    inputs,
    n_tokens=400,
    n_samples=1,
    temperature=1.0,
    top_k=4,
    top_p=1.0,
    batched=True,
)

Generate DNA sequences using the model.

This function performs sequence generation tasks using the loaded model, currently supporting CausalLM and EVO2 models for DNA sequence generation.

Parameters:

Name	Type	Description	Default
inputs	DataLoader | list[str]	DataLoader or List containing prompt sequences	required
n_tokens	int	Number of tokens to generate, default 400	400
n_samples	int	Do samples n times	1
temperature	float	Sampling temperature for generation, default 1.0	1.0
top_k	int	Top-k sampling parameter, default 4	4
top_p	float	Top-p sampling paramether, default 1	1.0
batched	bool	Do batched generation	True

Returns:

Type	Description
dict[Any, Any]	Dictionary containing generated sequences

Note

Currently only supports Causal language models for sequence generation

Source code in dnallm/inference/inference.py

def generate(
    self,
    inputs: DataLoader | list[str],
    n_tokens: int = 400,
    n_samples: int = 1,
    temperature: float = 1.0,
    top_k: int = 4,
    top_p: float = 1.0,
    batched: bool = True,
) -> dict[Any, Any]:
    """Generate DNA sequences using the model.

    This function performs sequence generation tasks using the loaded
    model, currently supporting CausalLM and EVO2 models for
    DNA sequence generation.

    Args:
        inputs: DataLoader or List containing prompt sequences
        n_tokens: Number of tokens to generate, default 400
        n_samples: Do samples n times
        temperature: Sampling temperature for generation, default 1.0
        top_k: Top-k sampling parameter, default 4
        top_p: Top-p sampling paramether, default 1
        batched: Do batched generation

    Returns:
        Dictionary containing generated sequences

    Note:
        Currently only supports Causal language models
        for sequence generation
    """
    # Prepare prompt sequences
    prompt_seqs: list[str] = []
    if isinstance(inputs, DataLoader):
        for data in tqdm(inputs, desc="Generating"):
            seqs = data["sequence"]
            if isinstance(prompt_seqs, list):
                seqs.extend([seq for seq in seqs if seq])
            if not seqs:
                continue
    else:
        prompt_seqs = inputs
    # Check if model supports generation
    if "evo2" in str(self.model).lower():
        # Generate sequences
        outputs = self.model.generate(
            prompt_seqs=prompt_seqs,
            n_tokens=n_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            batched=batched,
            cached_generation=True,
        )
        formatted_outputs = []
        for i, seq in enumerate(prompt_seqs):
            generated_seqs = outputs.sequences[i]
            scores = outputs.logprobs_mean[i]
            formatted_outputs.append({
                "Prompt": seq,
                "Output": generated_seqs,
                "Score": scores,
            })
        return formatted_outputs  # type: ignore
    elif "evo1" in str(self.model).lower():
        from evo import generate

        model = self.model.model
        tokenizer = self.tokenizer
        # Generate sequences
        outputs = generate(
            prompt_seqs * n_samples,
            model=model,
            tokenizer=tokenizer.raw_tokenizer,
            n_tokens=n_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            cached_generation=True,
            batched=batched,
            device=self.device,
            verbose=1,
        )
        formatted_outputs = []
        for i, seq in enumerate(prompt_seqs):
            generated_seqs = outputs[0][i]
            scores = outputs[1][i]
            formatted_outputs.append({
                "Prompt": seq,
                "Output": generated_seqs,
                "Score": scores,
            })
        return formatted_outputs  # type: ignore
    elif "megadna" in str(self.model).lower():
        model = self.model
        tokenizer = self.tokenizer
        formatted_outputs = []
        for seq in prompt_seqs:
            for _ in range(n_samples):
                input_ids = tokenizer(seq, return_tensors="pt").to(self.device)["input_ids"]
                output = model.generate(
                    input_ids,
                    seq_len=n_tokens,
                    temperature=temperature,
                    filter_thres=top_p,
                )
                decoded = tokenizer.decode(output.squeeze().cpu().int())
                formatted_outputs.append({
                    "Prompt": seq,
                    "Output": decoded.replace(" ", ""),
                })
        return formatted_outputs  # type: ignore
    elif "causallm" in str(self.model).lower() or "lmhead" in str(self.model).lower():
        outputs = []
        # Tokenize prompt sequences
        for seq in prompt_seqs:
            inputs = self.tokenizer(seq, return_tensors="pt").to(self.device)
            output = self.model.generate(  # type: ignore
                **inputs,
                max_new_tokens=n_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=True,
            )
            decoded = self.tokenizer.decode(output[0], skip_special_tokens=True)
            outputs.append({
                "Prompt": seq,
                "Output": decoded.replace(" ", ""),
            })
        return outputs  # type: ignore
    else:
        raise ValueError("This model is not supported for sequence generation.")

    return {}  # type: ignore[unreachable]

generate_dataset

generate_dataset(
    seq_or_path,
    batch_size=1,
    seq_col="sequence",
    label_col="labels",
    sep=None,
    fasta_sep="|",
    multi_label_sep=None,
    uppercase=False,
    lowercase=False,
    sampling=None,
    keep_seqs=True,
    padding=True,
    do_encode=True,
)

Generate dataset from sequences or file path.

This method creates a DNADataset and DataLoader from either a list of sequences or a file path, supporting various file formats and preprocessing options.

Parameters:

Name	Type	Description	Default
seq_or_path	str | list[str]	Single sequence, list of sequences, or path to a file containing sequences	required
batch_size	int	Batch size for DataLoader	1
seq_col	str	Column name for sequences in the file	'sequence'
label_col	str	Column name for labels in the file	'labels'
sep	str | None	Delimiter for CSV, TSV, or TXT files	None
fasta_sep	str	Delimiter for FASTA files	'|'
multi_label_sep	str | None	Delimiter for multi-label sequences	None
uppercase	bool	Whether to convert sequences to uppercase	False
lowercase	bool	Whether to convert sequences to lowercase	False
sampling	float | None	Fraction of data to randomly sample for inference	None
keep_seqs	bool	Whether to keep sequences in the dataset for later use	True
padding	str | bool	Padding strategy for encoding sequences	True
do_encode	bool	Whether to encode sequences for the model	True

Returns:

Type	Description
tuple[DNADataset, DataLoader]	Tuple containing: - DNADataset: Dataset object with sequences and labels - DataLoader: DataLoader object for batch processing

Raises:

Type	Description
ValueError	If input is neither a file path nor a list of sequences

Source code in dnallm/inference/inference.py

def generate_dataset(
    self,
    seq_or_path: str | list[str],
    batch_size: int = 1,
    seq_col: str = "sequence",
    label_col: str = "labels",
    sep: str | None = None,
    fasta_sep: str = "|",
    multi_label_sep: str | None = None,
    uppercase: bool = False,
    lowercase: bool = False,
    sampling: float | None = None,
    keep_seqs: bool = True,
    padding: str | bool = True,
    do_encode: bool = True,
) -> tuple[DNADataset, DataLoader]:
    """Generate dataset from sequences or file path.

    This method creates a DNADataset and DataLoader from either a list
    of sequences or a file path, supporting various file formats and
    preprocessing options.

    Args:
        seq_or_path: Single sequence, list of sequences, or path to a
            file containing sequences
        batch_size: Batch size for DataLoader
        seq_col: Column name for sequences in the file
        label_col: Column name for labels in the file
        sep: Delimiter for CSV, TSV, or TXT files
        fasta_sep: Delimiter for FASTA files
        multi_label_sep: Delimiter for multi-label sequences
        uppercase: Whether to convert sequences to uppercase
        lowercase: Whether to convert sequences to lowercase
        sampling: Fraction of data to randomly sample for inference
        keep_seqs: Whether to keep sequences in the dataset for later use
        padding: Padding strategy for encoding sequences
        do_encode: Whether to encode sequences for the model

    Returns:
        Tuple containing:
            - DNADataset: Dataset object with sequences and labels
            - DataLoader: DataLoader object for batch processing

    Raises:
        ValueError: If input is neither a file path nor a list of sequences
    """
    # Initialize dataset to None to avoid unbound variable issues
    dataset = None

    if isinstance(seq_or_path, str):
        suffix = seq_or_path.split(".")[-1]
        if suffix and os.path.isfile(seq_or_path):
            sequences = []
            dataset = DNADataset.load_local_data(
                seq_or_path,
                seq_col=seq_col,
                label_col=label_col,
                sep=sep,
                fasta_sep=fasta_sep,
                multi_label_sep=multi_label_sep,
                tokenizer=self.tokenizer,
                max_length=self.pred_config.max_length,
            )
        else:
            sequences = [seq_or_path]
    elif isinstance(seq_or_path, list):
        sequences = seq_or_path
    else:
        raise ValueError("Input should be a file path or a list of sequences.")

    # If sampling is specified, randomly sample the sequences
    if sampling:
        dataset = dataset.sampling(sampling) if dataset else None

    # Create dataset from sequences if we have any and no dataset was
    # loaded from file
    if len(sequences) > 0 and dataset is None:
        ds = Dataset.from_dict({"sequence": sequences})
        dataset = DNADataset(ds, self.tokenizer, max_length=self.pred_config.max_length)

    # Ensure dataset is not None before proceeding
    if not dataset:
        raise ValueError("No valid dataset could be created from the input.")
    # If labels are provided, keep labels
    if keep_seqs:
        self.sequences = dataset.dataset["sequence"]
    # Encode sequences
    if do_encode:
        task_type = self.task_config.task_type
        dataset.encode_sequences(
            padding=padding,  # type: ignore
            remove_unused_columns=True,
            task=task_type,
            uppercase=uppercase,
            lowercase=lowercase,
        )
        all_cols = dataset.dataset.features
        cols_drop = [c for c in all_cols if c not in self.accepted_args]
        dataset.dataset = dataset.dataset.remove_columns(cols_drop)
    else:
        all_cols = dataset.dataset.features
        # check if dataset is already encoded
        if "sequence" in all_cols:
            check_seq = dataset.dataset["sequence"][0]
            if isinstance(check_seq, torch.Tensor):
                # already encoded
                if "input_ids" not in all_cols:
                    dataset.dataset = dataset.dataset.rename_column("sequence", "input_ids")
                dataset.dataset.set_format(type="torch")
                cols_drop = [c for c in dataset.dataset.features if c not in self.accepted_args]
                dataset.dataset = dataset.dataset.remove_columns(cols_drop)
            else:
                cols_drop = [c for c in dataset.dataset.features if c not in self.accepted_args]
    # Check for labels in dataset - handle both Dataset and
    # DatasetDict cases
    if isinstance(dataset.dataset, DatasetDict):
        # For DatasetDict, check the first available split
        keys = list(dataset.dataset.keys())
        if keys and "labels" in dataset.dataset[keys[0]].features:
            self.labels = dataset.dataset[keys[0]]["labels"]
    else:
        # For single Dataset
        if "labels" in dataset.dataset.features:
            self.labels = dataset.dataset["labels"]
    # Create DataLoader
    dataloader: DataLoader = DataLoader(
        dataset,  # type: ignore[arg-type]
        batch_size=batch_size,
        num_workers=self.pred_config.num_workers,
    )

    return dataset, dataloader

get_available_outputs

get_available_outputs()

Get information about available model outputs.

Returns:

Type	Description
dict[str, Any]	Dict containing information about what outputs are available and
dict[str, Any]	collected

Source code in dnallm/inference/inference.py

def get_available_outputs(self) -> dict[str, Any]:
    """Get information about available model outputs.

    Returns:
        Dict containing information about what outputs are available and
        collected
    """
    capabilities = {
        "hidden_states_available": self._check_hidden_states_support(),
        "attentions_available": self._check_attention_support(),
        "hidden_states_collected": hasattr(self, "embeddings")
        and "hidden_states" in self.embeddings
        and self.embeddings["hidden_states"] is not None,
        "attentions_collected": hasattr(self, "embeddings")
        and "attentions" in self.embeddings
        and self.embeddings["attentions"] is not None,
    }
    return capabilities

get_embeddings

get_embeddings(
    inputs,
    do_reduce=False,
    reduce_strategy="mean",
    force=False,
)

Get embeddings from the last inference. This method performs inference on the provided inputs and extracts embeddings from the model's hidden states.

Parameters:

Name	Type	Description	Default
inputs	DataLoader | list[str] | str	DataLoader or list of sequences for inference	required
do_reduce	bool	Whether to reduce hidden states to 2D using PCA	False
reduce_strategy	str | int	Strategy to reduce hidden states ('mean', 'max', 'min', or int for center window size)	'mean'
force	bool	Whether to force re-computation of embeddings	False

Returns:

Type	Description
Any	Dict containing embeddings from the last inference

Source code in dnallm/inference/inference.py

def get_embeddings(
    self,
    inputs: DataLoader | list[str] | str,
    do_reduce: bool = False,
    reduce_strategy: str | int = "mean",
    force: bool = False,
) -> Any:
    """Get embeddings from the last inference.
    This method performs inference on the provided inputs and extracts
    embeddings from the model's hidden states.

    Args:
        inputs: DataLoader or list of sequences for inference
        do_reduce: Whether to reduce hidden states to 2D using PCA
        reduce_strategy: Strategy to reduce hidden states
            ('mean', 'max', 'min', or int for center window size)
        force: Whether to force re-computation of embeddings

    Returns:
        Dict containing embeddings from the last inference
    """
    # Initialize embeddings
    if force or not hasattr(self, "embeddings"):
        self.embeddings = {"hidden_states": None, "attention_mask": None}
    # Check specific models
    is_special = ""
    special_list = ["CustomEvo", "MEGADNA"]
    for name in special_list:
        if name in str(self.model):
            sequences = inputs
            is_special = name
            break
    if not is_special:
        if isinstance(inputs, list):
            _, dataloader = self.generate_dataset(
                inputs, batch_size=self.pred_config.batch_size
            )
        elif isinstance(inputs, str):
            # Assume it's a file path
            if os.path.isfile(inputs):
                file_path = os.path.abspath(inputs)
            else:
                raise ValueError(
                    f"Input {inputs} is not a valid file path. "
                    "Please provide a valid file path "
                    "or a list contains valid sequences."
                )
            _, dataloader = self.generate_dataset(
                file_path,
                do_encode=True,
                batch_size=self.pred_config.batch_size,
            )
        else:
            dataloader = inputs

    # Check if model supports generation
    if is_special.startswith("CustomEvo"):
        # Get model and tokenizer
        model = self.model.model
        tokenizer = self.tokenizer
        # Get layer names
        layers = []
        layer_prefix = "blocks"
        for name, _ in model.named_parameters():
            if name.startswith(layer_prefix):
                layer = layer_prefix + "." + name.split(".")[1]
                if layer not in layers:
                    layers.append(layer)
        # Get embeddings
        all_embeddings = [[] for _ in layers]  # type: ignore
        for sequence in tqdm(sequences):
            input_ids = (
                torch
                .tensor(
                    tokenizer.tokenize(sequence),
                    dtype=torch.int,
                )
                .unsqueeze(0)
                .to(self.device)
            )
            _, embeddings = self.model(input_ids, return_embeddings=True, layer_names=layers)
            for i, n in enumerate(layers):
                tmp = embeddings[n].detach().cpu().to(torch.float32)
                if do_reduce:
                    mean_emb = _compute_mean_embeddings(tmp, None).squeeze(0)
                    all_embeddings[i].append(mean_emb)
                else:
                    all_embeddings[i].append(tmp)
        for i, _ in enumerate(layers):
            all_embeddings[i] = np.stack(all_embeddings[i], axis=0)  # type: ignore[call-overload]
        if self.embeddings["hidden_states"] is None:
            self.embeddings["hidden_states"] = all_embeddings
        return all_embeddings

    elif is_special == "MEGADNA":
        model = self.model
        tokenizer = self.tokenizer
        all_embeddings = [None] * 3  # type: ignore
        out_embeddings = []
        for sequence in tqdm(sequences):
            input_ids = tokenizer(sequence, return_tensors="pt").to(self.device)["input_ids"]
            if not isinstance(input_ids, torch.LongTensor):
                input_ids = input_ids.long()
            with torch.no_grad():
                embeddings = model(input_ids, return_value="embedding")
            for i in range(len(embeddings)):
                if all_embeddings[i] is None:
                    all_embeddings[i] = []
                all_embeddings[i].append(embeddings[i].detach().cpu())
            out_embeddings = all_embeddings
        for i in range(len(all_embeddings)):
            emb = (
                np.stack(all_embeddings[i], axis=0)
                if i > 0
                else np.concatenate(all_embeddings[i], axis=0)
            )
            reshaped_emb = emb.reshape(emb.shape[0], -1, emb.shape[-1])
            if do_reduce:
                mean_emb = _compute_mean_embeddings(reshaped_emb, None)
            else:
                mean_emb = reshaped_emb
            # proj_emb = torch.nn.Linear(mean_emb.shape[-1], 128)
            all_embeddings[i] = mean_emb
        # Save embeddings
        if self.embeddings["hidden_states"] is None:
            self.embeddings["hidden_states"] = all_embeddings
        return out_embeddings

    if "hidden_states" not in self.embeddings or self.embeddings["hidden_states"] is None:
        _, _, embeddings = self.batch_infer(
            dataloader,
            do_pred=False,
            output_hidden_states=True,
            reduce_hidden_states=do_reduce,
            reduce_strategy=reduce_strategy,
        )
        self.embeddings = embeddings
    return self.embeddings["hidden_states"]

get_model_info

get_model_info()

Get information about the loaded model.

Returns:

Type	Description
dict[str, Any]	Dict containing model information including type, device, and
dict[str, Any]	attention support

Source code in dnallm/inference/inference.py

def get_model_info(self) -> dict[str, Any]:
    """Get information about the loaded model.

    Returns:
        Dict containing model information including type, device, and
        attention support
    """
    # Get basic model information
    info = self._get_basic_model_info()

    # Add model-specific configuration information
    info.update(self._get_model_config_info())

    # Add parameter information
    info["num_parameters"] = self._get_model_parameters_info()

    # Add configuration as a dictionary
    info["config"] = self._get_model_config_dict()

    return info

get_model_parameters

get_model_parameters()

Get information about model parameters.

Returns:

Type	Description
dict[str, int]	Dict containing parameter counts

Source code in dnallm/inference/inference.py

def get_model_parameters(self) -> dict[str, int]:
    """Get information about model parameters.

    Returns:
        Dict containing parameter counts
    """
    try:
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        frozen_params = total_params - trainable_params

        return {
            "total": total_params,
            "trainable": trainable_params,
            "frozen": frozen_params,
        }
    except Exception as e:
        return {"error": str(e)}  # type: ignore

infer

infer(
    sequences=None,
    file_path=None,
    evaluate=False,
    output_hidden_states=False,
    output_attentions=False,
    save_to_file=False,
    **kwargs,
)

Main inference method for sequences or files.

This is the primary entry point for performing inference. It automatically determines whether to process sequences directly or load from a file.

Parameters:

Name	Type	Description	Default
sequences	str | list[str] | None	Single sequence or list of sequences for inference	None
file_path	str | None	Path to file containing sequences for inference	None
evaluate	bool	Whether to evaluate predictions against true labels	False
output_hidden_states	bool	Whether to output hidden states for visualization	False
output_attentions	bool	Whether to output attention weights for visualization	False
save_to_file	bool	Whether to save predictions to output directory	False
**kwargs	Any	Additional arguments passed to specific inference methods	{}

Returns:

Name	Type	Description
Either	dict | tuple[dict, dict]	Dict: Dictionary containing predictions Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

Raises:

Type	Description
ValueError	If neither sequences nor file_path is provided

Source code in dnallm/inference/inference.py

def infer(
    self,
    sequences: str | list[str] | None = None,
    file_path: str | None = None,
    evaluate: bool = False,
    output_hidden_states: bool = False,
    output_attentions: bool = False,
    save_to_file: bool = False,
    **kwargs: Any,
) -> dict | tuple[dict, dict]:
    """Main inference method for sequences or files.

    This is the primary entry point for performing inference. It
    automatically determines whether to process sequences directly or
    load from a file.

    Args:
        sequences: Single sequence or list of sequences for inference
        file_path: Path to file containing sequences for inference
        evaluate: Whether to evaluate predictions against true labels
        output_hidden_states: Whether to output hidden states for
            visualization
        output_attentions: Whether to output attention weights for
            visualization
        save_to_file: Whether to save predictions to output directory
        **kwargs: Additional arguments passed to specific inference methods

    Returns:
        Either:
            - Dict: Dictionary containing predictions
            - Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

    Raises:
        ValueError: If neither sequences nor file_path is provided
    """
    if sequences is not None:
        return self.infer_seqs(
            sequences=sequences,
            evaluate=evaluate,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            save_to_file=save_to_file,
        )
    elif file_path is not None:
        return self.infer_file(
            file_path=file_path,
            evaluate=evaluate,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            save_to_file=save_to_file,
            **kwargs,
        )
    else:
        raise ValueError("Either sequences or file_path must be provided")

infer_file

infer_file(
    file_path,
    evaluate=False,
    output_hidden_states=False,
    output_attentions=False,
    seq_col="sequence",
    label_col="labels",
    sep=None,
    fasta_sep="|",
    multi_label_sep=None,
    uppercase=False,
    lowercase=False,
    sampling=None,
    do_encode=True,
    padding=True,
    save_to_file=False,
    plot_metrics=False,
)

Infer from a file containing sequences.

This method loads sequences from a file and performs inference, with optional evaluation, visualization, and saving capabilities.

Parameters:

Name	Type	Description	Default
file_path	str | Dataset	Path to the file containing sequences	required
evaluate	bool	Whether to evaluate predictions against true labels	False
output_hidden_states	bool	Whether to output hidden states for visualization	False
output_attentions	bool	Whether to output attention weights for visualization	False
seq_col	str	Column name for sequences in the file	'sequence'
label_col	str	Column name for labels in the file	'labels'
sep	str | None	Delimiter for CSV, TSV, or TXT files	None
fasta_sep	str	Delimiter for FASTA files	'|'
multi_label_sep	str | None	Delimiter for multi-label sequences	None
uppercase	bool	Whether to convert sequences to uppercase	False
lowercase	bool	Whether to convert sequences to lowercase	False
sampling	float | None	Fraction of data to randomly sample for inference	None
do_encode	bool	Whether to encode sequences for the model	True
padding	str | bool	Whether to apply padding to sequences	True
save_to_file	bool	Whether to save predictions and metrics to output directory	False
plot_metrics	bool	Whether to generate metric plots	False

Returns:

Name	Type	Description
Either	dict | tuple[dict, dict]	Dict: Dictionary containing predictions Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

Note

Setting output_attentions=True may consume significant memory

Source code in dnallm/inference/inference.py

def infer_file(
    self,
    file_path: str | Dataset,
    evaluate: bool = False,
    output_hidden_states: bool = False,
    output_attentions: bool = False,
    seq_col: str = "sequence",
    label_col: str = "labels",
    sep: str | None = None,
    fasta_sep: str = "|",
    multi_label_sep: str | None = None,
    uppercase: bool = False,
    lowercase: bool = False,
    sampling: float | None = None,
    do_encode: bool = True,
    padding: str | bool = True,
    save_to_file: bool = False,
    plot_metrics: bool = False,
) -> dict | tuple[dict, dict]:
    """Infer from a file containing sequences.

    This method loads sequences from a file and performs inference,
    with optional evaluation, visualization, and saving capabilities.

    Args:
        file_path: Path to the file containing sequences
        evaluate: Whether to evaluate predictions against true labels
        output_hidden_states: Whether to output hidden states for
            visualization
        output_attentions: Whether to output attention weights for
            visualization
        seq_col: Column name for sequences in the file
        label_col: Column name for labels in the file
        sep: Delimiter for CSV, TSV, or TXT files
        fasta_sep: Delimiter for FASTA files
        multi_label_sep: Delimiter for multi-label sequences
        uppercase: Whether to convert sequences to uppercase
        lowercase: Whether to convert sequences to lowercase
        sampling: Fraction of data to randomly sample for inference
        do_encode: Whether to encode sequences for the model
        padding: Whether to apply padding to sequences
        save_to_file: Whether to save predictions and metrics to
            output directory
        plot_metrics: Whether to generate metric plots

    Returns:
        Either:
            - Dict: Dictionary containing predictions
            - Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

    Note:
        Setting output_attentions=True may consume significant memory
    """
    # Get dataset and dataloader from file
    if isinstance(file_path, Dataset):
        dataloader = DataLoader(
            file_path,
            batch_size=self.pred_config.batch_size,
            num_workers=self.pred_config.num_workers,
        )
        self.labels = file_path["labels"] if "labels" in file_path.features else []
    else:
        _, dataloader = self.generate_dataset(
            file_path,
            seq_col=seq_col,
            label_col=label_col,
            sep=sep,
            fasta_sep=fasta_sep,
            multi_label_sep=multi_label_sep,
            uppercase=uppercase,
            lowercase=lowercase,
            sampling=sampling,
            do_encode=do_encode,
            padding=padding,
            batch_size=self.pred_config.batch_size,
        )
    # Do batch inference
    if output_attentions:
        warnings.warn(
            "Cautions: output_attentions may consume a lot of memory.\n",
            stacklevel=2,
        )
    logits, predictions, embeddings = self.batch_infer(
        dataloader,
        output_hidden_states=output_hidden_states,
        output_attentions=output_attentions,
    )
    # Keep hidden states
    if output_hidden_states or output_attentions:
        self.embeddings = embeddings
    # Save predictions
    if save_to_file and self.pred_config.output_dir:
        save_predictions(predictions, Path(self.pred_config.output_dir))  # type: ignore
    # Do evaluation
    if len(self.labels) == len(logits) and evaluate:
        metrics = self.calculate_metrics(logits, self.labels, plot=plot_metrics)
        metrics_save = dict(metrics)
        metrics_save.pop("curve", None)
        metrics_save.pop("scatter", None)
        if save_to_file and self.pred_config.output_dir:
            save_metrics(metrics, Path(self.pred_config.output_dir))
        # Whether to plot metrics
        if plot_metrics:
            return predictions, metrics  # type: ignore
        else:
            return predictions, metrics_save  # type: ignore

    return predictions  # type: ignore

infer_seqs

infer_seqs(
    sequences,
    do_pred=True,
    evaluate=False,
    output_hidden_states=False,
    output_attentions=False,
    save_to_file=False,
)

Infer for a list of sequences.

This method provides a convenient interface for performing inference on sequences, with optional evaluation and saving capabilities.

Parameters:

Name	Type	Description	Default
sequences	str | list[str]	Single sequence or list of sequences for inference	required
evaluate	bool	Whether to evaluate predictions against true labels	False
output_hidden_states	bool	Whether to output hidden states for visualization	False
output_attentions	bool	Whether to output attention weights for visualization	False
save_to_file	bool	Whether to save predictions to output directory	False

Returns:

Name	Type	Description
Either	dict | tuple[dict, dict]	Dict: Dictionary containing predictions Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

Note

Evaluation requires that labels are available in the dataset

Source code in dnallm/inference/inference.py

def infer_seqs(
    self,
    sequences: str | list[str],
    do_pred: bool = True,
    evaluate: bool = False,
    output_hidden_states: bool = False,
    output_attentions: bool = False,
    save_to_file: bool = False,
) -> dict | tuple[dict, dict]:
    """Infer for a list of sequences.

    This method provides a convenient interface for performing
    inference on sequences, with optional evaluation and saving
    capabilities.

    Args:
        sequences: Single sequence or list of sequences for inference
        evaluate: Whether to evaluate predictions against true labels
        output_hidden_states: Whether to output hidden states for
            visualization
        output_attentions: Whether to output attention weights for
            visualization
        save_to_file: Whether to save predictions to output directory

    Returns:
        Either:
            - Dict: Dictionary containing predictions
            - Tuple[Dict, Dict]: (predictions, metrics) if evaluate=True

    Note:
        Evaluation requires that labels are available in the dataset
    """
    # Get dataset and dataloader from sequences
    _, dataloader = self.generate_dataset(sequences, batch_size=self.pred_config.batch_size)
    # Do batch inference
    logits, predictions, embeddings = self.batch_infer(
        dataloader,
        output_hidden_states=output_hidden_states,
        output_attentions=output_attentions,
        do_pred=do_pred,
    )
    # Keep hidden states
    if output_hidden_states or output_attentions:
        self.embeddings = embeddings
    # Save predictions
    if save_to_file and self.pred_config.output_dir:
        save_predictions(predictions, Path(self.pred_config.output_dir))  # type: ignore
    # Do evaluation
    if len(self.labels) == len(logits) and evaluate:
        metrics = self.calculate_metrics(logits, self.labels)
        metrics_save = dict(metrics)
        metrics_save.pop("curve", None)
        metrics_save.pop("scatter", None)
        if save_to_file and self.pred_config.output_dir:
            save_metrics(metrics_save, Path(self.pred_config.output_dir))
        return predictions, metrics  # type: ignore

    return predictions  # type: ignore

logits_to_preds

logits_to_preds(logits)

Convert model logits to predictions and human-readable labels.

This method processes raw model outputs based on the task type to generate appropriate predictions and convert them to human-readable labels.

Parameters:

Name	Type	Description	Default
logits	Tensor	Model output logits tensor	required

Returns:

Type	Description
tuple[Tensor, list]	Tuple containing: - torch.Tensor: Model predictions (probabilities or raw values) - List: Human-readable labels corresponding to predictions

Raises:

Type	Description
ValueError	If task type is not supported

Source code in dnallm/inference/inference.py

def logits_to_preds(self, logits: torch.Tensor) -> tuple[torch.Tensor, list]:
    """Convert model logits to predictions and human-readable labels.

    This method processes raw model outputs based on the task type to
    generate appropriate predictions and convert them to human-readable
    labels.

    Args:
        logits: Model output logits tensor

    Returns:
        Tuple containing:
            - torch.Tensor: Model predictions (probabilities or raw values)
            - List: Human-readable labels corresponding to predictions

    Raises:
        ValueError: If task type is not supported
    """
    # Get task type and threshold from config
    task_type = self.task_config.task_type
    threshold = self.task_config.threshold
    label_names = self.task_config.label_names
    # Convert logits to predictions based on task type
    if task_type == "binary":
        probs = torch.softmax(logits, dim=-1)
        preds = (probs[:, 1] > threshold).long()
        labels = [label_names[pred] for pred in preds]
    elif task_type == "multiclass":
        probs = torch.softmax(logits, dim=-1)
        preds = torch.argmax(probs, dim=-1)
        labels = [label_names[pred] for pred in preds]
    elif task_type == "multilabel":
        probs = torch.sigmoid(logits)
        preds = (probs > threshold).long()
        labels = []
        for pred in preds:
            label = [label_names[i] for i in range(len(pred)) if pred[i] == 1]
            labels.append(label)
    elif task_type == "regression":
        preds = logits.squeeze(-1)
        probs = preds
        labels = label_names
    elif task_type == "token":
        probs = torch.softmax(logits, dim=-1)
        preds = torch.argmax(logits, dim=-1)
        labels = []
        for pred in preds:
            label = [label_names[pred[i]] for i in range(len(pred))]
            labels.append(label)
    else:
        raise ValueError(f"Unsupported task type: {task_type}")
    return probs, labels

plot_attentions

plot_attentions(
    seq_idx=0,
    layer=-1,
    head=-1,
    norm_method=None,
    skip_cls=True,
    width=800,
    height=800,
    save_path=None,
)

Plot attention map visualization.

This method creates a heatmap visualization of attention weights between tokens in a sequence, showing how the model attends to different parts of the input.

Parameters:

Name	Type	Description	Default
seq_idx	int	Index of the sequence to plot, default 0	0
layer	int	Layer index to visualize, default -1 (last layer)	-1
head	int	Attention head index to visualize, default -1 (last head)	-1
width	int	Width of the plot	800
height	int	Height of the plot	800
save_path	str | None	Path to save the plot. If None, plot will be shown interactively	None

Returns:

Type	Description
Any | None	Attention map visualization if available, otherwise None

Note

This method requires that attention weights were collected during inference by setting output_attentions=True in prediction methods

Source code in dnallm/inference/inference.py

def plot_attentions(
    self,
    seq_idx: int = 0,
    layer: int = -1,
    head: int = -1,
    norm_method: str | None = None,
    skip_cls=True,
    width: int = 800,
    height: int = 800,
    save_path: str | None = None,
) -> Any | None:
    """Plot attention map visualization.

    This method creates a heatmap visualization of attention weights
    between tokens in a sequence, showing how the model attends to
    different parts of the input.

    Args:
        seq_idx: Index of the sequence to plot, default 0
        layer: Layer index to visualize, default -1 (last layer)
        head: Attention head index to visualize, default -1 (last head)
        width: Width of the plot
        height: Height of the plot
        save_path: Path to save the plot. If None, plot will be shown
            interactively

    Returns:
        Attention map visualization if available, otherwise None

    Note:
        This method requires that attention weights were collected
        during inference by setting output_attentions=True in prediction
        methods
    """
    if hasattr(self, "embeddings"):
        attentions = self.embeddings["attentions"]
        if save_path:
            suffix = os.path.splitext(save_path)[-1]
            if suffix:
                heatmap = save_path.replace(suffix, "_heatmap" + suffix)
            else:
                heatmap = os.path.join(save_path, "heatmap.pdf")
        else:
            heatmap = None
        # Plot attention map
        attn_map = plot_attention_map(
            attentions,
            self.sequences,
            self.tokenizer,
            seq_idx=seq_idx,
            layer=layer,
            norm_method=norm_method,
            skip_cls=skip_cls,
            head=head,
            width=width,
            height=height,
            save_path=heatmap,
        )
        return attn_map
    else:
        logger.warning("No attention weights available to plot.")
        return None

plot_hidden_states

plot_hidden_states(
    reducer="t-SNE",
    reduced=False,
    quality="fast",
    ncols=4,
    width=300,
    height=300,
    point_size=10,
    save_path=None,
)

Visualize embeddings using dimensionality reduction.

This method creates 2D visualizations of high-dimensional embeddings from different model layers using PCA, t-SNE, or UMAP dimensionality reduction.

Parameters:

Name	Type	Description	Default
reducer	str	Dimensionality reduction method to use ('PCA', 't-SNE', 'UMAP')	't-SNE'
reduced	bool	Whether to use already reduced embeddings if available	False
quality	str	Quality/speed trade-off for reduction	'fast'
ncols	int	Number of columns in the plot grid	4
width	int	Width of each plot	300
height	int	Height of each plot	300
point_size	int	Size of points in the plot	10
save_path	str | None	Path to save the plot. If None, plot will be shown interactively	None

Returns:

Type	Description
Any | None	Embedding visualization if available, otherwise None

Note

This method requires that hidden states were collected during inference by setting output_hidden_states=True in prediction methods

Source code in dnallm/inference/inference.py

def plot_hidden_states(  # type: ignore
    self,
    reducer: str = "t-SNE",
    reduced: bool = False,
    quality: str = "fast",
    ncols: int = 4,
    width: int = 300,
    height: int = 300,
    point_size: int = 10,
    save_path: str | None = None,
) -> Any | None:
    """Visualize embeddings using dimensionality reduction.

    This method creates 2D visualizations of high-dimensional
    embeddings from different model layers using PCA, t-SNE, or UMAP
    dimensionality reduction.

    Args:
        reducer: Dimensionality reduction method to use
            ('PCA', 't-SNE', 'UMAP')
        reduced: Whether to use already reduced embeddings if available
        quality: Quality/speed trade-off for reduction
        ncols: Number of columns in the plot grid
        width: Width of each plot
        height: Height of each plot
        point_size: Size of points in the plot
        save_path: Path to save the plot. If None, plot will be shown
            interactively

    Returns:
        Embedding visualization if available, otherwise None

    Note:
        This method requires that hidden states were collected during
        inference by setting output_hidden_states=True in prediction
        methods
    """
    if hasattr(self, "embeddings"):
        hidden_states = self.embeddings["hidden_states"]
        # attention_mask = torch.unsqueeze(
        #     self.embeddings["attention_mask"], dim=-1
        # )
        attention_mask = self.embeddings["attention_mask"]
        labels = self.embeddings["labels"]
        if save_path:
            suffix = os.path.splitext(save_path)[-1]
            if suffix:
                embedding = save_path.replace(suffix, "_embedding" + suffix)
            else:
                embedding = os.path.join(save_path, "embedding.pdf")
        else:
            embedding = None
        # Plot hidden states
        label_names = self.task_config.label_names
        embeddings_vis = plot_embeddings(
            hidden_states,
            attention_mask,
            reducer=reducer,
            quality=quality,
            labels=labels,
            label_names=label_names,
            ncols=ncols,
            width=width,
            height=height,
            point_size=point_size,
            save_path=embedding,
            reduced=reduced,
        )
        return embeddings_vis
    else:
        logger.warning("No hidden states available to plot.")

scoring

scoring(
    inputs,
    reduce_method="mean",
    score_type="embedding",
    reduce_hidden_states=False,
)

Score sequences using the model.

This function computes scores for input sequences using the loaded model. It supports specific scoring methods for EVO2, EVO1, and MegaDNA models, as well as a general scoring approach for other base models.

Parameters:

Name	Type	Description	Default
inputs	DataLoader | list[str]	DataLoader or List containing sequences to score	required
reduce_method	str	Method to reduce scores ('mean', 'max', 'min', 'last'), default 'mean'	'mean'
score_type	str	Type of score to compute ('embedding', 'logits', 'probability', 'loss'), default 'embedding'	'embedding'
reduce_hidden_states	bool	Whether to reduce hidden states across layers, default False	False

Returns: Dictionary containing sequences and their corresponding scores

Source code in dnallm/inference/inference.py

def scoring(
    self,
    inputs: DataLoader | list[str],
    reduce_method: str = "mean",
    score_type: str = "embedding",
    reduce_hidden_states: bool = False,
) -> dict[Any, Any]:
    """Score sequences using the model.

    This function computes scores for input sequences using the loaded
    model. It supports specific scoring methods for EVO2, EVO1, and
    MegaDNA models, as well as a general scoring approach for other
    base models.

    Args:
        inputs: DataLoader or List containing sequences to score
        reduce_method: Method to reduce scores ('mean', 'max', 'min',
            'last'), default 'mean'
        score_type: Type of score to compute ('embedding', 'logits',
            'probability', 'loss'), default 'embedding'
        reduce_hidden_states: Whether to reduce hidden states across
            layers, default False
    Returns:
        Dictionary containing sequences and their corresponding scores
    """
    # Prepare score sequences
    score_seqs = []
    if isinstance(inputs, DataLoader):
        for data in tqdm(inputs, desc="Scoring"):
            seqs = (
                data.get("sequence", None)
                if isinstance(data, dict)
                else getattr(data, "sequence", None)
            )
            if not seqs:
                continue
            score_seqs.extend([s for s in seqs if s])
    else:
        score_seqs = inputs
    # Check if model supports scoring
    model_name = str(self.model).lower()
    if "evo2" in model_name:
        outputs = self.model.score_sequences(score_seqs, reduce_method=reduce_method)
        outputs = [{"Input": score_seqs[i], "Score": s} for i, s in enumerate(outputs)]
        return outputs  # type: ignore
    elif "evo1" in model_name:
        from evo import score_sequences

        model = self.model.model
        tokenizer = self.tokenizer
        outputs = score_sequences(
            score_seqs,
            model=model,
            tokenizer=tokenizer.raw_tokenizer,
            reduce_method=reduce_method,
            device=self.device,
        )
        outputs = [{"Input": score_seqs[i], "Score": s} for i, s in enumerate(outputs)]
        return outputs  # type: ignore
    elif "megadna" in model_name:
        model = self.model
        tokenizer = self.tokenizer
        outputs = []
        for seq in score_seqs:
            input_ids = tokenizer(seq, return_tensors="pt").to(self.device)["input_ids"]
            with torch.no_grad():
                loss = model(input_ids, return_value="loss")
            outputs.append({"Input": seq, "Score": loss})
        return outputs  # type: ignore

    # General scoring for other base models (No classification head)
    # Use batch_infer to get embeddings and compute scores
    if isinstance(inputs, list):
        _, dataloader = self.generate_dataset(
            inputs,
            batch_size=self.pred_config.batch_size,
            padding="longest",
        )
    elif isinstance(inputs, dict):
        _, dataloader = self.generate_dataset(
            inputs["sequence"],
            batch_size=self.pred_config.batch_size,
            padding="longest",
        )
    elif isinstance(inputs, DataLoader):
        if "sequence" in inputs.dataset.dataset.column_names:  # type: ignore
            seqs = inputs.dataset["sequence"]
            if len(seqs) != len(inputs.dataset):  # type: ignore
                raise ValueError("Some sequences are missing in the dataset.")
            score_seqs = [s for s in seqs if s]
        if score_seqs:
            _, dataloader = self.generate_dataset(
                score_seqs,
                batch_size=self.pred_config.batch_size,
                padding="longest",
            )
        else:
            dataloader = inputs
    else:
        dataloader = inputs  # type: ignore[unreachable]
    all_logits, _, embeddings = self.batch_infer(
        dataloader,
        output_hidden_states=True if score_type == "embedding" else False,
        reduce_hidden_states=reduce_hidden_states,
        reduce_strategy=reduce_method,
        do_pred=False,
    )
    # Prepare logits list for scoring
    logits_list = []
    if isinstance(all_logits, torch.Tensor):
        # assume shape (N, L, V)
        for i in range(all_logits.size(0)):
            logits_list.append(all_logits[i].detach().cpu())
    elif isinstance(all_logits, (list, tuple)):  # type: ignore[unreachable]
        for item in all_logits:
            if item is not None:
                logits_list.append(
                    item.detach().cpu()
                    if isinstance(item, torch.Tensor)
                    else torch.tensor(item)
                )
    else:
        logits_list = []
    # Compute scores
    scores = []
    if "hidden_states" in embeddings and score_type == "embedding":
        hidden_states = embeddings["hidden_states"]
        for i in range(len(score_seqs)):
            # (layers, seq_len, dim)
            if reduce_method == "last":
                seq_hidden = hidden_states[-1][i]
            elif reduce_method == "first":
                seq_hidden = hidden_states[0][i]
            else:
                seq_hidden = torch.stack([h[i] for h in hidden_states], dim=0)
            if reduce_method == "mean":
                score = seq_hidden.mean().item()
            elif reduce_method == "max":
                score = seq_hidden.max().item()
            elif reduce_method == "min":
                score = seq_hidden.min().item()
            else:
                score = seq_hidden.mean().item()
            scores.append({"Input": score_seqs[i], "Score": score})
        return scores  # type: ignore
    elif logits_list and score_type == "logits":
        # use logits as scores
        for i in range(len(score_seqs)):
            logits = logits_list[i]
            if logits.dim() == 3:
                logits = logits.squeeze(0)
            if reduce_method == "mean":
                score = logits.mean().item()
            elif reduce_method == "max":
                score = logits.max().item()
            elif reduce_method == "min":
                score = logits.min().item()
            else:
                score = logits.mean().item()
            scores.append({"Input": score_seqs[i], "Score": score})
        return scores  # type: ignore
    elif logits_list and score_type == "probability":
        tokenizer = self.tokenizer
        for i, seq in enumerate(score_seqs):
            logits = logits_list[i].to(self.device)  # (L, V) or (1,L,V)
            if logits.dim() == 3 and logits.size(0) == 1:
                logits = logits.squeeze(0)
            # Re-tokenize the sequence to obtain input_ids & attention_mask
            enc = tokenizer(seq, return_tensors="pt", padding=False)
            input_ids = enc["input_ids"].to(self.device)  # (1, L)
            attention_mask = enc.get("attention_mask", None)
            if attention_mask is not None:
                attention_mask = attention_mask.to(self.device)
            # Decide model type by presence of mask token
            mask_token_id = getattr(tokenizer, "mask_token_id", None)
            is_mlm = False
            if mask_token_id is not None:
                if (input_ids == mask_token_id).any():
                    is_mlm = True
            with torch.no_grad():
                # (L, V)
                logprobs = torch.log_softmax(logits.to(self.device), dim=-1)
            tgt_ids = input_ids.squeeze(0)
            if is_mlm:
                # MLM: logits[t] predicts token at t
                # Gather token logprobs at each position
                gathered = logprobs.gather(1, tgt_ids.unsqueeze(-1)).squeeze(-1)  # (L,)
                # only keep positions where mask appears
                mask_positions = tgt_ids == mask_token_id
                if mask_positions.any():
                    token_logprobs = gathered[mask_positions]
                else:
                    # No explicit mask found; fall back to using all tokens
                    token_logprobs = gathered
            else:
                # Causal LM: logits[t] predicts token at t+1
                # Align: drop last logit, drop first input id
                if logprobs.size(0) >= 2 and tgt_ids.size(0) >= 2:
                    lp = logprobs[:-1, :]  # (L-1, V)
                    tgt = tgt_ids[1:]  # (L-1,)
                    gathered = lp.gather(1, tgt.unsqueeze(-1)).squeeze(-1)
                    # apply attention_mask if available (exclude padding)
                    if attention_mask is not None:
                        attn = attention_mask.squeeze(0)[1:].to(torch.bool)
                        if attn.any():
                            token_logprobs = gathered[attn]
                        else:
                            token_logprobs = gathered
                    else:
                        token_logprobs = gathered
                else:
                    # fallback: gather directly (if short)
                    gathered = logprobs.gather(1, tgt_ids.unsqueeze(-1)).squeeze(-1)
                    token_logprobs = gathered
            if token_logprobs.numel() == 0:
                score = float("nan")  # no tokens to score
            else:
                if reduce_method == "mean":
                    score = float(token_logprobs.mean().item())
                elif reduce_method == "sum":
                    score = float(token_logprobs.sum().item())
                elif reduce_method == "max":
                    score = float(token_logprobs.max().item())
                elif reduce_method == "min":
                    score = float(token_logprobs.min().item())
                else:
                    score = float(token_logprobs.mean().item())
            scores.append({"Input": seq, "Score": score})
        return scores  # type: ignore

    return {}

Functions:

save_metrics

save_metrics(metrics, output_dir)

Save evaluation metrics to JSON file.

This function saves computed evaluation metrics in JSON format to the specified output directory.

Parameters:

Name	Type	Description	Default
metrics	dict	Dictionary containing metrics to save	required
output_dir	Path	Directory path where metrics will be saved	required

Source code in dnallm/inference/inference.py

def save_metrics(metrics: dict, output_dir: Path) -> None:
    """Save evaluation metrics to JSON file.

    This function saves computed evaluation metrics in JSON format to the
    specified output directory.

    Args:
        metrics: Dictionary containing metrics to save
        output_dir: Directory path where metrics will be saved
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save metrics
    with open(output_dir / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=4)

save_predictions

save_predictions(predictions, output_dir)

Save predictions to JSON file.

This function saves model predictions in JSON format to the specified output directory.

Parameters:

Name	Type	Description	Default
predictions	dict	Dictionary containing predictions to save	required
output_dir	Path	Directory path where predictions will be saved	required

Source code in dnallm/inference/inference.py

def save_predictions(predictions: dict, output_dir: Path) -> None:
    """Save predictions to JSON file.

    This function saves model predictions in JSON format to the specified
    output directory.

    Args:
        predictions: Dictionary containing predictions to save
        output_dir: Directory path where predictions will be saved
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save predictions
    with open(output_dir / "predictions.json", "w") as f:
        json.dump(predictions, f, indent=4)