[McClain-Thiel]

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[Rocketknight1]

Hi @McClain-Thiel, I'm the ex-postdoc who handles a lot of bio models at HF so feel free to ping me if you have any questions, or whenever this is ready for review!

[McClain-Thiel]

Awesome thanks @Rocketknight1. One quick question: not sure if youve read through the evo repo but they use a custom inference engine + library called vortex. What is the huggingface philosophy on including this as a dependency? Am i going to have to reimplement all the forward stuff in that in base torch? Or can I add it as a dep?

[Rocketknight1]

Hi @McClain-Thiel, good question! We generally prefer not to add whole external inference engines. The right approach is probably:

1. Try to port the core model architecture to a modular file
2. If the model has weird kernels in vortex, then firstly see if we have an existing similar model you can copy from for the modular (we definitely have other SSM/Hyena-like models - there might be one with equivalent architecture!)
3. As a last resort, we can try to implement missing kernels via something like Kernel Hub

I haven't investigated too deeply, but if the inference lib also does other stuff like searching for multiple sequence alignments, we probably don't want to put all that directly in Transformers, but we could just add info to the model card. It's fine to show examples in the model card using vortex as long as we can run the core model without it.

[McClain-Thiel]

@Rocketknight1 So I think I have it mostly working. Gemini 3 did a lot of the heavy lifting but if you look at the tests I actually got the logits from the vortex implementation then use those to verify that the implementation here is close. I can remove that file ofc but thought you might want to take a look.

It also passes a simple smoke test (tokenizer is single letter (not constrained to ATGC) but when sampled it seems to give reasonable results (ATCG). Open to more rigorous testing.

Any other suggestions?

[Rocketknight1]

The main thing I'd suggest now before I do a full review is to convert to modular. Basically, we want each modeling file in Transformers to be standalone, to make it easy for users to hack on them, and so that users don't need to understand ten layers of inheritance and abstraction to understand the code.

However, there's a lot of duplicated code this way, and the way we handle that is "modular" files. They're like scaffolds for the actual modeling files - unlike modeling files, imports and inheritance from other models is allowed there. Ideally, when a model is similar to existing models in the codebase, the modular file is very short and mostly just imports/inheritance, but it still can be used to auto-generate a full modeling file (which is the code that actually gets called)

Take a look at other modular files in the codebase to get an idea of how it works. In general, "# Copied from" syntax is deprecated now, and we prefer to import in the modular file instead. Once you've made the modular file, running make fix-copies will autogenerate a modeling file for you. You shouldn't edit that file directly, just tweak the modular file and rerun make fix-copies.

[ArthurZucker]

let's make this as simple as possible

[ArthurZucker]

Dead against the 7 released vortex configs: num_key_value_heads (proj_groups: 1 everywhere), the three dropouts, hyena_order (always 3), hyena_flip_x1x2 (false everywhere) and hyena_filter_configurations (shapes follow from the layer kind). Released configs have short_filter_length: 3 (default here is 8), and the 262k / 1M checkpoints need rope_parameters={"rope_type": "linear", "factor": 32 / 128}. layer_types is derived from mixer_types so DynamicCache(config) builds the hybrid cache.

Suggested change

	"""Evo2 model configuration."""
	from __future__ import annotations
	from typing import Optional, Sequence, List, Dict, Any
	from ...configuration_utils import PretrainedConfig
	from ...modeling_rope_utils import standardize_rope_params
	from ...utils import logging
	logger = logging.get_logger(__name__)
	__all__ = ["Evo2Config"]
	class Evo2Config(PretrainedConfig):
	r"""
	This is the configuration class to store the configuration of a [`Evo2Model`]. It is used to instantiate an Evo2
	model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
	defaults will yield a similar configuration to that of the Evo2-1b-base model.
	Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
	documentation from [`PretrainedConfig`] for more information.
	Args:
	vocab_size (`int`, *optional*, defaults to 512):
	Vocabulary size of the Evo2 model. Defines the number of different tokens that can be represented by the
	`inputs_ids` passed when calling [`Evo2Model`]
	hidden_size (`int`, *optional*, defaults to 2048):
	Dimension of the hidden representations.
	intermediate_size (`int`, *optional*):
	Dimension of the MLP representations.
	num_hidden_layers (`int`, *optional*, defaults to 24):
	Number of hidden layers in the Transformer encoder.
	num_attention_heads (`int`, *optional*, defaults to 16):
	Number of attention heads for each attention layer in the Transformer encoder.
	num_key_value_heads (`int`, *optional*):
	This is the number of key_value heads that should be used to implement Grouped Query Attention. If
	`num_key_value_heads=None`, the model will use the same number of key/value heads as the number of
	query heads.
	max_position_embeddings (`int`, *optional*, defaults to 2048):
	The maximum sequence length that this model might ever be used with.
	rope_theta (`float`, *optional*, defaults to 1000000.0):
	The base period of the RoPE embeddings.
	rms_norm_eps (`float`, *optional*, defaults to 1e-6):
	The epsilon used by the rms normalization layers.
	attn_dropout (`float`, *optional*, defaults to 0.0):
	The dropout ratio for the attention probabilities.
	hidden_dropout (`float`, *optional*, defaults to 0.0):
	The dropout ratio for the hidden units.
	mlp_dropout (`float`, *optional*, defaults to 0.0):
	The dropout ratio for the MLP layers.
	layer_types (`Sequence[str]`, *optional*):
	List of layer types ("attention" or "hyena") for each layer. If None, defaults to all "attention".
	hyena_filters (`int`, *optional*, defaults to 256):
	Number of Hyena filter groups.
	hyena_kernel_size (`int`, *optional*, defaults to 8):
	Kernel size for the short convolution in Hyena.
	hyena_hidden_size (`int`, *optional*):
	Hidden size for Hyena layers.
	hyena_order (`int`, *optional*, defaults to 4):
	Order of the Hyena recurrence.
	hyena_flip_x1x2 (`bool`, *optional*, defaults to False):
	Whether to flip x1 and x2 in the Hyena gating mechanism.
	initializer_range (`float`, *optional*, defaults to 0.02):
	The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
	use_cache (`bool`, *optional*, defaults to True):
	Whether or not the model should return the last key/values attentions (not used by all models).
	pad_token_id (`int`, *optional*, defaults to 1):
	Padding token id.
	bos_token_id (`int`, *optional*):
	Beginning of stream token id.
	eos_token_id (`int`, *optional*, defaults to 0):
	End of stream token id.
	tie_word_embeddings (`bool`, *optional*, defaults to True):
	Whether to tie weight embeddings
	"""
	model_type = "evo2"
	keys_to_ignore_at_inference = ["past_key_values"]
	def __init__(
	self,
	vocab_size: int = 512,
	hidden_size: int = 2048,
	intermediate_size: Optional[int] = None,
	num_hidden_layers: int = 24,
	num_attention_heads: int = 16,
	num_key_value_heads: Optional[int] = None,
	max_position_embeddings: int = 2048,
	rope_theta: float = 1_000_000.0,
	rms_norm_eps: float = 1e-6,
	attn_dropout: float = 0.0,
	hidden_dropout: float = 0.0,
	mlp_dropout: float = 0.0,
	layer_types: Optional[Sequence[str]] = None,
	hyena_filters: int = 256,
	hyena_kernel_size: int = 8,
	hyena_hidden_size: Optional[int] = None,
	hyena_order: int = 4,
	hyena_flip_x1x2: bool = False,
	hyena_filter_configurations: Optional[List[Dict[str, Any]]] = None,
	initializer_range: float = 0.02,
	use_cache: bool = True,
	pad_token_id: int = 1,
	bos_token_id: Optional[int] = None,
	eos_token_id: int = 0,
	tie_word_embeddings: bool = True,
	**kwargs,
	) -> None:
	super().__init__(
	pad_token_id=pad_token_id,
	bos_token_id=bos_token_id,
	eos_token_id=eos_token_id,
	tie_word_embeddings=tie_word_embeddings,
	**kwargs,
	)
	self.vocab_size = vocab_size
	self.hidden_size = hidden_size
	self.intermediate_size = intermediate_size if intermediate_size is not None else hidden_size * 4
	self.num_hidden_layers = num_hidden_layers
	self.num_attention_heads = num_attention_heads
	self.num_key_value_heads = num_key_value_heads or num_attention_heads
	self.max_position_embeddings = max_position_embeddings
	self.rope_theta = rope_theta
	self.rms_norm_eps = rms_norm_eps
	self.attn_dropout = attn_dropout
	self.hidden_dropout = hidden_dropout
	self.mlp_dropout = mlp_dropout
	self.hyena_filters = hyena_filters
	self.hyena_kernel_size = hyena_kernel_size
	self.hyena_hidden_size = hyena_hidden_size if hyena_hidden_size is not None else hidden_size
	self.hyena_order = hyena_order
	self.hyena_flip_x1x2 = hyena_flip_x1x2
	self.hyena_filter_configurations = hyena_filter_configurations
	self.initializer_range = initializer_range
	self.use_cache = use_cache
	if layer_types is None:
	self.layer_types = ["attention"] * num_hidden_layers
	else:
	self.layer_types = list(layer_types)
	standardize_rope_params(self, rope_theta=self.rope_theta)
	if len(self.layer_types) != self.num_hidden_layers:
	raise ValueError(
	"The length of `layer_types` must match `num_hidden_layers` (received"
	f" {len(self.layer_types)} and {self.num_hidden_layers})."
	)
	for layer_type in self.layer_types:
	if layer_type not in {"attention", "hyena"}:
	raise ValueError(f"Unsupported layer type: {layer_type}. Expected 'attention' or 'hyena'.")
	if self.num_attention_heads <= 0 or self.hidden_size % self.num_attention_heads != 0:
	raise ValueError("`hidden_size` must be divisible by `num_attention_heads`.")
	if self.num_key_value_heads <= 0 or self.hidden_size % self.num_key_value_heads != 0:
	raise ValueError("`hidden_size` must be divisible by `num_key_value_heads`.")
	logger.info("Initialized Evo2Config with %s layers (%s).", self.num_hidden_layers, ", ".join(self.layer_types))
	@property
	def head_dim(self) -> int:
	return self.hidden_size // self.num_attention_heads
	@property
	def kv_head_dim(self) -> int:
	return self.hidden_size // self.num_key_value_heads
	@property
	def num_attention_layers(self) -> int:
	return sum(layer_type == "attention" for layer_type in self.layer_types)
	@property
	def num_hyena_layers(self) -> int:
	return sum(layer_type == "hyena" for layer_type in self.layer_types)
	def to_dict(self) -> dict:
	output = super().to_dict()
	output["layer_types"] = list(self.layer_types)
	return output
	# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
	# This file was automatically generated from src/transformers/models/evo2/modular_evo2.py.
	# Do NOT edit this file manually as any edits will be overwritten by the generation of
	# the file from the modular. If any change should be done, please apply the change to the
	# modular_evo2.py file directly. One of our CI enforces this.
	# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
	# Copyright 2025 Arc Institute and The HuggingFace Inc. team. All rights reserved.
	#
	# 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.
	from huggingface_hub.dataclasses import strict
	from ...configuration_utils import PreTrainedConfig
	from ...modeling_rope_utils import RopeParameters
	from ...utils import auto_docstring
	@auto_docstring(checkpoint="arcinstitute/evo2_1b_base")
	@strict
	class Evo2Config(PreTrainedConfig):
	r"""
	mixer_types (`list[str]`, *optional*):
	Mixer of each layer: `"attention"`, `"hcs"` (short explicit filter), `"hcm"` (medium explicit filter) or
	`"hcl"` (long implicit filter). Defaults to the repeating `hcs, hcm, hcl, attention, hcs, hcm, hcl` stripe.
	short_conv_kernel_size (`int`, *optional*, defaults to 3):
	Kernel size of the depthwise convolution applied to the hyena input projections.
	hcs_filter_length (`int`, *optional*, defaults to 7):
	Length of the explicit filter of `"hcs"` layers.
	hcm_filter_length (`int`, *optional*, defaults to 128):
	Length of the explicit filter of `"hcm"` layers.
	filter_groups (`int`, *optional*, defaults to 128):
	Number of channel groups sharing one explicit filter in `"hcs"` / `"hcm"` layers.
	state_size (`int`, *optional*, defaults to 16):
	Number of modes (poles/residues) per channel of the `"hcl"` long filter.
	"""
	model_type = "evo2"
	keys_to_ignore_at_inference = ["past_key_values"]
	attribute_map = {"num_key_value_heads": "num_attention_heads"}
	# q/k/v have no bias in every released checkpoint, o_proj always has one (set in `Evo2Attention`)
	attention_bias = False
	attention_dropout = 0.0
	mlp_bias = False
	# hyena layers cache two states: the short conv window, and the explicit filter window or the modal state
	number_of_conv_states = 2
	vocab_size: int = 512
	hidden_size: int = 1920
	intermediate_size: int = 5120
	num_hidden_layers: int = 25
	num_attention_heads: int = 15
	hidden_act: str = "gelu"
	max_position_embeddings: int = 8192
	initializer_range: float = 0.02
	norm_eps: float = 1e-6
	use_cache: bool = True
	pad_token_id: int | None = 1
	bos_token_id: int | None = None
	eos_token_id: int | list[int] | None = 0
	tie_word_embeddings: bool = True
	rope_parameters: RopeParameters | dict | None = None
	mixer_types: list[str] | None = None
	short_conv_kernel_size: int = 3
	hcs_filter_length: int = 7
	hcm_filter_length: int = 128
	filter_groups: int = 128
	state_size: int = 16
	def __post_init__(self, **kwargs):
	if self.mixer_types is None:
	stripe = ["hcs", "hcm", "hcl", "attention", "hcs", "hcm", "hcl"]
	self.mixer_types = [stripe[i % len(stripe)] for i in range(self.num_hidden_layers)]
	super().__post_init__(**kwargs)
	@property
	def layer_types(self):
	return ["full_attention" if mixer == "attention" else "conv" for mixer in self.mixer_types]
	__all__ = ["Evo2Config"]

[ArthurZucker]

Generated imports for the modular in the review body.

Suggested change

	"""PyTorch Evo2 model."""
	from __future__ import annotations
	import math
	from collections.abc import Callable
	from typing import Optional, Tuple, Union
	import torch
	import torch.nn.functional as F
	from torch import nn
	from torch.utils.checkpoint import checkpoint
	from ...cache_utils import Cache, DynamicCache
	from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
	from ...generation import GenerationMixin
	from ...modeling_layers import GradientCheckpointingLayer
	from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
	from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
	from ...modeling_utils import PreTrainedModel
	from ...utils import logging
	from .configuration_evo2 import Evo2Config
	logger = logging.get_logger(__name__)
	__all__ = ["Evo2Model", "Evo2ForCausalLM", "Evo2PreTrainedModel"]
	# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
	# This file was automatically generated from src/transformers/models/evo2/modular_evo2.py.
	# Do NOT edit this file manually as any edits will be overwritten by the generation of
	# the file from the modular. If any change should be done, please apply the change to the
	# modular_evo2.py file directly. One of our CI enforces this.
	# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
	# Copyright 2025 Arc Institute and The HuggingFace Inc. team. All rights reserved.
	#
	# 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.
	from collections.abc import Callable
	import torch
	import torch.nn.functional as F
	from torch import nn
	from ... import initialization as init
	from ...activations import ACT2FN
	from ...cache_utils import Cache, DynamicCache
	from ...generation import GenerationMixin
	from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
	from ...masking_utils import create_causal_mask, create_recurrent_attention_mask
	from ...modeling_layers import GradientCheckpointingLayer
	from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
	from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
	from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
	from ...processing_utils import Unpack
	from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
	from ...utils.generic import merge_with_config_defaults
	from ...utils.output_capturing import capture_outputs
	from .configuration_evo2 import Evo2Config

[ArthurZucker]

This is LlamaRMSNorm (class Evo2RMSNorm(LlamaRMSNorm): pass in modular). Note vortex divides by rms + eps instead of sqrt(ms + eps); negligible at 1e-6.

Suggested change

	# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Evo2
	class Evo2RMSNorm(nn.Module):
	def __init__(self, hidden_size, eps=1e-6):
	"""Evo2RMSNorm is equivalent to T5LayerNorm."""
	super().__init__()
	self.weight = nn.Parameter(torch.ones(hidden_size))
	self.variance_epsilon = eps
	def forward(self, hidden_states):
	input_dtype = hidden_states.dtype
	hidden_states = hidden_states.to(torch.float32)
	variance = hidden_states.pow(2).mean(-1, keepdim=True)
	hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
	return self.weight * hidden_states.to(input_dtype)
	def extra_repr(self):
	return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
	@use_kernel_forward_from_hub("RMSNorm")
	class Evo2RMSNorm(nn.Module):
	def __init__(self, hidden_size, eps: float = 1e-6) -> None:
	"""
	Evo2RMSNorm is equivalent to T5LayerNorm
	"""
	super().__init__()
	self.weight = nn.Parameter(torch.ones(hidden_size))
	self.variance_epsilon = eps
	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	input_dtype = hidden_states.dtype
	hidden_states = hidden_states.to(torch.float32)
	variance = hidden_states.pow(2).mean(-1, keepdim=True)
	hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
	return self.weight * hidden_states.to(input_dtype)
	def extra_repr(self):
	return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"

[ArthurZucker]

Llama's helpers, they come with Evo2Attention(LlamaAttention) in modular.

Suggested change

	# Copied from transformers.models.llama.modeling_llama.rotate_half
	def rotate_half(x):
	"""Rotates half the hidden dims of the input."""
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2 :]
	return torch.cat((-x2, x1), dim=-1)
	# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
	def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
	"""Applies Rotary Position Embedding to the query and key tensors."""
	del position_ids
	cos = cos.unsqueeze(unsqueeze_dim)
	sin = sin.unsqueeze(unsqueeze_dim)
	q_embed = (q * cos) + (rotate_half(q) * sin)
	k_embed = (k * cos) + (rotate_half(k) * sin)
	return q_embed, k_embed
	def rotate_half(x):
	"""Rotates half the hidden dims of the input."""
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2 :]
	return torch.cat((-x2, x1), dim=-1)
	@use_kernel_forward_from_hub("rotary_pos_emb")
	def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
	"""Applies Rotary Position Embedding to the query and key tensors.
	Args:
	q (`torch.Tensor`): The query tensor.
	k (`torch.Tensor`): The key tensor.
	cos (`torch.Tensor`): The cosine part of the rotary embedding.
	sin (`torch.Tensor`): The sine part of the rotary embedding.
	unsqueeze_dim (`int`, *optional*, defaults to 1):
	The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
	sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
	that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
	k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
	cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
	the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
	Returns:
	`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
	"""
	cos = cos.unsqueeze(unsqueeze_dim)
	sin = sin.unsqueeze(unsqueeze_dim)
	q_embed = (q * cos) + (rotate_half(q) * sin)
	k_embed = (k * cos) + (rotate_half(k) * sin)
	return q_embed, k_embed

[ArthurZucker]

LlamaRotaryEmbedding; inv_freq does not need to be a persistent buffer. The 262k / 1M checkpoints use linear position interpolation (rotary_emb_scaling_factor 32 / 128), which is the existing rope_type: "linear" and is currently ignored here.

Suggested change

	class Evo2RotaryEmbedding(nn.Module):
	inv_freq: torch.Tensor
	def __init__(self, config: Evo2Config, device=None):
	super().__init__()
	self.max_seq_len_cached = getattr(config, "max_position_embeddings", 2048)
	self.original_max_seq_len = self.max_seq_len_cached
	self.config = config
	self.rope_type = self.config.rope_parameters["rope_type"]
	rope_init_fn: Callable = self.compute_default_rope_parameters
	if self.rope_type != "default":
	rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
	inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
	# Register inv_freq as persistent so it can be loaded from checkpoints
	self.register_buffer("inv_freq", inv_freq, persistent=True)
	self.original_inv_freq = inv_freq
	@staticmethod
	def compute_default_rope_parameters(
	config: Optional[Evo2Config] = None,
	device: Optional[torch.device] = None,
	seq_len: Optional[int] = None,
	) -> Tuple[torch.Tensor, float]:
	del seq_len
	rope_params = getattr(config, "rope_parameters", None)
	base = rope_params.get("rope_theta") if rope_params is not None else config.rope_theta
	dim = config.head_dim
	inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float, device=device) / dim))
	return inv_freq, 1.0
	@torch.no_grad()
	@dynamic_rope_update
	def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
	inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
	position_ids_expanded = position_ids[:, None, :].float()
	device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
	with torch.autocast(device_type=device_type, enabled=False):
	freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
	emb = torch.cat((freqs, freqs), dim=-1)
	cos = emb.cos() * self.attention_scaling
	sin = emb.sin() * self.attention_scaling
	return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
	class Evo2RotaryEmbedding(nn.Module):
	def __init__(self, config: Evo2Config):
	super().__init__()
	self.max_seq_len_cached = config.max_position_embeddings
	self.original_max_seq_len = config.max_position_embeddings
	self.config = config
	self.rope_type = self.config.rope_parameters["rope_type"]
	rope_init_fn: Callable = self.compute_default_rope_parameters
	if self.rope_type != "default":
	rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
	inv_freq, self.attention_scaling = rope_init_fn(self.config)
	self.inv_freq = nn.Buffer(inv_freq, persistent=False)
	self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)
	@staticmethod
	def compute_default_rope_parameters(config: Evo2Config, **kwargs) -> tuple[torch.Tensor, float]:
	"""
	Computes the inverse frequencies according to the original RoPE implementation
	Args:
	config ([`~transformers.PreTrainedConfig`]):
	The model configuration.
	Returns:
	Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
	post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
	"""
	base = config.rope_parameters["rope_theta"]
	dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
	attention_factor = 1.0 # Unused in this type of RoPE
	# Compute the inverse frequencies
	inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
	return inv_freq, attention_factor
	@torch.no_grad()
	@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
	def forward(self, x, position_ids):
	freqs = position_ids[..., None].float() * self.inv_freq.to(device=x.device, dtype=torch.float)
	emb = torch.cat((freqs, freqs), dim=-1)
	cos = emb.cos() * self.attention_scaling
	sin = emb.sin() * self.attention_scaling
	return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

[ArthurZucker]

Three classes for one pre-norm block: this is Lfm2DecoderLayer (attention or conv mixer + MLP). The fp32 residual casts and dropouts are not part of the reference inference path.

Suggested change

	class Evo2AttentionBlock(nn.Module):
	def __init__(self, config: Evo2Config, layer_idx: int):
	super().__init__()
	self.attention = Evo2Attention(config, layer_idx)
	self.input_layernorm = Evo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.post_attention_layernorm = Evo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.mlp = Evo2ParallelGatedMLP(config, layer_idx)
	self.hidden_dropout = nn.Dropout(config.hidden_dropout)
	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor],
	position_ids: torch.Tensor,
	past_key_value: Optional[Cache],
	output_attentions: bool,
	use_cache: bool,
	cache_position: Optional[torch.LongTensor],
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
	# Keep residual in float32 for numerical stability (Mamba2 technique)
	residual = hidden_states.to(torch.float32) if hidden_states.dtype == torch.bfloat16 else hidden_states
	hidden_states = self.input_layernorm(hidden_states)
	attn_output, attn_weights, present_kv = self.attention(
	hidden_states,
	attention_mask,
	position_ids,
	past_key_value=past_key_value,
	output_attentions=output_attentions,
	use_cache=use_cache,
	cache_position=cache_position,
	)
	hidden_states = (residual + self.hidden_dropout(attn_output).to(residual.dtype)).to(hidden_states.dtype)
	# Keep residual in float32 for numerical stability
	residual = hidden_states.to(torch.float32) if hidden_states.dtype == torch.bfloat16 else hidden_states
	hidden_states = self.post_attention_layernorm(hidden_states)
	hidden_states = self.mlp(hidden_states)
	hidden_states = (residual + self.hidden_dropout(hidden_states).to(residual.dtype)).to(hidden_states.dtype)
	return hidden_states, attn_weights, present_kv
	class Evo2HyenaBlock(nn.Module):
	def __init__(self, config: Evo2Config, layer_idx: int):
	super().__init__()
	self.input_layernorm = Evo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.filter = Evo2HyenaFilter(config, layer_idx)
	self.post_attention_layernorm = Evo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.mlp = Evo2ParallelGatedMLP(config, layer_idx)
	self.hidden_dropout = nn.Dropout(config.hidden_dropout)
	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor],
	position_ids: torch.Tensor,
	past_key_value: Optional[Cache],
	output_attentions: bool,
	use_cache: bool,
	cache_position: Optional[torch.LongTensor],
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
	del attention_mask, past_key_value, output_attentions, use_cache, cache_position, position_ids
	# Keep residual in float32 for numerical stability (Mamba2 technique)
	residual = hidden_states.to(torch.float32) if hidden_states.dtype == torch.bfloat16 else hidden_states
	hidden_states = self.input_layernorm(hidden_states)
	hidden_states = self.filter(hidden_states)
	hidden_states = (residual + self.hidden_dropout(hidden_states).to(residual.dtype)).to(hidden_states.dtype)
	# Keep residual in float32 for numerical stability
	residual = hidden_states.to(torch.float32) if hidden_states.dtype == torch.bfloat16 else hidden_states
	hidden_states = self.post_attention_layernorm(hidden_states)
	hidden_states = self.mlp(hidden_states)
	hidden_states = (residual + self.hidden_dropout(hidden_states).to(residual.dtype)).to(hidden_states.dtype)
	return hidden_states, None, None
	class Evo2DecoderLayer(GradientCheckpointingLayer):
	def __init__(self, config: Evo2Config, layer_type: str, layer_idx: int):
	super().__init__()
	self.layer_type = layer_type
	if layer_type == "attention":
	self.block = Evo2AttentionBlock(config, layer_idx)
	else:
	self.block = Evo2HyenaBlock(config, layer_idx)
	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor],
	position_ids: torch.Tensor,
	past_key_value: Optional[Cache],
	output_attentions: bool,
	use_cache: bool,
	cache_position: Optional[torch.LongTensor],
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
	return self.block(
	hidden_states,
	attention_mask,
	position_ids,
	past_key_value,
	output_attentions,
	use_cache,
	cache_position,
	)
	class Evo2DecoderLayer(GradientCheckpointingLayer):
	def __init__(self, config: Evo2Config, layer_idx: int):
	super().__init__()
	self.is_attention_layer = config.layer_types[layer_idx] == "full_attention"
	if self.is_attention_layer:
	self.self_attn = Evo2Attention(config, layer_idx)
	else:
	self.conv = Evo2ShortConv(config, layer_idx)
	self.feed_forward = Evo2MLP(config, layer_idx)
	self.operator_norm = Evo2RMSNorm(config.hidden_size, eps=config.norm_eps)
	self.ffn_norm = Evo2RMSNorm(config.hidden_size, eps=config.norm_eps)
	def forward(
	self,
	hidden_states: torch.Tensor,
	position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
	attention_mask: torch.Tensor | None = None,
	position_ids: torch.LongTensor | None = None,
	past_key_values: Cache | None = None,
	**kwargs,
	) -> torch.Tensor:
	residual = hidden_states
	if self.is_attention_layer:
	hidden_states, _ = self.self_attn(
	hidden_states=self.operator_norm(hidden_states),
	position_embeddings=position_embeddings,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	**kwargs,
	)
	else:
	hidden_states = self.conv(
	hidden_states=self.operator_norm(hidden_states),
	past_key_values=past_key_values,
	attention_mask=attention_mask,
	seq_idx=kwargs.get("seq_idx"),
	)
	hidden_states = hidden_states + residual
	hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states))
	return hidden_states

[ArthurZucker]

Lfm2PreTrainedModel (attention backends, output recording) + init of the hyena parameters.

Suggested change

	class Evo2PreTrainedModel(PreTrainedModel):
	config_class = Evo2Config
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["Evo2DecoderLayer"]
	def _init_weights(self, module):
	if isinstance(module, nn.Linear):
	nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
	if module.bias is not None:
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
	elif isinstance(module, Evo2RMSNorm):
	module.weight.data.fill_(1.0)
	@auto_docstring
	class Evo2PreTrainedModel(PreTrainedModel):
	config: Evo2Config
	base_model_prefix = "model"
	supports_gradient_checkpointing = True
	_no_split_modules = ["Evo2DecoderLayer"]
	_skip_keys_device_placement = ["past_key_values"]
	_supports_flash_attn = True
	_supports_sdpa = True
	_supports_flex_attn = True
	_can_compile_fullgraph = True
	_supports_attention_backend = True
	_can_record_outputs = {
	"hidden_states": Evo2DecoderLayer,
	"attentions": Evo2Attention,
	}
	@torch.no_grad()
	def _init_weights(self, module):
	super()._init_weights(module)
	if isinstance(module, Evo2ShortConv):
	for param in module.parameters(recurse=False):
	init.normal_(param, std=self.config.initializer_range)

[ArthurZucker]

Lfm2Model as is: DynamicCache(config), create_causal_mask / create_recurrent_attention_mask, capture_outputs, no manual checkpointing or output collection.

Suggested change

	class Evo2Model(Evo2PreTrainedModel):
	def __init__(self, config: Evo2Config):
	super().__init__(config)
	self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
	self.layers = nn.ModuleList(
	[Evo2DecoderLayer(config, layer_type, layer_idx) for layer_idx, layer_type in enumerate(config.layer_types)]
	)
	self.norm = Evo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout)
	self.gradient_checkpointing = False
	self.post_init()
	def get_input_embeddings(self):
	return self.embed_tokens
	def set_input_embeddings(self, value):
	self.embed_tokens = value
	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Cache] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	**kwargs,
	) -> BaseModelOutputWithPast:
	output_attentions = (
	output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False)
	)
	output_hidden_states = (
	output_hidden_states
	if output_hidden_states is not None
	else getattr(self.config, "output_hidden_states", False)
	)
	return_dict = return_dict if return_dict is not None else True
	use_cache = use_cache if use_cache is not None else self.config.use_cache
	if input_ids is not None and inputs_embeds is not None:
	raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
	elif input_ids is not None:
	batch_size, seq_length = input_ids.shape
	elif inputs_embeds is not None:
	batch_size, seq_length = inputs_embeds.shape[:2]
	else:
	raise ValueError("You must specify either input_ids or inputs_embeds")
	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)
	if use_cache:
	if past_key_values is None:
	past_key_values = DynamicCache(config=self.config)
	elif not isinstance(past_key_values, Cache):
	raise TypeError("`past_key_values` must be a `Cache` when `use_cache` is True.")
	else:
	past_key_values = None
	past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
	attention_mask = _prepare_4d_causal_attention_mask(
	attention_mask,
	(batch_size, seq_length),
	inputs_embeds,
	past_length,
	)
	hidden_states = self.dropout(inputs_embeds)
	if cache_position is None:
	cache_position = torch.arange(past_length, past_length + seq_length, device=hidden_states.device)
	else:
	cache_position = cache_position.to(hidden_states.device)
	if position_ids is None:
	position_ids = cache_position.unsqueeze(0).expand(batch_size, -1)
	all_hidden_states = () if output_hidden_states else None
	all_attentions = () if output_attentions else None
	for layer_idx, (decoder_layer, layer_type) in enumerate(zip(self.layers, self.config.layer_types)):
	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)
	layer_past = past_key_values if layer_type == "attention" else None
	if self.gradient_checkpointing and self.training:
	if use_cache:
	use_cache = False
	def create_custom_forward(module):
	def custom_forward(*inputs):
	return module(*inputs)
	return custom_forward
	hidden_states, attn, present_kv = checkpoint(
	create_custom_forward(decoder_layer),
	hidden_states,
	attention_mask,
	position_ids,
	layer_past,
	output_attentions,
	use_cache,
	cache_position,
	)
	else:
	hidden_states, attn, present_kv = decoder_layer(
	hidden_states,
	attention_mask,
	position_ids,
	layer_past,
	output_attentions,
	use_cache,
	cache_position,
	)
	if layer_type == "attention" and present_kv is not None and use_cache:
	past_key_values = present_kv
	if output_attentions:
	all_attentions = all_attentions + (attn,)
	hidden_states = self.norm(hidden_states)
	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)
	if not return_dict:
	outputs = (hidden_states, past_key_values)
	if output_hidden_states:
	outputs += (all_hidden_states,)
	if output_attentions:
	outputs += (all_attentions,)
	return outputs
	return BaseModelOutputWithPast(
	last_hidden_state=hidden_states,
	past_key_values=past_key_values if use_cache else None,
	hidden_states=all_hidden_states,
	attentions=all_attentions,
	)
	@auto_docstring
	class Evo2Model(Evo2PreTrainedModel):
	def __init__(self, config: Evo2Config):
	super().__init__(config)
	self.padding_idx = config.pad_token_id
	self.vocab_size = config.vocab_size
	self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
	self.layers = nn.ModuleList(
	[Evo2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
	)
	self.rotary_emb = Evo2RotaryEmbedding(config=config)
	self.gradient_checkpointing = False
	self.embedding_norm = Evo2RMSNorm(config.hidden_size, eps=config.norm_eps)
	# Initialize weights and apply final processing
	self.post_init()
	@merge_with_config_defaults
	@capture_outputs
	@auto_docstring
	def forward(
	self,
	input_ids: torch.LongTensor | None = None,
	attention_mask: torch.Tensor | None = None,
	position_ids: torch.LongTensor | None = None,
	past_key_values: Cache | None = None,
	inputs_embeds: torch.FloatTensor | None = None,
	use_cache: bool | None = None,
	**kwargs: Unpack[TransformersKwargs],
	) -> BaseModelOutputWithPast:
	if (input_ids is None) ^ (inputs_embeds is not None):
	raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
	if inputs_embeds is None:
	inputs_embeds = self.embed_tokens(input_ids)
	if use_cache and past_key_values is None:
	past_key_values = DynamicCache(config=self.config)
	if position_ids is None:
	past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
	position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
	position_ids = position_ids.unsqueeze(0)
	if not isinstance(causal_mask_mapping := attention_mask, dict):
	mask_kwargs = {
	"config": self.config,
	"inputs_embeds": inputs_embeds,
	"attention_mask": attention_mask,
	"past_key_values": past_key_values,
	"position_ids": position_ids,
	}
	causal_mask_mapping = {
	"full_attention": create_causal_mask(**mask_kwargs),
	"conv": create_recurrent_attention_mask(**mask_kwargs),
	}
	hidden_states = inputs_embeds
	position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
	# decoder layers
	for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
	hidden_states = decoder_layer(
	hidden_states,
	attention_mask=causal_mask_mapping[self.config.layer_types[i]],
	position_embeddings=position_embeddings,
	position_ids=position_ids,
	past_key_values=past_key_values,
	**kwargs,
	)
	hidden_states = self.embedding_norm(hidden_states)
	return BaseModelOutputWithPast(
	last_hidden_state=hidden_states,
	past_key_values=past_key_values,
	)

[ArthurZucker]

Lfm2ForCausalLM as is (self.loss_function, logits_to_keep, tied head, no _reorder_cache).

Suggested change

	class Evo2ForCausalLM(Evo2PreTrainedModel, GenerationMixin):
	_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
	def __init__(self, config: Evo2Config):
	super().__init__(config)
	self.model = Evo2Model(config)
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
	self.post_init()
	def get_input_embeddings(self):
	return self.model.embed_tokens
	def set_input_embeddings(self, value):
	self.model.embed_tokens = value
	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	past_key_values: Optional[Cache] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	cache_position: Optional[torch.LongTensor] = None,
	logits_to_keep: int | torch.Tensor = 0,
	**kwargs,
	) -> CausalLMOutputWithPast:
	return_dict = return_dict if return_dict is not None else True
	outputs = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	cache_position=cache_position,
	**kwargs,
	)
	hidden_states = outputs[0]
	if isinstance(logits_to_keep, int):
	slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None)
	# Mamba2 technique: convert to lm_head dtype, then to float32 for numerical stability
	logits = self.lm_head(hidden_states[:, slice_indices, :].to(self.lm_head.weight.dtype)).float()
	else:
	# Mamba2 technique: convert to lm_head dtype, then to float32 for numerical stability
	logits = self.lm_head(hidden_states[:, logits_to_keep, :].to(self.lm_head.weight.dtype)).float()
	loss = None
	if labels is not None:
	shift_logits = logits[..., :-1, :].contiguous()
	shift_labels = labels[..., 1:].contiguous()
	loss_fct = nn.CrossEntropyLoss()
	loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
	if not return_dict:
	output = (logits,) + outputs[1:]
	return ((loss,) + output) if loss is not None else output
	return CausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)
	def _reorder_cache(self, past_key_values: Cache, beam_idx: torch.LongTensor) -> Cache:
	past_key_values.reorder_cache(beam_idx)
	return past_key_values
	@auto_docstring
	class Evo2ForCausalLM(Evo2PreTrainedModel, GenerationMixin):
	_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
	_tp_plan = {"lm_head": "colwise_gather_output"}
	_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
	_fsdp_plan = {"lm_head": "keep_full_weight"}
	def __init__(self, config):
	super().__init__(config)
	self.model = Evo2Model(config)
	self.vocab_size = config.vocab_size
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
	# Initialize weights and apply final processing
	self.post_init()
	@can_return_tuple
	@auto_docstring
	def forward(
	self,
	input_ids: torch.LongTensor | None = None,
	attention_mask: torch.Tensor | None = None,
	position_ids: torch.LongTensor | None = None,
	past_key_values: Cache | None = None,
	inputs_embeds: torch.FloatTensor | None = None,
	labels: torch.LongTensor | None = None,
	use_cache: bool | None = None,
	logits_to_keep: int | torch.Tensor = 0,
	**kwargs: Unpack[TransformersKwargs],
	) -> CausalLMOutputWithPast:
	r"""
	Example:
	```python
	>>> from transformers import AutoTokenizer, Evo2ForCausalLM
	>>> model = Evo2ForCausalLM.from_pretrained("meta-evo2/Evo2-2-7b-hf")
	>>> tokenizer = AutoTokenizer.from_pretrained("meta-evo2/Evo2-2-7b-hf")
	>>> prompt = "Hey, are you conscious? Can you talk to me?"
	>>> inputs = tokenizer(prompt, return_tensors="pt")
	>>> # Generate
	>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
	>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
	"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
	```"""
	outputs: BaseModelOutputWithPast = self.model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	past_key_values=past_key_values,
	inputs_embeds=inputs_embeds,
	use_cache=use_cache,
	**kwargs,
	)
	hidden_states = outputs.last_hidden_state
	# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
	slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
	logits = self.lm_head(hidden_states[:, slice_indices, :])
	loss = None
	if labels is not None:
	loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
	return CausalLMOutputWithPast(
	loss=loss,
	logits=logits,
	past_key_values=outputs.past_key_values,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)
	__all__ = ["Evo2PreTrainedModel", "Evo2Model", "Evo2ForCausalLM"]

[ArthurZucker]

needs. alot of cleanup on what is actually used

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

• The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
• A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/ — this covers files the PR deletes or renames, not only the ones it edits
• A new security issue was detected in the changed Python files (Bandit check) — a high-severity one, shell execution, exec/eval, a URL fetch, or disabled TLS verification
• The PR touches a path this repository protects from untrusted PRs, such as the file that decides who reviews it — a maintainer must make that change in a separate PR

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, evo2