### Install AHN and Dependencies for Training (Bash)

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Installs the AHN repository and its required forked libraries (flash-linear-attention, LLaMA-Factory) using pip. It also provides an optional installation for the Mamba version and installs AHN in editable mode with training extras. Ensure Python 3.11, CUDA 12.4, and PyTorch 2.5.1+cu124 are set up.

```bash
# Install dependencies and set up AHN for training

# 1. Clone the AHN repository and move into it
git clone https://github.com/ByteDance-Seed/AHN.git
cd AHN

# 2. Install required forked libraries
pip install "git+https://github.com/Seerkfang/flash-linear-attention.git@main#egg=flash-linear-attention"
pip install "git+https://github.com/Seerkfang/LLaMA-Factory.git@main#egg=llamafactory"
# (Optional) Install the forked Mamba version if you plan to use AHN-Mamba2
# MAMBA_FORCE_BUILD=TRUE pip install "git+https://github.com/yuweihao/mamba.git"

# 3. Install AHN in editable mode with training extras
pip install -e ".[train]"
```

--------------------------------

### Install AHN Dependencies and Core Libraries

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This script outlines the steps to clone the AHN repository, install forked dependencies for AHN support (flash-linear-attention, LLaMA-Factory), and install the AHN package in editable mode with optional training or evaluation dependencies.

```bash
# Clone repository
git clone https://github.com/ByteDance-Seed/AHN.git
cd AHN

# Install forked dependencies for AHN support
pip install "git+https://github.com/Seerkfang/flash-linear-attention.git@main#egg=flash-linear-attention"
pip install "git+https://github.com/Seerkfang/LLaMA-Factory.git@main#egg=llamafactory"

# Optional: Install Mamba2 if using AHN-Mamba2
# MAMBA_FORCE_BUILD=TRUE pip install "git+https://github.com/yuweihao/mamba.git"

# Install AHN in editable mode with training dependencies
pip install -e ".[train]"

# For evaluation only
pip install -e ".[eval]"

# Verify installation
python -c "from ahn.rnn import GatedDeltaNet, DeltaNet, Mamba2; print('AHN installed successfully')"
```

--------------------------------

### Install Evaluation Dependencies with Pip

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Installs the necessary dependencies for evaluation using pip. This command ensures that all required packages, including those for evaluation, are available in the current Python environment.

```bash
pip install -e ".[eval]"
```

--------------------------------

### Bash Script for LongBench Evaluation

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This bash script details the process for evaluating AHN models on the LongBench benchmark, supporting up to 32K context length. It includes steps for installing dependencies, setting up model paths and parameters, and running full benchmark evaluations or individual dataset predictions.

```bash
# Install evaluation dependencies
pip install -e ".[eval]"

# LongBench evaluation (32K context)
cd eval/longbench

MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
METHOD=ahn
MAX_INPUT_LENGTH=32000
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=8064
OUTPUT_DIR=./eval_results

# Run full benchmark
bash eval.sh $MODEL_NAME $MERGED_MODEL_PATH $METHOD \
  $MAX_INPUT_LENGTH $NUM_ATTENTION_SINK $SLIDING_WINDOW $OUTPUT_DIR

# Or evaluate individual dataset
CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name $MODEL_NAME \
  --model_path $MERGED_MODEL_PATH \
  --max_length $MAX_INPUT_LENGTH \
  --method $METHOD \
  --dataset hotpotqa \
  --attention_sink $NUM_ATTENTION_SINK \
  --sliding_window $SLIDING_WINDOW \
  --results_dir $OUTPUT_DIR

python eval.py \
  --model $MODEL_NAME \
  --results_path $OUTPUT_DIR/${MODEL_NAME}_${MAX_INPUT_LENGTH}

```

--------------------------------

### Merge Base Model and AHN Weights for Inference (Bash)

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Provides a bash command example for merging base model weights with AHN weights for inference. This example specifically shows how to merge the Qwen2.5-3B-Instruct model with the GatedDeltaNet AHN module. The `BASE_MODEL` variable should be set to the repository ID or local path of the base model.

```bash
# Base model (repo_id or local path)
BASE_MODEL=Qwen/Qwen2.5-3B-Instruct

# (Example command for merging weights - actual command omitted as it's not provided in text)
```

--------------------------------

### Python Training Entry Point for AHN Models

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This Python script serves as the entry point for training AHN-augmented models using the LLaMA-Factory framework. It registers custom Qwen2 and Qwen3 models and then initiates the training process, expecting configuration via command-line arguments.

```python
# src/train.py - Training entry point
from llamafactory.train.tuner import run_exp
from ahn.transformer.qwen2_ahn import register_customized_qwen2
from ahn.transformer.qwen3_ahn import register_customized_qwen3

def main():
    # Register custom AHN-augmented models
    register_customized_qwen2()
    register_customized_qwen3()

    # Run training via LLaMA-Factory
    # All config passed via CLI args
    run_exp()

if __name__ == "__main__":
    main()

```

--------------------------------

### Run LongBench Inference and Evaluation

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Executes LongBench inference and evaluation together. This script requires specifying model details, inference method, input length, attention sink tokens, sliding window size, and output directory.

```bash
cd eval/longbench

MODEL_NAME=qwen2.5-3b-ahn-gdn # model identifier, the results will be save in $OUTPUT_DIR/$MODEL_NAME
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN # path to weights of base model and AHN
METHOD=ahn # inference method
MAX_INPUT_LENGTH=32000 # maximum input tokens
NUM_ATTENTION_SINK=128 # number of attention sink tokens
SLIDING_WINDOW=8064 # sliding window size
OUTPUT_DIR=./eval_results # directory to save predictions

bash eval.sh $MODEL_NAME $MERGED_MODEL_PATH $METHOD $MAX_INPUT_LENGTH $NUM_ATTENTION_SINK $SLIDING_WINDOW $OUTPUT_DIR
```

--------------------------------

### Run InfiniteBench Inference Separately with Bash

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Executes the inference step for InfiniteBench separately using a bash script. This script requires model details, maximum input length, inference method, dataset, split, attention sink tokens, sliding window size, and the directory for results.

```bash
cd eval/longbench

MODEL_NAME=qwen2.5-3b-ahn-gdn # model identifier, the results will be save in $OUTPUT_DIR/$MODEL_NAME
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN # path to weights of base model and AHN
METHOD=ahn # inference method
MAX_INPUT_LENGTH=128000 # maximum input tokens
NUM_ATTENTION_SINK=128 # number of attention sink tokens
SLIDING_WINDOW=32640 # sliding window size
DATASET=infinitebench
SPLIT=longbook_qa_eng # or longbook_qa_chn
OUTPUT_DIR=./eval_results # directory to save predictions


CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name $MODEL_NAME \
  --model_path $MERGED_MODEL_PATH \
  --max_length $MAX_INPUT_LENGTH \
  --method $METHOD \
  --dataset $DATASET \
  --split $SPLIT \
  --attention_sink $NUM_ATTENTION_SINK \
  --sliding_window $SLIDING_WINDOW \
  --results_dir $OUTPUT_DIR
```

--------------------------------

### AHN Project Configuration in pyproject.toml

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This TOML file defines the project's metadata, including its name, version, description, and Python version requirements. It also specifies optional dependencies for training and evaluation.

```toml
# pyproject.toml - Project configuration
[project]
name = "ahn"
version = "0.1.0"
description = "Artificial Hippocampus Network"
requires-python = ">=3.9"

[project.optional-dependencies]
train = [
  "tensorflow==2.18.0",
  "tf-keras==2.18.0",
  "deepspeed==0.16.2",
  "numpy==1.26.4",
  "datasets==3.6.0",
  "transformers==4.51.0",
  "liger-kernel"
]

eval = [
  "fuzzywuzzy",
  "jieba",
  "rouge"
]
```

--------------------------------

### Run LongBench Inference Separately with Bash

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Executes the inference step for LongBench separately using a bash script. This script requires model details, maximum input length, inference method, dataset, attention sink tokens, sliding window size, and the directory for results.

```bash
cd eval/longbench

MODEL_NAME=qwen2.5-3b-ahn-gdn # model identifier, the results will be save in $OUTPUT_DIR/$MODEL_NAME
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN # path to weights of base model and AHN
METHOD=ahn # inference method
MAX_INPUT_LENGTH=32000 # maximum input tokens
NUM_ATTENTION_SINK=128 # number of attention sink tokens
SLIDING_WINDOW=8064 # sliding window size
DATASET=hotpotqa # one of ("dureader" "hotpotqa" "musique" "narrativeqa" "qmsum" "triviaqa")
OUTPUT_DIR=./eval_results # directory to save predictions

CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name $MODEL_NAME \
  --model_path $MERGED_MODEL_PATH \
  --max_length $MAX_INPUT_LENGTH \
  --method $METHOD \
  --dataset $DATASET \
  --attention_sink $NUM_ATTENTION_SINK \
  --sliding_window $SLIDING_WINDOW \
  --results_dir $OUTPUT_DIR
```

--------------------------------

### Run LongBench Evaluation Separately with Bash

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Executes the evaluation step for LongBench separately using a bash command. This command requires the model name and the path to the inference results.

```bash
python eval.py \
  --model $MODEL_NAME \
  --results_path $OUTPUT_DIR/${MODEL_NAME}_${MAX_INPUT_LENGTH}
```

--------------------------------

### Train Qwen2.5-3B-Instruct with AHN-GDN

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Initiates training for the Qwen2.5-3B-Instruct model using GatedDeltaNet (GDN) as the AHN module on the ChatQA2 dataset. This script handles configuration for model name, dataset, loss type, AHN implementation, AHN position, sliding window type, saving AHN-only weights, and sequence length filtering/cutoff.

```bash
bash ./examples/scripts/train_qwen2.5_3b_ahn_gdn.sh
```

--------------------------------

### Command-line Inference with AHN Models

Source: https://context7.com/bytedance-seed/ahn/llms.txt

Executes text generation using AHN-augmented models via the command line. This script allows specifying model paths, sliding window parameters, and generation configurations for inference.

```bash
# Command-line inference
PROMPT="Summarize the key points from the long document above."
CUDA_VISIBLE_DEVICES=0 python ./examples/scripts/inference.py \
  --model ./merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN \
  --sliding-window 32640 \
  --num-attention-sink 128 \
  --dtype bfloat16 \
  --max-new-tokens 1024 \
  --prompt "$PROMPT"
```

--------------------------------

### Evaluate Long-Context Models with InfiniteBench

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This script performs inference and evaluation for long-context models on the InfiniteBench dataset. It requires specifying model paths, maximum context length, method, dataset split, and attention/sliding window parameters. The results are saved to a specified directory.

```bash
cd eval/longbench

DATASET=infinitebench
SPLIT=longbook_qa_eng  # or longbook_qa_chn

CUDA_VISIBLE_DEVICES=0 python pred.py \
  --model_name qwen2.5-3b-ahn-gdn \
  --model_path ../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN \
  --max_length 128000 \
  --method ahn \
  --dataset $DATASET \
  --split $SPLIT \
  --attention_sink 128 \
  --sliding_window 32640 \
  --results_dir ./eval_results

python eval.py \
  --model qwen2.5-3b-ahn-gdn \
  --results_path ./eval_results/qwen2.5-3b-ahn-gdn_128000
```

--------------------------------

### Train Model with Multiple Custom Datasets and Sampling

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This command initiates the training process using multiple custom datasets with specified sampling probabilities. Key parameters include model path, dataset names, interleaving probabilities, filter length, cutoff length, and maximum samples.

```bash
# Use multiple datasets with sampling probabilities
torchrun $DISTRIBUTED_ARGS src/train.py \
    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
    --dataset chatqa2,my_custom_dataset \
    --interleave_probs 0.7,0.3 \
    --filter_len 288 \
    --cutoff_len 24576 \
    --max_samples 100000 \
    # ... other training args
```

--------------------------------

### Run LV-Eval Inference and Evaluation with Python

Source: https://github.com/bytedance-seed/ahn/blob/main/eval/README.md

Performs inference and evaluation using the LV-Eval framework with Python scripts. It requires detailed model and configuration parameters, including model name, path, method, maximum length, attention sizes, and output directory.

```python
cd eval/lveval

MODEL_NAME=qwen2.5-3b-ahn-gdn # model identifier, the results will be save in $OUTPUT_DIR/$MODEL_NAME
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN # path to weights of base model and AHN
METHOD=ahn # inference method
MAX_INPUT_LENGTH=256000 # maximum input tokens
NUM_ATTENTION_SINK=128 # number of attention sink tokens
SLIDING_WINDOW=32640 # sliding window size
OUTPUT_DIR=./eval_results # directory to save predictions

python3 pred.py \
  --model-name $MODEL_NAME \
  --model-path $MERGED_MODEL_PATH \
  --method $METHOD \
  --model-max-len $MAX_INPUT_LENGTH \
  --start_size $NUM_ATTENTION_SINK \
  --recent_size $SLIDING_WINDOW \
  --output-dir $OUTPUT_DIR

python3 eval.py \
  --input-dir "${OUTPUT_DIR}/${MODEL_NAME}_${METHOD}_as${NUM_ATTENTION_SINK}_sw${SLIDING_WINDOW}"
```

--------------------------------

### Configure Custom Datasets with LLaMA-Factory

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This JSON configuration file defines custom datasets for training, specifying file names, formatting, and column mappings. It allows for flexible integration of user-provided data into the training pipeline.

```json
// data/dataset_info.json
{
  "chatqa2": {
    "file_name": "chatqa2.jsonl",
    "formatting": "sharegpt",
    "columns": {
      "messages": "conversations"
    },
    "tags": {
      "role_tag": "role",
      "content_tag": "content",
      "user_tag": "user",
      "assistant_tag": "assistant"
    }
  },
  "my_custom_dataset": {
    "file_name": "my_data.jsonl",
    "formatting": "sharegpt",
    "columns": {
      "messages": "messages"
    }
  }
}
```

--------------------------------

### Bash Script for LV-Eval Evaluation

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This bash script outlines the evaluation process for AHN models on the LV-Eval benchmark, supporting up to 256K context length. It covers generating predictions and evaluating them, specifying parameters like model name, path, method, and attention/sliding window sizes.

```bash
# LV-Eval evaluation (256K context)
cd eval/lveval

MODEL_NAME=qwen2.5-3b-ahn-gdn
MERGED_MODEL_PATH=../../merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN
METHOD=ahn
MAX_INPUT_LENGTH=256000
NUM_ATTENTION_SINK=128
SLIDING_WINDOW=32640
OUTPUT_DIR=./eval_results

# Generate predictions
python3 pred.py \
  --model-name $MODEL_NAME \
  --model-path $MERGED_MODEL_PATH \
  --method $METHOD \
  --model-max-len $MAX_INPUT_LENGTH \
  --start_size $NUM_ATTENTION_SINK \
  --recent_size $SLIDING_WINDOW \
  --output-dir $OUTPUT_DIR

# Evaluate predictions
python3 eval.py \
  --input-dir ${OUTPUT_DIR}/${MODEL_NAME}_${METHOD}_as${NUM_ATTENTION_SINK}_sw${SLIDING_WINDOW}

# Expected output format:
# {
#   "avg_score": 0.xyz,
#   "task_scores": {
#     "topic_retrieval": 0.abc,
#     "longdialogue_qa": 0.def,
#     ...
#   }
# }

```

--------------------------------

### Debug Script on Single GPU

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Executes a debugging script on a single specified GPU (GPU 0). This is a quick way to identify and resolve issues within the project's scripts.

```bash
CUDA_VISIBLE_DEVICES=0 bash ./examples/scripts/debug.sh
```

--------------------------------

### Python Implementation of GatedDeltaNet Memory Module

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This Python code demonstrates the initialization and forward pass of the GatedDeltaNet, a recommended RNN memory module for AHN models. It compresses attention cache into fixed-size states, supporting chunk or fused recurrent kernel modes and optional output gating and short convolutions.

```python
# GatedDeltaNet - Recommended for best performance
from ahn.rnn import GatedDeltaNet
import torch

# Initialize GatedDeltaNet layer
gdn = GatedDeltaNet(
    hidden_size=2048,        # Model dimension
    expand_v=1.0,            # Value expansion ratio
    head_dim=256,            # Dimension per head
    num_heads=6,             # Number of attention heads
    mode="chunk",            # Kernel: "chunk" or "fused_recurrent"
    use_gate=True,           # Use output gating
    use_short_conv=True,     # Use short convolutions
    conv_size=4,             # Convolution kernel size
    layer_idx=0              # Layer index in model
)

# Forward pass
batch_size, seq_len, hidden_dim = 2, 1024, 2048
hidden_states = torch.randn(batch_size, seq_len, hidden_dim).cuda()
attention_mask = torch.ones(batch_size, seq_len).cuda()

output, past_key_value = gdn(
    hidden_states=hidden_states,
    attention_mask=attention_mask,
    past_key_values=None,  # Initialize empty for first pass
    use_cache=True
)

# Output shape: (batch_size, seq_len, hidden_dim)
# past_key_value: Compressed memory state for next iteration

```

--------------------------------

### Run Model Inference

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Performs inference using a merged model on a specified GPU. Takes a text prompt as input and generates a response. This is useful for testing the model's capabilities with specific queries.

```bash
PROMPT="When was the concept of AI introduced?"
CUDA_VISIBLE_DEVICES=0 python ./examples/scripts/inference.py \
  --model $MERGED_MODEL_PATH \
  --prompt "$PROMPT"
```

--------------------------------

### Python Implementation of Mamba2 Memory Module

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This Python code illustrates the initialization and forward pass of Mamba2, a state-space model integration for AHN memory. It supports selective state updates and incorporates convolutional biases for enhanced performance.

```python
# Mamba2 - State space model integration
from ahn.rnn import Mamba2

mamba2 = Mamba2(
    hidden_size=2048,
    num_heads=8,
    head_dim=256,
    expand=2,
    chunk_size=256,
    use_conv_bias=False,
    layer_idx=0
)

# Mamba2 forward with selective state updates
output = mamba2(
    hidden_states=hidden_states,
    attention_mask=attention_mask
)

```

--------------------------------

### Merge Base and AHN Model Weights

Source: https://github.com/bytedance-seed/ahn/blob/main/README.md

Merges base model weights with AHN (Artificial Hippocampus Network) parameters to create a unified model. Requires the path to the base model and the AHN model, outputting the merged model to a specified directory. This process enhances the model's long-context understanding capabilities.

```bash
python ./examples/scripts/utils/merge_weights.py \
  --base-model $BASE_MODEL \
  --ahn-path $AHN_PATH \
  --output-path $MERGED_MODEL_PATH
```

--------------------------------

### Train AHN Memory Modules via Self-Distillation

Source: https://context7.com/bytedance-seed/ahn/llms.txt

Trains AHN memory modules using self-distillation while keeping base model weights frozen. This script configures distributed training, specifies AHN module implementations, dataset, and training parameters.

```bash
# examples/scripts/train_qwen2.5_3b_ahn_gdn.sh
# Train Qwen2.5-3B-Instruct with GatedDeltaNet AHN module on 32 GPUs

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export WANDB_PROJECT="AHN"

DISTRIBUTED_ARGS="
    --nnodes=4 \
    --node_rank=0 \
    --nproc_per_node=8 \
    --master_addr=localhost \
    --master_port=29500
"

torchrun $DISTRIBUTED_ARGS src/train.py \
    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
    --stage sft \
    --finetuning_type freeze \
    --freeze_extra_modules ahn \
    --loss_type kl \
    --use_normalized_l2 True \
    --deepspeed examples/deepspeed/ds_z3_config.json \
    --do_train \
    --flash_attn fa2 \
    --sliding_window 256 \
    --sliding_window_type random \
    --ahn_position random \
    --layer_implementation Qwen2MemDecoderLayer \
    --ahn_implementation GatedDeltaNet \
    --dataset chatqa2 \
    --template qwen \
    --filter_len 288 \
    --cutoff_len 24576 \
    --max_samples 100000 \
    --output_dir ./ckpt/qwen2.5-3b_ahn_gdn \
    --save_ahn_only True \
    --logging_steps 1 \
    --save_steps 500 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 4 \
    --learning_rate 1e-4 \
    --num_train_epochs 1 \
    --lr_scheduler_type cosine \
    --warmup_ratio 0.1 \
    --weight_decay 0.1 \
    --bf16 true \
    --report_to wandb

# Training outputs:

```

--------------------------------

### Python Implementation of DeltaNet Memory Module

Source: https://context7.com/bytedance-seed/ahn/llms.txt

This Python code shows the initialization and forward pass of DeltaNet, an efficient linear attention variant for AHN memory. It supports chunk or fused recurrent modes and uses beta gating with optional L2 normalization for query-key activations.

```python
# DeltaNet - Efficient linear attention variant
from ahn.rnn import DeltaNet

deltanet = DeltaNet(
    mode='chunk',
    hidden_size=1024,
    expand_k=1.0,
    expand_v=1.0,
    num_heads=4,
    head_dim=128,
    use_beta=True,
    use_gate=True,
    qk_activation='silu',
    qk_norm='l2'
)

# Forward pass with streaming
output, cache = deltanet(
    hidden_states=hidden_states,
    past_key_values=None
)

```

--------------------------------

### Merge AHN Model Weights with Base Transformer

Source: https://context7.com/bytedance-seed/ahn/llms.txt

Merges the weights of a trained AHN module with a base Transformer model to create a unified model for inference. This script integrates AHN module weights, tokenizer, and generation configuration into the base model.

```bash
# Merge Qwen2.5-3B-Instruct with GatedDeltaNet AHN weights
BASE_MODEL=Qwen/Qwen2.5-3B-Instruct
AHN_PATH=ByteDance-Seed/AHN-GDN-for-Qwen-2.5-Instruct-3B
MERGED_MODEL_PATH=./merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN

python ./examples/scripts/utils/merge_weights.py \
  --base-model $BASE_MODEL \
  --ahn-path $AHN_PATH \
  --output-path $MERGED_MODEL_PATH

# The merged model includes:
# - All base model parameters
# - AHN module weights properly integrated
# - Tokenizer and generation config
# - Ready for direct inference with AutoModelForCausalLM
```

--------------------------------

### Python Inference with AHN-Augmented Models

Source: https://context7.com/bytedance-seed/ahn/llms.txt

Performs text generation using merged AHN-augmented models. It registers custom AHN model architectures, loads the merged model and tokenizer, and generates text with sliding window attention for long contexts.

```python
# examples/scripts/inference.py
from ahn.transformer.qwen2_ahn import register_customized_qwen2
from ahn.transformer.qwen3_ahn import register_customized_qwen3
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Register AHN-augmented model architectures
register_customized_qwen2()
register_customized_qwen3()

# Load merged model
model_path = "./merged_ckpt/Qwen-2.5-Instruct-3B-AHN-GDN"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="cuda:0"
)
model.eval()

# Generate with long context
prompt = "Write a detailed analysis of the provided 50,000-word document..."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=1024,
        num_beams=1,
        do_sample=False
    )[0]

    context_length = inputs.input_ids.shape[-1]
    prediction = tokenizer.decode(output[context_length:], skip_special_tokens=True)
    print(prediction)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.