# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2024 Arc Institute. All rights reserved.
# Copyright (c) 2024 Michael Poli. All rights reserved.
# Copyright (c) 2024 Stanford University. 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.


import os
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Literal, Optional, Type

import torch
from megatron.core import parallel_state
from megatron.core.inference.contexts import StaticInferenceContext
from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import GPTInferenceWrapper
from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import InferenceWrapperConfig
from megatron.core.transformer.enums import AttnBackend
from megatron.core.transformer.transformer_config import TransformerConfig

from nemo.collections.llm.gpt.model.base import GPTModel, gpt_data_step
from nemo.collections.llm.gpt.model.megatron.hyena.hyena_layer_specs import get_hyena_stack_spec
from nemo.collections.llm.gpt.model.megatron.hyena.hyena_model import HyenaModel as MCoreHyenaModel
from nemo.collections.llm.gpt.model.megatron.hyena.hyena_utils import hyena_no_weight_decay_cond
from nemo.lightning import get_vocab_size, io, teardown
from nemo.lightning.base import NEMO_MODELS_CACHE
from nemo.lightning.io.state import TransformFns
from nemo.utils import logging


class HyenaInferenceContext(StaticInferenceContext):
    """Hyena-specific inference context."""

    def reset(self):
        """Reset the inference context."""
        super().reset()  # standard state reset for GPT models
        for key in dir(self):
            # Remove all of the state that we add in hyena.py
            if "filter_state_dict" in key:
                delattr(self, key)


class HyenaModel(GPTModel):
    """
    This is a wrapper around the MCoreHyenaModel to allow for inference. Our model follows the same API
      as the GPTModel, but the megatron model class is different so we need to handle the inference wrapper
      slightly differently.
    """

    def get_inference_wrapper(
        self, params_dtype, inference_batch_times_seqlen_threshold, inference_max_seq_length=None
    ) -> torch.Tensor:
        """
        Gets the inference wrapper for the Hyena model.

        Args:
            params_dtype: The data type for model parameters
            inference_batch_times_seqlen_threshold: Threshold for batch size * sequence length during inference
            inference_max_seq_length: Maximum sequence length for inference

        Returns:
            GPTInferenceWrapper: The inference wrapper for the model

        Raises:
            ValueError: If MCoreHyenaModel instance not found or vocab size cannot be determined
        """
        # This is to get the MCore model required in GPTInferenceWrapper.
        mcore_model = self.module
        while mcore_model:
            if type(mcore_model) is MCoreHyenaModel:
                break
            mcore_model = getattr(mcore_model, "module", None)
        if mcore_model is None or type(mcore_model) is not MCoreHyenaModel:
            raise ValueError("Exact MCoreHyenaModel instance not found in the model structure.")

        vocab_size = None
        if self.tokenizer is not None:
            vocab_size = self.tokenizer.vocab_size
        elif hasattr(self.config, 'vocab_size'):
            vocab_size = self.config.vocab_size
        else:
            raise ValueError(
                'Unable to find vocab size.'
                ' Either pass in a tokenizer with vocab size, or set vocab size in the model config'
            )

        inference_wrapper_config = InferenceWrapperConfig(
            hidden_size=mcore_model.config.hidden_size,
            params_dtype=params_dtype,
            inference_batch_times_seqlen_threshold=inference_batch_times_seqlen_threshold,
            padded_vocab_size=vocab_size,
            inference_max_seq_length=inference_max_seq_length,
            inference_max_requests=1,
        )

        inference_context = HyenaInferenceContext.from_config(inference_wrapper_config)
        model_inference_wrapper = GPTInferenceWrapper(mcore_model, inference_wrapper_config, inference_context)
        return model_inference_wrapper

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        decoder_input: Optional[torch.Tensor] = None,
        loss_mask: Optional[torch.Tensor] = None,
        inference_context=None,
        packed_seq_params=None,
    ) -> torch.Tensor:
        """
        Forward pass of the Hyena model.

        Args:
            input_ids: Input token IDs
            position_ids: Position IDs for input tokens
            attention_mask: Optional attention mask
            labels: Optional labels for loss computation
            decoder_input: Optional decoder input
            loss_mask: Optional loss mask
            inference_context: Optional inference parameters
            packed_seq_params: Optional parameters for packed sequences

        Returns:
            torch.Tensor: Output tensor from the model
        """
        extra_kwargs = {'packed_seq_params': packed_seq_params} if packed_seq_params is not None else {}
        output_tensor = self.module(
            input_ids,
            position_ids,
            attention_mask,
            decoder_input=decoder_input,
            labels=labels,
            inference_context=inference_context,
            loss_mask=loss_mask,
            **extra_kwargs,
        )
        return output_tensor


def hyena_forward_step(model, batch) -> torch.Tensor:
    """
    Performs a forward step for the Hyena model.

    Args:
        model: The Hyena model
        batch: Dictionary containing input batch data with keys:
            - tokens: Input token IDs
            - position_ids: Position IDs
            - labels: Labels for loss computation
            - loss_mask: Mask for loss computation

    Returns:
        torch.Tensor: Output from the model forward pass
    """
    forward_args = {
        "input_ids": batch["tokens"],
        "position_ids": batch["position_ids"],
        "labels": batch["labels"],
        "loss_mask": batch["loss_mask"],
    }
    forward_args["attention_mask"] = None
    return model(**forward_args)


@dataclass
class HyenaConfig(TransformerConfig, io.IOMixin):
    """
    Configuration dataclass for Hyena.

    For adjusting ROPE when doing context extension, set seq_len_interpolation_factor relative to 8192.
    For example, if your context length is 512k, then set the factor to 512k / 8k = 64.
    """

    # From megatron.core.models.hyena.hyena_model.HyenaModel
    fp16_lm_cross_entropy: bool = False
    parallel_output: bool = True
    params_dtype: torch.dtype = torch.bfloat16
    fp16: bool = False
    bf16: bool = True
    num_layers: int = 2
    hidden_size: int = 1024
    num_attention_heads: int = 8
    num_groups_hyena: int = None
    num_groups_hyena_medium: int = None
    num_groups_hyena_short: int = None
    hybrid_attention_ratio: float = 0.0
    hybrid_mlp_ratio: float = 0.0
    hybrid_override_pattern: str = None
    post_process: bool = True
    pre_process: bool = True
    seq_length: int = 2048
    position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'rope'
    rotary_percent: float = 1.0
    rotary_base: int = 10000
    seq_len_interpolation_factor: Optional[float] = None
    apply_rope_fusion: bool = True
    make_vocab_size_divisible_by: int = 128
    gated_linear_unit: bool = True
    fp32_residual_connection: bool = True
    normalization: str = 'RMSNorm'
    add_bias_linear: bool = False
    hidden_dropout: float = 0.0
    attention_dropout: float = 0.0
    layernorm_epsilon: float = 1e-6
    attention_backend: AttnBackend = AttnBackend.flash
    # TODO: Move this to better places?
    get_attention_mask_from_fusion: bool = False
    recompute_granularity: str = 'full'
    recompute_method: str = 'uniform'
    recompute_num_layers: int = 4
    forward_step_fn: Callable = hyena_forward_step
    data_step_fn: Callable = gpt_data_step
    tokenizer_model_path: str = None
    hyena_init_method: str = None
    hyena_output_layer_init_method: str = None
    hyena_filter_no_wd: bool = True
    remove_activation_post_first_layer: bool = True
    add_attn_proj_bias: bool = True
    cross_entropy_loss_fusion: bool = False  # Faster but lets default to False for more precision
    tp_comm_overlap: bool = False
    bias_activation_fusion: bool = True
    bias_dropout_add_fusion: bool = True
    add_bias_output: bool = False
    use_te: bool = True
    to_upper: str = "normalized_weighted"  # choose between "weighted" and "normalized_weighted"
    use_short_conv_bias: bool = False
    # Use this if you want to turn FP8 on for the linear layer in the mixer only. When using this, do not set
    #  Fp8 in the mixed precision plugin.
    vortex_style_fp8: bool = False
    use_subquadratic_ops: bool = False
    share_embeddings_and_output_weights: bool = True
    unfused_rmsnorm: bool = False  # Use unfused RMSNorm + TELinear for dense projection
    plain_row_linear: bool = False  # Use plain pytorch implementation instead of Megatron's row parallel linears

    def __post_init__(self):
        """
        Post-initialization hook that sets up weight decay conditions.
        """
        super().__post_init__()
        self.hyena_no_weight_decay_cond_fn = hyena_no_weight_decay_cond if self.hyena_filter_no_wd else None

    def configure_model(self, tokenizer, vp_stage: Optional[int] = None) -> "MCoreHyenaModel":
        """
        Configures and returns a Hyena model instance based on the config settings.

        Args:
            tokenizer: Tokenizer to use for the model
            vp_stage: Virtual pipeline stage

        Returns:
            MCoreHyenaModel: Configured Hyena model instance
        """
        self.bias_activation_fusion = False if self.remove_activation_post_first_layer else self.bias_activation_fusion

        assert (
            getattr(self, "virtual_pipeline_model_parallel_size", None) is None and vp_stage is None
        ), "Virtual pipeline model parallelism is temporarily unsupported in Hyena."

        model = MCoreHyenaModel(
            self,
            hyena_stack_spec=get_hyena_stack_spec(
                use_te=self.use_te,
                vortex_style_fp8=self.vortex_style_fp8,
                unfused_rmsnorm=self.unfused_rmsnorm,
                plain_row_linear=self.plain_row_linear,
            ),
            vocab_size=get_vocab_size(self, tokenizer.vocab_size, self.make_vocab_size_divisible_by),
            max_sequence_length=self.seq_length,
            num_groups_hyena=self.num_groups_hyena,
            num_groups_hyena_medium=self.num_groups_hyena_medium,
            num_groups_hyena_short=self.num_groups_hyena_short,
            hybrid_override_pattern=self.hybrid_override_pattern,
            position_embedding_type=self.position_embedding_type,
            rotary_percent=self.rotary_percent,
            rotary_base=self.rotary_base,
            seq_len_interpolation_factor=self.seq_len_interpolation_factor,
            pre_process=parallel_state.is_pipeline_first_stage(),
            post_process=parallel_state.is_pipeline_last_stage(),
            share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,
            hyena_init_method=self.hyena_init_method,
            hyena_output_layer_init_method=self.hyena_output_layer_init_method,
            remove_activation_post_first_layer=self.remove_activation_post_first_layer,
            add_attn_proj_bias=self.add_attn_proj_bias,
        )
        return model


@dataclass
class HyenaTestConfig(HyenaConfig):
    """Configuration for testing Hyena models."""

    hybrid_override_pattern: str = "SDH*"
    num_layers: int = 4
    seq_length: int = 8192
    hidden_size: int = 4096
    num_groups_hyena: int = 4096
    num_groups_hyena_medium: int = 256
    num_groups_hyena_short: int = 256
    make_vocab_size_divisible_by: int = 8
    tokenizer_library: str = 'byte-level'
    mapping_type: str = "base"
    ffn_hidden_size: int = 11008
    gated_linear_unit: bool = True
    num_attention_heads: int = 32
    use_cpu_initialization: bool = False
    hidden_dropout: float = 0.0
    attention_dropout: float = 0.0
    params_dtype: torch.dtype = torch.bfloat16
    normalization: str = "RMSNorm"
    add_qkv_bias: bool = False
    add_bias_linear: bool = False
    layernorm_epsilon: float = 1e-6
    recompute_granularity: str = 'full'
    recompute_method: str = 'uniform'
    recompute_num_layers: int = 2
    hyena_init_method: str = 'small_init'
    hyena_output_layer_init_method: str = 'wang_init'
    hyena_filter_no_wd: bool = True
    use_short_conv_bias: bool = False
    use_subquadratic_ops: bool = False


@dataclass
class HyenaNVTestConfig(HyenaTestConfig):
    """
    Several unintentional design choices were made to the original Arc implementation that are required to use the
    original Arc checkpoints, but may result in less stable model training. If you are training from scratch,
    these are the recommended configs.
    """

    remove_activation_post_first_layer: bool = False
    add_attn_proj_bias: bool = False
    use_short_conv_bias: bool = True


@dataclass
class Hyena1bConfig(HyenaConfig):
    """Config matching the 1b 8k context Evo2 model"""

    hybrid_override_pattern: str = "SDH*SDHSDH*SDHSDH*SDHSDH*"
    num_layers: int = 25
    recompute_num_layers: int = 5  # needs to be a multiple of num_layers
    seq_length: int = 8192
    hidden_size: int = 1920
    num_groups_hyena: int = 1920
    num_groups_hyena_medium: int = 128
    num_groups_hyena_short: int = 128
    make_vocab_size_divisible_by: int = 8
    tokenizer_library: str = 'byte-level'
    mapping_type: str = "base"
    ffn_hidden_size: int = 5120
    gated_linear_unit: bool = True
    num_attention_heads: int = 15
    use_cpu_initialization: bool = False
    hidden_dropout: float = 0.0
    attention_dropout: float = 0.0
    params_dtype: torch.dtype = torch.bfloat16
    normalization: str = "RMSNorm"
    add_qkv_bias: bool = False
    add_bias_linear: bool = False
    layernorm_epsilon: float = 1e-6
    recompute_granularity: str = 'full'
    recompute_method: str = 'uniform'
    recompute_num_layers: int = 5
    hyena_init_method: str = 'small_init'
    hyena_output_layer_init_method: str = 'wang_init'
    hyena_filter_no_wd: bool = True


@dataclass
class HyenaNV1bConfig(Hyena1bConfig):
    """
    Several unintentional design choices were made to the original Arc implementation that are required to use the
    original Arc checkpoints, but may result in less stable model training. If you are training from scratch,
    these are the recommended configs.
    """

    remove_activation_post_first_layer: bool = False
    add_attn_proj_bias: bool = False
    use_short_conv_bias: bool = True


@dataclass
class Hyena7bConfig(HyenaConfig):
    """Config matching the 7b 8k context Evo2 model"""

    hybrid_override_pattern: str = "SDH*SDHSDH*SDHSDH*SDHSDH*SDHSDH*"
    num_layers: int = 32
    seq_length: int = 8192
    hidden_size: int = 4096
    num_groups_hyena: int = 4096
    num_groups_hyena_medium: int = 256
    num_groups_hyena_short: int = 256
    make_vocab_size_divisible_by: int = 8
    tokenizer_library: str = 'byte-level'
    mapping_type: str = "base"
    ffn_hidden_size: int = 11008
    gated_linear_unit: bool = True
    num_attention_heads: int = 32
    use_cpu_initialization: bool = False
    hidden_dropout: float = 0.0
    attention_dropout: float = 0.0
    params_dtype: torch.dtype = torch.bfloat16
    normalization: str = "RMSNorm"
    add_qkv_bias: bool = False
    add_bias_linear: bool = False
    layernorm_epsilon: float = 1e-6
    recompute_granularity: str = 'full'
    recompute_method: str = 'uniform'
    recompute_num_layers: int = 4
    hyena_init_method: str = 'small_init'
    hyena_output_layer_init_method: str = 'wang_init'
    hyena_filter_no_wd: bool = True


@dataclass
class HyenaNV7bConfig(Hyena7bConfig):
    """
    Several unintentional design choices were made to the original Arc implementation that are required to use the
    original Arc checkpoints, but may result in less stable model training. If you are training from scratch,
    these are the recommended configs.
    """

    remove_activation_post_first_layer: bool = False
    add_attn_proj_bias: bool = False
    use_short_conv_bias: bool = True
    ffn_hidden_size: int = 11264  # start with the larger FFN hidden size to avoid having to pad during extension.
    rotary_base: int = 1_000_000


@dataclass
class Hyena40bConfig(HyenaConfig):
    """Config matching the 40b 8k context Evo2 model"""

    hybrid_override_pattern: str = "SDH*SDHSDH*SDHSDH*SDHSDH*SDHSDH*SDH*SDHSDH*SDHSDH*"
    num_layers: int = 50
    seq_length: int = 8192
    hidden_size: int = 8192
    num_groups_hyena: int = 8192
    num_groups_hyena_medium: int = 512
    num_groups_hyena_short: int = 512
    make_vocab_size_divisible_by: int = 8
    tokenizer_library: str = 'byte-level'
    mapping_type: str = "base"
    ffn_hidden_size: int = 21888
    gated_linear_unit: bool = True
    num_attention_heads: int = 64
    use_cpu_initialization: bool = False
    hidden_dropout: float = 0.0
    attention_dropout: float = 0.0
    params_dtype: torch.dtype = torch.bfloat16
    normalization: str = "RMSNorm"
    add_qkv_bias: bool = False
    add_bias_linear: bool = False
    layernorm_epsilon: float = 1e-6
    recompute_granularity: str = 'full'
    recompute_method: str = 'uniform'
    recompute_num_layers: int = 2
    hyena_init_method: str = 'small_init'
    hyena_output_layer_init_method: str = 'wang_init'
    hyena_filter_no_wd: bool = True
    rotary_base: int = 1_000_000


@dataclass
class HyenaNV40bConfig(Hyena40bConfig):
    """
    Several unintentional design choices were made to the original Arc implementation that are required to use the
    original Arc checkpoints, but may result in less stable model training. If you are training from scratch,
    these are the recommended configs.
    """

    remove_activation_post_first_layer: bool = False
    add_attn_proj_bias: bool = False
    use_short_conv_bias: bool = True
    ffn_hidden_size: int = 22528  # start with the larger FFN hidden size to avoid having to pad during extension.


@dataclass
class Hyena7bARCLongContextConfig(Hyena7bConfig):
    """The checkpoint from ARC requires padding to the FFN dim
    due to constraintes from large TP size for training."""

    seq_length: int = 1_048_576
    ###
    # Hand verification of RoPE base for 7B-1M @antonvnv
    # >>> def inv_freq(base, dim):
    #     return 1.0 / (base ** (torch.arange(0, dim, 2, device="cuda", dtype=torch.float32) / dim))
    #
    # >>> pt = torch.load("evo2_7b.pt", weights_only=False, mmap=True, map_location="cpu")
    #
    # >>> torch.mean(pt["blocks.3.inner_mha_cls.rotary_emb.inv_freq"] - inv_freq(1e11, 128).cpu())
    # tensor(0.0688)
    #
    # >>> torch.mean(pt["blocks.3.inner_mha_cls.rotary_emb.inv_freq"] - inv_freq(1e6, 128).cpu())
    # tensor(0.0361)
    #
    # >>> torch.mean(pt["blocks.3.inner_mha_cls.rotary_emb.inv_freq"] - inv_freq(1e4, 128).cpu())
    # tensor(5.2014e-06)
    rotary_base: int = 10_000
    ffn_hidden_size: int = 11264
    seq_len_interpolation_factor: float = 128


@dataclass
class Hyena40bARCLongContextConfig(Hyena40bConfig):
    """The checkpoint from ARC requires padding to the FFN dim
    due to constraintes from large TP size for training."""

    seq_length: int = 1_048_576
    ####
    # For 40B-1M hand verification of RoPE base @antonvnv
    # >>> def inv_freq(base, dim):
    #     return 1.0 / (base ** (torch.arange(0, dim, 2, device="cuda", dtype=torch.float32) / dim))
    #
    # >>> pt = torch.load("evo2_40b.pt", weights_only=False, mmap=True, map_location="cpu")
    #
    # >>> torch.mean(pt["blocks.3.inner_mha_cls.rotary_emb.inv_freq"] - inv_freq(1e11, 128).cpu())
    # tensor(0.0326)
    #
    # >>> torch.mean(pt["blocks.3.inner_mha_cls.rotary_emb.inv_freq"] - inv_freq(1e6, 128).cpu())
    # tensor(-2.5294e-05)
    rotary_base: int = 1_000_000
    ffn_hidden_size: int = 22528
    seq_len_interpolation_factor: float = 128


@io.model_importer(HyenaModel, "pytorch")
class PyTorchHyenaImporter(io.ModelConnector["HyenaModel", HyenaModel]):
    """
    Importer class for converting PyTorch Hyena models to NeMo format.
    """

    def __new__(cls, path: str, model_config=None):
        """
        Creates a new importer instance.

        Args:
            path: Path to the PyTorch model
            model_config: Optional model configuration

        Returns:
            PyTorchHyenaImporter instance
        """
        instance = super().__new__(cls, path)
        instance.model_config = model_config
        return instance

    def init(self) -> HyenaModel:
        """
        Initializes a new HyenaModel instance.

        Returns:
            HyenaModel: Initialized model
        """
        return HyenaModel(self.config, tokenizer=self.tokenizer)

    def get_source_model(self):
        """
        Returns the source model.
        """
        return torch.load(str(self), map_location='cpu')

    def apply(self, output_path: Path, checkpoint_format: str = 'torch_dist') -> Path:
        """
        Applies the model conversion from PyTorch to NeMo format.

        Args:
            output_path: Path to save the converted model
            checkpoint_format: Format for saving checkpoints

        Returns:
            Path: Path to the saved NeMo model
        """
        source = self.get_source_model()

        if 'model' in source:
            source = source['model']

        class ModelState:
            """Wrapper around the source model state dictionary that also handles some weight transformations."""

            def __init__(self, state_dict, num_layers, fp32_suffixes):
                """Wrapper around the source model state dictionary that also handles some weight transformations.

                Args:
                    state_dict: original state dictionary from the source model
                    num_layers: number of layers in the source model
                """
                self.num_layers = num_layers
                state_dict = self.transform_source_dict(state_dict)
                self._state_dict = state_dict
                self.fp32_suffixes = fp32_suffixes

            def state_dict(self):
                """Return the state dictionary."""
                return self._state_dict

            def to(self, dtype):
                """Convert the state dictionary to the target dtype."""
                for k, v in self._state_dict.items():
                    if "_extra" not in k:
                        if v.dtype != dtype:
                            logging.warning(f"Converting {k} from {v.dtype} (source model) to {dtype} (target model)")
                        k_suffix = k.split('.')[-1]
                        if k_suffix in self.fp32_suffixes:
                            _dtype = torch.float32
                        else:
                            _dtype = dtype
                        self._state_dict[k] = v.to(_dtype)

            def adjust_medium_filter(self, updated_data):
                """Adjust the medium filter."""
                from nemo.collections.llm.gpt.model.megatron.hyena.hyena_config import HyenaConfig

                for k, v in updated_data.items():
                    if "filter.h" in k or "filter.decay" in k:
                        updated_data[k] = v[:, : HyenaConfig().hyena_medium_conv_len]
                return updated_data

            def transform_source_dict(self, source):
                """Transform the source state dictionary, applying some challenging layer name re-mappings and
                removing extra keys, as well as truncating a filter that didn't need to extend to the full
                sequence length dim.
                """
                import re

                layer_map = {i + 2: i for i in range(self.num_layers)}
                layer_map[self.num_layers + 3] = self.num_layers
                updated_data = {}

                for key in list(source['module'].keys()):
                    if "_extra" in key:
                        source['module'].pop(key)
                    else:
                        match = re.search(r'sequential\.(\d+)', key)
                        if match:
                            original_layer_num = int(match.group(1))
                            if original_layer_num in layer_map:
                                # Create the updated key by replacing the layer number
                                new_key = re.sub(rf'\b{original_layer_num}\b', str(layer_map[original_layer_num]), key)
                                updated_data[new_key] = source['module'][key]
                            else:
                                # Keep the key unchanged if no mapping exists
                                updated_data[key] = source['module'][key]
                        else:
                            updated_data[key] = source['module'][key]
                updated_data = self.adjust_medium_filter(updated_data)
                return updated_data

        target = self.init()
        trainer = self.nemo_setup(target, ckpt_async_save=False, save_ckpt_format=checkpoint_format)
        target.to(self.config.params_dtype)
        fp32_suffixes = {n.split('.')[-1] for n, p in target.named_parameters() if p.dtype == torch.float32}
        source = ModelState(source, self.config.num_layers, fp32_suffixes)
        source.to(self.config.params_dtype)
        self.convert_state(source, target)
        self.nemo_save(output_path, trainer)

        logging.info(f"Converted Hyena model to Nemo, model saved to {output_path}")

        teardown(trainer, target)
        del trainer, target

        return output_path

    def convert_state(self, source, target):
        """
        Converts the state dictionary from source format to target format.

        Args:
            source: Source model state
            target: Target model

        Returns:
            Result of applying state transforms
        """
        mapping = {}
        mapping['sequential.0.word_embeddings.weight'] = 'embedding.word_embeddings.weight'
        mapping[f'sequential.{len(self.config.hybrid_override_pattern)}.norm.weight'] = 'decoder.final_norm.weight'
        te_enabled = self.config.use_te
        for i, symbol in enumerate(self.config.hybrid_override_pattern):
            if te_enabled:
                mapping[f'sequential.{i}.pre_mlp_layernorm.weight'] = (
                    f'decoder.layers.{i}.mlp.linear_fc1.layer_norm_weight'
                )
            else:
                mapping[f'sequential.{i}.pre_mlp_layernorm.weight'] = f'decoder.layers.{i}.pre_mlp_layernorm.weight'
            mapping[f'sequential.{i}.mlp.w3.weight'] = f'decoder.layers.{i}.mlp.linear_fc2.weight'

            if symbol != '*':
                if te_enabled:
                    mapping[f'sequential.{i}.input_layernorm.weight'] = (
                        f'decoder.layers.{i}.mixer.dense_projection.layer_norm_weight'
                    )
                else:
                    mapping[f'sequential.{i}.input_layernorm.weight'] = f'decoder.layers.{i}.norm.weight'

                mapping[f'sequential.{i}.mixer.dense_projection.weight'] = (
                    f'decoder.layers.{i}.mixer.dense_projection.weight'
                )
                mapping[f'sequential.{i}.mixer.hyena_proj_conv.short_conv_weight'] = (
                    f'decoder.layers.{i}.mixer.hyena_proj_conv.short_conv_weight'
                )
                mapping[f'sequential.{i}.mixer.dense.weight'] = f'decoder.layers.{i}.mixer.dense.weight'
                mapping[f'sequential.{i}.mixer.dense.bias'] = f'decoder.layers.{i}.mixer.dense.bias'

                if symbol == 'S':
                    mapping[f'sequential.{i}.mixer.mixer.short_conv.short_conv_weight'] = (
                        f'decoder.layers.{i}.mixer.mixer.short_conv.short_conv_weight'
                    )

                elif symbol == 'D':
                    mapping[f'sequential.{i}.mixer.mixer.conv_bias'] = f'decoder.layers.{i}.mixer.mixer.conv_bias'
                    mapping[f'sequential.{i}.mixer.mixer.filter.h'] = f'decoder.layers.{i}.mixer.mixer.filter.h'
                    mapping[f'sequential.{i}.mixer.mixer.filter.decay'] = (
                        f'decoder.layers.{i}.mixer.mixer.filter.decay'
                    )

                elif symbol == 'H':
                    mapping[f'sequential.{i}.mixer.mixer.conv_bias'] = f'decoder.layers.{i}.mixer.mixer.conv_bias'
                    mapping[f'sequential.{i}.mixer.mixer.filter.gamma'] = (
                        f'decoder.layers.{i}.mixer.mixer.filter.gamma'
                    )
                    mapping[f'sequential.{i}.mixer.mixer.filter.R'] = f'decoder.layers.{i}.mixer.mixer.filter.R'
                    mapping[f'sequential.{i}.mixer.mixer.filter.p'] = f'decoder.layers.{i}.mixer.mixer.filter.p'

            elif symbol == '*':
                if te_enabled:
                    mapping[f'sequential.{i}.input_layernorm.weight'] = (
                        f'decoder.layers.{i}.self_attention.linear_qkv.layer_norm_weight'
                    )
                else:
                    mapping[f'sequential.{i}.input_layernorm.weight'] = f'decoder.layers.{i}.input_layernorm.weight'

                mapping[f'sequential.{i}.mixer.dense_projection.weight'] = (
                    f'decoder.layers.{i}.self_attention.linear_qkv.weight'
                )
                mapping[f'sequential.{i}.mixer.dense.weight'] = f'decoder.layers.{i}.self_attention.linear_proj.weight'
                mapping[f'sequential.{i}.mixer.dense.bias'] = f'decoder.layers.{i}.self_attention.linear_proj.bias'
            else:
                raise ValueError(f'Unknown symbol: {symbol}')

        return io.apply_transforms(
            source,
            target,
            mapping=mapping,
            transforms=[
                # Transforms that are more complicated than a simple mapping of an old key name to a new one:
                io.state_transform(
                    source_key=("sequential.*.mlp.w1.weight", "sequential.*.mlp.w2.weight"),
                    target_key="decoder.layers.*.mlp.linear_fc1.weight",
                    fn=TransformFns.merge_fc1,
                )
            ],
        )

    @property
    def tokenizer(self):
        """
        Gets the tokenizer for the model.

        Returns:
            Tokenizer instance
        """
        from nemo.collections.common.tokenizers.tokenizer_utils import get_nmt_tokenizer

        tokenizer = get_nmt_tokenizer(
            library=self.model_config.tokenizer_library,
        )

        return tokenizer

    @property
    def config(self) -> HyenaConfig:
        """
        Gets the model configuration.

        Returns:
            HyenaConfig: Model configuration
        """
        return self.model_config


@io.model_importer(HyenaModel, "hf")
class HuggingFaceSavannaHyenaImporter(PyTorchHyenaImporter):
    """
    Importer class for converting HuggingFace Savanna Hyena models to NeMo format.
        See: https://huggingface.co/arcinstitute/savanna_evo2_7b for an example of a savanna model that this can
        import and convert to NeMo format. Any of the Arc models that start with "savanna_" should work.
    """

    def get_source_model(self):
        """
        Returns the source model.
        """
        import huggingface_hub.errors
        from huggingface_hub import hf_hub_download

        if os.path.exists(str(self)):
            logging.info(f"Loading model from local path {str(self)}")
            return torch.load(str(self), map_location='cpu')
        else:
            if ":" in str(self):
                repo_id, revision = str(self).split(":")
            else:
                repo_id = str(self)
                revision = None
            # See HF download logic here:
            #   https://github.com/ArcInstitute/evo2/blob/96ac9d9cd/evo2/models.py#L191-L231
            modelname = repo_id.split("/")[-1]
            download_dir = str(NEMO_MODELS_CACHE / repo_id)
            weights_filename = f"{modelname}.pt"
            try:
                weights_path = hf_hub_download(
                    repo_id=repo_id, local_dir=download_dir, revision=revision, filename=weights_filename
                )
            except Exception:
                # Try downloading multi-part
                # If file is split, download and join parts
                logging.warning(f"Single path download failed, try loading checkpoint shards for {modelname}")
                # If file is split, get the first part's directory to use the same cache location
                weights_path = os.path.join(download_dir, weights_filename)
                if os.path.exists(weights_path):
                    logging.info(f"Found {weights_path}")
                else:
                    # Download and join parts
                    parts = []
                    part_num = 0
                    while True:
                        try:
                            part_path = hf_hub_download(
                                repo_id=repo_id,
                                local_dir=download_dir,
                                revision=revision,
                                filename=f"{weights_filename}.part{part_num}",
                            )
                            parts.append(part_path)
                            part_num += 1
                        except huggingface_hub.errors.EntryNotFoundError:
                            break

                    # Join in the same directory
                    with open(weights_path, 'wb') as outfile:
                        for part in parts:
                            with open(part, 'rb') as infile:
                                while True:
                                    chunk = infile.read(8192 * 1024)
                                    if not chunk:
                                        break
                                    outfile.write(chunk)

                    # Cleaning up the parts
                    for part in parts:
                        try:
                            os.remove(part)
                        except OSError as e:
                            print(f"Error removing {part}: {e}")
                        print("Cleaned up shards, final checkpoint saved to", weights_path)

        return torch.load(weights_path, map_location='cpu')


HYENA_MODEL_OPTIONS: dict[str, Type[HyenaConfig]] = {
    "1b": Hyena1bConfig,
    "1b_nv": HyenaNV1bConfig,
    "7b": Hyena7bConfig,
    "7b_arc_longcontext": Hyena7bARCLongContextConfig,
    "7b_nv": HyenaNV7bConfig,
    "40b": Hyena40bConfig,
    "40b_arc_longcontext": Hyena40bARCLongContextConfig,
    "40b_nv": HyenaNV40bConfig,
    "test": HyenaTestConfig,
    "test_nv": HyenaNVTestConfig,
}


__all__ = [
    "HyenaConfig",
    "Hyena7bConfig",
    "HyenaNV7bConfig",
    "Hyena1bConfig",
    "HyenaNV1bConfig",
    "Hyena40bConfig",
    "HyenaNV40bConfig",
    "Hyena7bARCLongContextConfig",
    "Hyena40bARCLongContextConfig",
    "HyenaTestConfig",
    "HyenaNVTestConfig",
    "HYENA_MODEL_OPTIONS",
]
