Updated py files

__init__.py

@@ -1,7 +1,21 @@
-from .configuration_sapnous import SapnousConfig
-from .modeling_sapnous import SapnousModel
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 typing import TYPE_CHECKING
 from transformers.utils import _LazyModule
+from transformers.models.auto import CONFIG_MAPPING, MODEL_MAPPING, MODEL_FOR_CAUSAL_LM_MAPPING
+from transformers.models.auto import AutoConfig, AutoModel, AutoModelForCausalLM
 
 _import_structure = {
     "configuration_sapnous": ["SAPNOUS_PRETRAINED_CONFIG_ARCHIVE_MAP", "SapnousT1Config"],
@@ -13,4 +27,13 @@ if TYPE_CHECKING:
     from .modeling_sapnous import SapnousT1Model, SapnousT1ForCausalLM
 else:
     import sys
-    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+
+# Register model in auto classes
+CONFIG_MAPPING["sapnous_t1"] = SapnousT1Config
+MODEL_MAPPING["sapnous_t1"] = SapnousT1Model
+MODEL_FOR_CAUSAL_LM_MAPPING["sapnous_t1"] = SapnousT1ForCausalLM
+
+AutoConfig.register("sapnous_t1", SapnousT1Config)
+AutoModel.register(SapnousT1Config, SapnousT1Model)
+AutoModelForCausalLM.register(SapnousT1Config, SapnousT1ForCausalLM)

attention_sapnous.py

@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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.
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
+    """Precompute the frequency tensor for complex rotation."""
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply rotary position embeddings to the input tensor."""
+    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.view(1, *freqs_cis.shape)
+    x_rotated = x_complex * freqs_cis
+    return torch.view_as_real(x_rotated).flatten(-2)
+
+class SapnousAttention(nn.Module):
+    """Multi-head attention with rotary position embeddings and sliding window attention."""
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.sliding_window = config.sliding_window if config.use_sliding_window else None
+
+        if (self.head_dim * self.num_attention_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_attention_heads (got {self.hidden_size} and {self.num_attention_heads})"
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_attention_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.attention_dropout = nn.Dropout(config.attention_dropout)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor:
+        return tensor.view(bsz, seq_len, self.num_attention_heads, self.head_dim).transpose(1, 2)
+
+    def _kv_shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor:
+        return tensor.view(bsz, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = self._shape(query_states, q_len, bsz)
+        key_states = self._kv_shape(key_states, q_len, bsz)
+        value_states = self._kv_shape(value_states, q_len, bsz)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        # Apply rotary position embeddings
+        if position_ids is None:
+            position_ids = torch.arange(kv_seq_len, device=hidden_states.device)
+        cos, sin = freqs_cis[position_ids]
+        query_states, key_states = apply_rotary_emb(query_states, cos), apply_rotary_emb(key_states, sin)
+
+        if past_key_value is not None:
+            # Reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # Repeat k/v heads if n_kv_heads < n_heads
+        key_states = torch.repeat_interleave(key_states, self.num_key_value_groups, dim=1)
+        value_states = torch.repeat_interleave(value_states, self.num_key_value_groups, dim=1)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        # Sliding window attention if configured
+        if self.sliding_window is not None and kv_seq_len > self.sliding_window:
+            # Create sliding window mask
+            window_mask = torch.ones_like(attn_weights, dtype=torch.bool)
+            for i in range(q_len):
+                window_start = max(0, i - self.sliding_window // 2)
+                window_end = min(kv_seq_len, i + self.sliding_window // 2)
+                window_mask[:, :, i, window_start:window_end] = False
+            attn_weights = attn_weights.masked_fill(window_mask, float('-inf'))
+
+        # Causal mask for autoregressive generation
+        if self.config.scoring_func == "softmax":
+            causal_mask = torch.triu(torch.ones((q_len, kv_seq_len), dtype=torch.bool), diagonal=1)
+            causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)
+            attn_weights = attn_weights.masked_fill(causal_mask.to(attn_weights.device), float('-inf'))
+            attn_weights = F.softmax(attn_weights, dim=-1)
+        else:
+            # Alternative scoring functions (e.g., RoPE-only, cosine similarity)
+            attn_weights = F.relu(attn_weights)
+            attn_weights = attn_weights / (attn_weights.sum(dim=-1, keepdim=True) + 1e-6)
+
+        attn_weights = self.attention_dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+class SapnousBlock(nn.Module):
+    """Transformer block with attention, layer norm, and feed-forward network."""
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = SapnousAttention(config)
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        self.mlp = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size, bias=False),
+            nn.SiLU(),
+            nn.Linear(config.intermediate_size, config.hidden_size, bias=False),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # Self Attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            freqs_cis=freqs_cis,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+class SapnousVisionEmbeddings(nn.Module):
+    """Vision embeddings for multimodal support."""
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        
+        # Vision embedding layers
+        self.patch_embed = nn.Conv2d(3, self.hidden_size, kernel_size=16, stride=16)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.hidden_size))
+        self.pos_embed = nn.Parameter(torch.zeros(1, (224 // 16) ** 2 + 1, self.hidden_size))
+        
+        # Layer normalization and dropout
+        self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps)
+        self.dropout = nn.Dropout(config.attention_dropout)
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        B = pixel_values.shape[0]
+        
+        # Create patch embeddings
+        x = self.patch_embed(pixel_values)
+        x = x.flatten(2).transpose(1, 2)  # B, N, C
+        
+        # Add cls token and position embeddings
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        
+        # Apply normalization and dropout
+        x = self.norm(x)
+        x = self.dropout(x)
+        
+        return x

configuration_sapnous.py

@@ -1,11 +1,25 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
-from transformers import AutoConfig  # ✅ Correct Import
+from transformers import AutoConfig  
 
 logger = logging.get_logger(__name__)
 
 SAPNOUS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "Sapnous-AI/Sapnous-6B": "https://huggingface.co/Sapnous-AI/Sapnous-6B/resolve/main/config.json",
+    "Sapnous-AI/Sapnous-VR-6B": "https://huggingface.co/Sapnous-AI/Sapnous-VR-6B/resolve/main/config.json",
 }
 
 class SapnousT1Config(PretrainedConfig):

convert_to_gguf.py

@@ -1,7 +1,24 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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.
 import os
 import torch
+import json
+from pathlib import Path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from ctransformers import AutoModelForCausalLM as GGUFModel
+from models.sapnous import SapnousT1Config
 
 def convert_to_gguf(model_path, output_path):
     # Load the model and tokenizer with vision-language support
@@ -16,29 +33,86 @@ def convert_to_gguf(model_path, output_path):
         trust_remote_code=True
     )
     
-    # Save in GGUF format
+    # Get model configuration
+    config = model.config
+    if not isinstance(config, SapnousT1Config):
+        raise ValueError("Model must be a SapnousT1 model")
+    
+    # Save in intermediate format
     model.save_pretrained(output_path, safe_serialization=True)
     tokenizer.save_pretrained(output_path)
     
-    # Convert to GGUF using ctransformers with SapnousT1 architecture settings
+    # Convert to GGUF using custom SapnousT1 architecture settings
     gguf_model = GGUFModel.from_pretrained(
         output_path,
-        model_type='llama',  # Base architecture type
+        model_type='sapnous_t1',  # Custom architecture type
         gpu_layers=0,  # CPU only for conversion
         config={
-            'context_length': 32768,  # Match model's sliding window size
-            'attention_type': 'multiquery',  # For efficient attention
-            'num_attention_heads': 40,  # Match model's head count
-            'num_key_value_heads': 8,  # Match model's KV head count
-            'hidden_size': 5120,  # Match model's hidden size
-            'intermediate_size': 20480,  # Match model's intermediate size
-            'max_position_embeddings': 128000  # Match model's max positions
+            'context_length': config.sliding_window,
+            'attention_type': 'multihead',  # Custom attention implementation
+            'num_attention_heads': config.num_attention_heads,
+            'num_key_value_heads': config.num_key_value_heads,
+            'hidden_size': config.hidden_size,
+            'intermediate_size': config.intermediate_size,
+            'max_position_embeddings': config.max_position_embeddings,
+            'vocab_size': config.vocab_size,
+            'num_hidden_layers': config.num_hidden_layers,
+            'rms_norm_eps': config.rms_norm_eps,
+            'rope_theta': config.rope_theta,
+            # Vision model parameters
+            'vision_config': {
+                'hidden_size': config.vision_hidden_size,
+                'num_hidden_layers': config.vision_layers,
+                'num_attention_heads': config.vision_heads,
+                'intermediate_size': config.vision_intermediate_size,
+                'patch_size': config.patch_size,
+                'image_size': config.image_size
+            }
         }
     )
     
     print(f"Model converted and saved to {output_path}")
     return gguf_model
 
+def convert_to_hf(gguf_path, output_path):
+    """Convert GGUF model back to Hugging Face format"""
+    # Load GGUF model configuration
+    config_path = Path(gguf_path) / "config.json"
+    with open(config_path, 'r') as f:
+        gguf_config = json.load(f)
+    
+    # Create SapnousT1 configuration
+    config = SapnousT1Config(
+        vocab_size=gguf_config['vocab_size'],
+        hidden_size=gguf_config['hidden_size'],
+        num_hidden_layers=gguf_config['num_hidden_layers'],
+        num_attention_heads=gguf_config['num_attention_heads'],
+        num_key_value_heads=gguf_config['num_key_value_heads'],
+        intermediate_size=gguf_config['intermediate_size'],
+        max_position_embeddings=gguf_config['max_position_embeddings'],
+        rms_norm_eps=gguf_config['rms_norm_eps'],
+        rope_theta=gguf_config['rope_theta'],
+        # Vision configuration
+        vision_hidden_size=gguf_config['vision_config']['hidden_size'],
+        vision_layers=gguf_config['vision_config']['num_hidden_layers'],
+        vision_heads=gguf_config['vision_config']['num_attention_heads'],
+        vision_intermediate_size=gguf_config['vision_config']['intermediate_size'],
+        patch_size=gguf_config['vision_config']['patch_size'],
+        image_size=gguf_config['vision_config']['image_size']
+    )
+    
+    # Load GGUF model
+    gguf_model = GGUFModel.from_pretrained(gguf_path)
+    
+    # Convert weights to HF format
+    model = AutoModelForCausalLM.from_config(config)
+    model.load_state_dict(gguf_model.state_dict())
+    
+    # Save converted model
+    model.save_pretrained(output_path)
+    print(f"Model converted back to Hugging Face format at {output_path}")
+    return model
+
 if __name__ == '__main__':
     model_path = os.path.dirname(os.path.abspath(__file__))
     output_path = os.path.join(model_path, 'gguf_model')

model.py

@@ -1,3 +1,17 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 transformers import PreTrainedModel, AutoConfig
 import torch
 import torch.nn as nn

modeling_sapnous.py

@@ -1,53 +1,271 @@
+import math
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List, Union
 from transformers import PreTrainedModel, AutoModelForCausalLM
-from configuration_sapnous import SapnousT1Config  # Ensure this file is correct
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from .configuration_sapnous import SapnousT1Config
+from .attention_sapnous import SapnousAttention, SapnousBlock, SapnousVisionEmbeddings, precompute_freqs_cis
 
 class SapnousT1PreTrainedModel(PreTrainedModel):
     """Base class for all Sapnous-T1 models."""
     config_class = SapnousT1Config
+    base_model_prefix = "sapnous"
 
     def __init__(self, config: SapnousT1Config):
         super().__init__(config)
         self.config = config
 
     def _init_weights(self, module):
-        """Initialize weights if required."""
+        """Initialize weights using the model's initialization configuration."""
+        std = self.config.initializer_range
         if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, SapnousAttention):
+            module.q_proj.weight.data.normal_(mean=0.0, std=std)
+            module.k_proj.weight.data.normal_(mean=0.0, std=std)
+            module.v_proj.weight.data.normal_(mean=0.0, std=std)
+            module.o_proj.weight.data.normal_(mean=0.0, std=std)
 
 class SapnousT1Model(SapnousT1PreTrainedModel):
-    """Base Transformer Model"""
+    """Base Transformer Model with advanced attention mechanisms and optional vision support."""
     def __init__(self, config: SapnousT1Config):
         super().__init__(config)
+        
         self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.encoder = nn.TransformerEncoder(
-            nn.TransformerEncoderLayer(
-                d_model=config.hidden_size,
-                nhead=config.num_attention_heads
-            ),
-            num_layers=config.num_hidden_layers
+        self.layers = nn.ModuleList([SapnousBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        # Vision support
+        self.vision_embed = SapnousVisionEmbeddings(config) if getattr(config, 'vision_config', None) else None
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+        
+        # Compute and cache RoPE frequencies
+        self.freqs_cis = precompute_freqs_cis(
+            self.config.hidden_size // self.config.num_attention_heads,
+            self.config.max_position_embeddings,
+            self.config.rope_theta,
         )
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
 
-    def forward(self, input_ids):
-        x = self.embeddings(input_ids)
-        x = self.encoder(x)
-        return self.lm_head(x)
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embeddings
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.embeddings = 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[List[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds")
+
+        # Process text input
+        if input_ids is not None:
+            inputs_embeds = self.embeddings(input_ids)
+            batch_size, seq_length = input_ids.shape[:2]
+        else:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+
+        # Process vision input if available
+        if pixel_values is not None and self.vision_embed is not None:
+            vision_embeds = self.vision_embed(pixel_values)
+            inputs_embeds = torch.cat([vision_embeds, inputs_embeds], dim=1)
+            seq_length = inputs_embeds.shape[1]
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0)
+
+        # Prepare attention mask
+        if attention_mask is not None:
+            attention_mask = attention_mask.view(batch_size, -1)
+            attention_mask = attention_mask[:, None, None, :]
+            attention_mask = attention_mask.to(dtype=inputs_embeds.dtype)
+            attention_mask = (1.0 - attention_mask) * torch.finfo(inputs_embeds.dtype).min
+
+        freqs_cis = self.freqs_cis.to(inputs_embeds.device)
+        
+        hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                freqs_cis=freqs_cis,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [
+                hidden_states,
+                next_decoder_cache,
+                all_hidden_states,
+                all_self_attns,
+            ] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
 
 class SapnousT1ForCausalLM(SapnousT1PreTrainedModel):
-    """Sapnous-T1 Model for Causal LM (Text Generation)"""
+    """Sapnous-T1 Model for Causal Language Modeling with vision support."""
+    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
+
     def __init__(self, config: SapnousT1Config):
         super().__init__(config)
         self.model = SapnousT1Model(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-    def forward(self, input_ids):
-        hidden_states = self.model(input_ids)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[List[Tuple[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if position_ids is None:
+            position_ids = (attention_mask.long().cumsum(-1) - 1) if attention_mask is not None else None
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "pixel_values": kwargs.get("pixel_values", None),
+        }
+
+    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[List[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = 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,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""Labels for computing the masked language modeling loss."""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+            pixel_values=pixel_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-        return logits
 
-# ✅ Register the model properly
+        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 tie_weights(self):
+        """Tie the weights between the input embeddings and the output embeddings."""
+        self.lm_head.weight = self.model.embeddings.weight
+
+# Register the model
 AutoModelForCausalLM.register(SapnousT1Config, SapnousT1ForCausalLM)

models/sapnous/__init__.py

@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 typing import TYPE_CHECKING
+from transformers.utils import _LazyModule
+from transformers.models.auto import CONFIG_MAPPING, MODEL_MAPPING, MODEL_FOR_CAUSAL_LM_MAPPING
+from transformers.models.auto import AutoConfig, AutoModel, AutoModelForCausalLM
+
+_import_structure = {
+    "configuration_sapnous": ["SAPNOUS_PRETRAINED_CONFIG_ARCHIVE_MAP", "SapnousT1Config"],
+    "modeling_sapnous": ["SapnousT1Model", "SapnousT1ForCausalLM"],
+    "tokenization_sapnous": ["SapnousT1Tokenizer"],
+}
+
+if TYPE_CHECKING:
+    from .configuration_sapnous import SAPNOUS_PRETRAINED_CONFIG_ARCHIVE_MAP, SapnousT1Config
+    from .modeling_sapnous import SapnousT1Model, SapnousT1ForCausalLM
+    from .tokenization_sapnous import SapnousT1Tokenizer
+else:
+    import sys
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure)
+
+# Import configuration and models
+from .configuration_sapnous import SapnousT1Config
+from .modeling_sapnous import SapnousT1Model, SapnousT1ForCausalLM
+
+# Register model in auto classes
+CONFIG_MAPPING["sapnous_t1"] = SapnousT1Config
+MODEL_MAPPING["sapnous_t1"] = SapnousT1Model
+MODEL_FOR_CAUSAL_LM_MAPPING["sapnous_t1"] = SapnousT1ForCausalLM

models/sapnous/configuration_sapnous.py

@@ -0,0 +1,131 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers import AutoConfig
+
+logger = logging.get_logger(__name__)
+
+SAPNOUS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "Sapnous-AI/Sapnous-6B": "https://huggingface.co/Sapnous-AI/Sapnous-6B/resolve/main/config.json",
+}
+
+class SapnousT1Config(PretrainedConfig):
+    """Configuration class for Sapnous-T1 model with vision-language capabilities.
+    
+    This configuration class handles both text and vision modalities, supporting multimodal
+    tasks like image understanding, video processing, and vision-language reasoning.
+    """
+    
+    model_type = "sapnous_t1"
+
+    def __init__(
+        self,
+        # Text model parameters
+        vocab_size=151936,
+        hidden_size=5120,
+        intermediate_size=20480,
+        num_hidden_layers=36,
+        num_attention_heads=40,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=128000,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=151643,
+        eos_token_id=151645,
+        tie_word_embeddings=True,
+        
+        # Vision model parameters
+        vision_start_token_id=151652,
+        vision_end_token_id=151653,
+        vision_token_id=151654,
+        image_token_id=151655,
+        video_token_id=151656,
+        vision_config=None,
+        patch_size=14,
+        image_size=224,
+        num_channels=3,
+        vision_layers=24,
+        vision_heads=16,
+        vision_hidden_size=1024,
+        vision_intermediate_size=4096,
+        vision_act="gelu",
+        vision_layer_norm_eps=1e-5,
+        vision_dropout=0.0,
+        vision_attention_dropout=0.0,
+        vision_embedding_dropout=0.0,
+        
+        # Cross-attention parameters
+        num_cross_attention_layers=12,
+        cross_attention_heads=16,
+        cross_attention_dropout=0.0,
+        use_cross_attention=True,
+        
+        # Positional encoding and attention parameters
+        rope_theta=1000000.0,
+        sliding_window=32768,
+        use_sliding_window=False,
+        max_window_layers=70,
+        attention_dropout=0.0,
+        rope_scaling=None,
+        scoring_func="softmax",
+        
+        # Training parameters
+        aux_loss_alpha=0.001,
+        seq_aux=True,
+        **kwargs
+    ):
+        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.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        self.vision_token_id = vision_token_id
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_config = vision_config
+        self.rope_theta = rope_theta
+        self.sliding_window = sliding_window
+        self.use_sliding_window = use_sliding_window
+        self.max_window_layers = max_window_layers
+        self.attention_dropout = attention_dropout
+        self.rope_scaling = rope_scaling
+        self.scoring_func = scoring_func
+        self.aux_loss_alpha = aux_loss_alpha
+        self.seq_aux = seq_aux
+
+    model_type = "sapnous_t1"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+AutoConfig.register("sapnous_t1", SapnousT1Config)

models/sapnous/modeling_sapnous.py

@@ -0,0 +1,535 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers import AutoModelForCausalLM, AutoModel
+
+from .configuration_sapnous import SapnousT1Config
+
+logger = logging.get_logger(__name__)
+
+class SapnousT1Attention(nn.Module):
+    def __init__(self, config: SapnousT1Config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.dropout = nn.Dropout(config.attention_dropout)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        if past_key_value is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        # repeat k/v heads if n_kv_heads < n_heads
+        key_states = self._repeat_kv(key_states)
+        value_states = self._repeat_kv(value_states)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = self.dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _repeat_kv(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.num_key_value_heads != self.num_heads:
+            hidden_states = hidden_states.repeat_interleave(self.num_heads // self.num_key_value_heads, dim=1)
+        return hidden_states
+
+class SapnousT1MLP(nn.Module):
+    def __init__(self, config: SapnousT1Config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+class SapnousT1DecoderLayer(nn.Module):
+    def __init__(self, config: SapnousT1Config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = SapnousT1Attention(config=config)
+        self.mlp = SapnousT1MLP(config)
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+class SapnousT1PreTrainedModel(PreTrainedModel):
+    config_class = SapnousT1Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["SapnousT1DecoderLayer"]
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, SapnousT1Model):
+            module.gradient_checkpointing = value
+
+class SapnousT1Model(SapnousT1PreTrainedModel):
+    def __init__(self, config: SapnousT1Config):
+        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([SapnousT1DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        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: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        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
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if attention_mask is not None:
+            if batch_size <= 0:
+                raise ValueError("batch_size has to be defined and > 0")
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+class SapnousT1ForCausalLM(SapnousT1PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SapnousT1Model(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()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = 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,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`)
+            attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*)
+            position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*)
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*)
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*)
+            use_cache (`bool`, *optional*)
+            output_attentions (`bool`, *optional*)
+            output_hidden_states (`bool`, *optional*)
+            return_dict (`bool`, *optional*)
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, SapnousT1ForCausalLM
+
+        >>> model = SapnousT1ForCausalLM.from_pretrained("Sapnous-AI/Sapnous-VR-6B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Sapnous-AI/Sapnous-VR-6B")
+
+        >>> 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? Yes, I am an AI language model capable of engaging in conversation."
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        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,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        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, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+AutoModel.register(SapnousT1Config, SapnousT1Model)
+AutoModelForCausalLM.register(SapnousT1Config, SapnousT1ForCausalLM)

models/sapnous/test_tokenization_sapnous.py

@@ -0,0 +1,52 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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.
+import unittest
+from transformers import AutoTokenizer
+from .configuration_sapnous import SapnousT1Config
+from .tokenization_sapnous import SapnousT1Tokenizer
+
+class TestSapnousTokenizer(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tokenizer = SapnousT1Tokenizer(
+            vocab_file="vocab.json",
+            merges_file="merges.txt"
+        )
+
+    def test_tokenizer_from_pretrained(self):
+        tokenizer = AutoTokenizer.from_pretrained(
+            "Sapnous-AI/Sapnous-VR-6B",
+            trust_remote_code=True
+        )
+        self.assertIsInstance(tokenizer, SapnousT1Tokenizer)
+
+    def test_save_load_pretrained(self):
+        vocab = self.tokenizer.get_vocab()
+        self.assertIsInstance(vocab, dict)
+        self.assertGreater(len(vocab), 0)
+
+    def test_tokenization(self):
+        text = "Hello, world!"
+        tokens = self.tokenizer.tokenize(text)
+        self.assertIsInstance(tokens, list)
+        self.assertGreater(len(tokens), 0)
+
+    def test_special_tokens(self):
+        self.assertIsNotNone(self.tokenizer.unk_token)
+        self.assertIsNotNone(self.tokenizer.bos_token)
+        self.assertIsNotNone(self.tokenizer.eos_token)
+
+if __name__ == '__main__':
+    unittest.main()

models/sapnous/tokenization_sapnous.py

@@ -0,0 +1,91 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 typing import List, Optional, Tuple
+
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+from transformers import AutoTokenizer
+
+logger = logging.get_logger(__name__)
+
+SAPNOUS_PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "Sapnous-AI/Sapnous-VR-6B": "https://huggingface.co/Sapnous-AI/Sapnous-VR-6B/resolve/main/vocab.json",
+    },
+    "merges_file": {
+        "Sapnous-AI/Sapnous-VR-6B": "https://huggingface.co/Sapnous-AI/Sapnous-VR-6B/resolve/main/merges.txt",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "Sapnous-AI/Sapnous-VR-6B": 128000,
+}
+
+class SapnousT1Tokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json", "merges_file": "merges.txt"}
+    pretrained_vocab_files_map = SAPNOUS_PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token="<|endoftext|>",
+        eos_token="<|endoftext|>",
+        pad_token=None,
+        add_prefix_space=False,
+        **kwargs
+    ):
+        super().__init__(
+            errors=errors,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self.merges_file = merges_file
+        self.add_prefix_space = add_prefix_space
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    def _tokenize(self, text: str) -> List[str]:
+        """ Tokenize a string. """
+        raise NotImplementedError("Implement in subclass")
+
+    def _convert_token_to_id(self, token: str) -> int:
+        """ Converts a token to an id using the vocab. """
+        raise NotImplementedError("Implement in subclass")
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """ Converts an index (integer) to a token. """
+        raise NotImplementedError("Implement in subclass")
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str, str]:
+        """ Save the vocabulary and special tokens file to a directory. """
+        raise NotImplementedError("Implement in subclass")
+
+AutoTokenizer.register(SapnousT1Config, SapnousT1Tokenizer)

test_modeling_sapnous.py

@@ -0,0 +1,92 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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.
+import unittest
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from .modeling_sapnous import SapnousT1ForCausalLM
+from .configuration_sapnous import SapnousT1Config
+
+class TestSapnousModel(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.config = SapnousT1Config(
+            vocab_size=32000,
+            hidden_size=768,
+            num_hidden_layers=12,
+            num_attention_heads=12,
+            intermediate_size=3072
+        )
+        cls.model = SapnousT1ForCausalLM(cls.config)
+
+    def test_model_forward(self):
+        input_ids = torch.randint(0, self.config.vocab_size, (1, 10))
+        outputs = self.model(input_ids)
+        
+        self.assertIsNotNone(outputs)
+        self.assertTrue(hasattr(outputs, 'logits'))
+        self.assertEqual(outputs.logits.shape, (1, 10, self.config.vocab_size))
+
+    def test_weight_tying(self):
+        self.model.tie_weights()
+        self.assertTrue(torch.equal(self.model.lm_head.weight, self.model.model.embeddings.weight))
+
+    def test_auto_model_registration(self):
+        model = AutoModelForCausalLM.from_config(self.config)
+        self.assertIsInstance(model, SapnousT1ForCausalLM)
+
+    def test_vision_embeddings(self):
+        # Test vision input processing
+        batch_size = 1
+        pixel_values = torch.randn(batch_size, 3, 224, 224)
+        input_ids = torch.randint(0, self.config.vocab_size, (batch_size, 10))
+        
+        outputs = self.model(input_ids=input_ids, pixel_values=pixel_values)
+        self.assertIsNotNone(outputs)
+        self.assertTrue(hasattr(outputs, 'logits'))
+        
+        # Vision input should increase sequence length
+        expected_seq_length = 10 + (224 // 16) ** 2 + 1  # text_len + num_patches + cls_token
+        self.assertEqual(outputs.logits.shape, (batch_size, expected_seq_length, self.config.vocab_size))
+    
+    def test_attention_mask(self):
+        # Test attention mask handling
+        batch_size = 2
+        seq_length = 15
+        input_ids = torch.randint(0, self.config.vocab_size, (batch_size, seq_length))
+        attention_mask = torch.ones(batch_size, seq_length)
+        attention_mask[:, -5:] = 0  # Mask out last 5 tokens
+        
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        self.assertIsNotNone(outputs)
+        self.assertEqual(outputs.logits.shape, (batch_size, seq_length, self.config.vocab_size))
+    
+    def test_generation_with_vision(self):
+        # Test text generation with vision input
+        pixel_values = torch.randn(1, 3, 224, 224)
+        input_ids = torch.randint(0, self.config.vocab_size, (1, 5))
+        
+        outputs = self.model.generate(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            max_length=20,
+            num_beams=1
+        )
+        
+        self.assertIsInstance(outputs, torch.Tensor)
+        self.assertEqual(outputs.dim(), 2)
+        self.assertTrue(outputs.size(1) <= 20)
+
+if __name__ == '__main__':
+    unittest.main()

test_tokenization_sapnous.py

@@ -0,0 +1,157 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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.
+import unittest
+import torch
+from pathlib import Path
+from transformers import AutoTokenizer
+from .tokenization_sapnous import SapnousTokenizer
+
+class TestSapnousTokenizer(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # Create temporary vocab and merges files for testing
+        cls.temp_dir = Path('test_tokenizer_files')
+        cls.temp_dir.mkdir(exist_ok=True)
+        
+        # Create a simple test vocabulary
+        cls.vocab_file = cls.temp_dir / 'vocab.json'
+        cls.vocab = {
+            '<|endoftext|>': 0,
+            '<|startoftext|>': 1,
+            '<|pad|>': 2,
+            '<|vision_start|>': 3,
+            '<|vision_end|>': 4,
+            '<|image|>': 5,
+            '<|video|>': 6,
+            'hello': 7,
+            'world': 8,
+            'test': 9,
+        }
+        with cls.vocab_file.open('w', encoding='utf-8') as f:
+            import json
+            json.dump(cls.vocab, f)
+        
+        # Create test merges file
+        cls.merges_file = cls.temp_dir / 'merges.txt'
+        merges_content = "#version: 0.2\nh e\ne l\nl l\no w\nw o\no r\nr l\nl d"
+        cls.merges_file.write_text(merges_content)
+        
+        # Initialize tokenizer
+        cls.tokenizer = SapnousTokenizer(
+            str(cls.vocab_file),
+            str(cls.merges_file),
+        )
+    
+    @classmethod
+    def tearDownClass(cls):
+        # Clean up temporary files
+        import shutil
+        shutil.rmtree(cls.temp_dir)
+    
+    def test_tokenizer_initialization(self):
+        self.assertEqual(self.tokenizer.vocab_size, len(self.vocab))
+        self.assertEqual(self.tokenizer.get_vocab(), self.vocab)
+        
+        # Test special tokens
+        self.assertEqual(self.tokenizer.unk_token, '<|endoftext|>')
+        self.assertEqual(self.tokenizer.bos_token, '<|startoftext|>')
+        self.assertEqual(self.tokenizer.eos_token, '<|endoftext|>')
+        self.assertEqual(self.tokenizer.pad_token, '<|pad|>')
+    
+    def test_tokenization(self):
+        text = "hello world test"
+        tokens = self.tokenizer.tokenize(text)
+        self.assertIsInstance(tokens, list)
+        self.assertTrue(all(isinstance(token, str) for token in tokens))
+        
+        # Test encoding
+        input_ids = self.tokenizer.encode(text, add_special_tokens=False)
+        self.assertIsInstance(input_ids, list)
+        self.assertEqual(len(input_ids), 3)  # 'hello', 'world', 'test'
+        
+        # Test decoding
+        decoded_text = self.tokenizer.decode(input_ids)
+        self.assertEqual(decoded_text.strip(), text)
+    
+    def test_special_tokens_handling(self):
+        text = "hello world"
+        # Test with special tokens
+        tokens_with_special = self.tokenizer.encode(text, add_special_tokens=True)
+        self.assertTrue(tokens_with_special[0] == self.tokenizer.bos_token_id)
+        self.assertTrue(tokens_with_special[-1] == self.tokenizer.eos_token_id)
+        
+        # Test without special tokens
+        tokens_without_special = self.tokenizer.encode(text, add_special_tokens=False)
+        self.assertNotEqual(tokens_without_special[0], self.tokenizer.bos_token_id)
+        self.assertNotEqual(tokens_without_special[-1], self.tokenizer.eos_token_id)
+    
+    def test_vision_tokens(self):
+        # Test vision-specific token methods
+        text = "This is an image description"
+        vision_text = self.tokenizer.prepare_for_vision(text)
+        self.assertTrue(vision_text.startswith('<|vision_start|>'))
+        self.assertTrue(vision_text.endswith('<|vision_end|>'))
+        
+        image_text = self.tokenizer.prepare_for_image(text)
+        self.assertTrue(image_text.startswith('<|image|>'))
+        
+        video_text = self.tokenizer.prepare_for_video(text)
+        self.assertTrue(video_text.startswith('<|video|>'))
+    
+    def test_batch_encoding(self):
+        texts = ["hello world", "test hello"]
+        batch_encoding = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+        
+        self.assertIsInstance(batch_encoding["input_ids"], torch.Tensor)
+        self.assertIsInstance(batch_encoding["attention_mask"], torch.Tensor)
+        self.assertEqual(batch_encoding["input_ids"].shape[0], len(texts))
+        self.assertEqual(batch_encoding["attention_mask"].shape[0], len(texts))
+    
+    def test_save_and_load(self):
+        # Test saving vocabulary
+        save_dir = Path('test_save_tokenizer')
+        save_dir.mkdir(exist_ok=True)
+        
+        try:
+            vocab_files = self.tokenizer.save_vocabulary(str(save_dir))
+            self.assertTrue(all(Path(f).exists() for f in vocab_files))
+            
+            # Test loading saved vocabulary
+            loaded_tokenizer = SapnousTokenizer(*vocab_files)
+            self.assertEqual(loaded_tokenizer.get_vocab(), self.tokenizer.get_vocab())
+            
+            # Test encoding/decoding with loaded tokenizer
+            text = "hello world test"
+            original_encoding = self.tokenizer.encode(text)
+            loaded_encoding = loaded_tokenizer.encode(text)
+            self.assertEqual(original_encoding, loaded_encoding)
+        finally:
+            # Clean up
+            import shutil
+            shutil.rmtree(save_dir)
+    
+    def test_auto_tokenizer_registration(self):
+        # Test if the tokenizer can be loaded using AutoTokenizer
+        config = {
+            "model_type": "sapnous",
+            "vocab_file": str(self.vocab_file),
+            "merges_file": str(self.merges_file)
+        }
+        
+        tokenizer = AutoTokenizer.from_pretrained(str(self.temp_dir), **config)
+        self.assertIsInstance(tokenizer, SapnousTokenizer)
+
+if __name__ == '__main__':
+    unittest.main()

tokenization_sapnous.py

@@ -0,0 +1,197 @@
+# coding=utf-8
+# Copyright 2025-present, the HuggingFace Inc. Team and AIRAS 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 typing import List, Optional, Tuple, Union
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers import AutoTokenizer
+import json
+import regex as re
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+BYTES_TO_UNICODE_REGEX = re.compile(r"'([^']+)':\s*([0-9]+)")
+
+def bytes_to_unicode():
+    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+def get_pairs(word):
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+class SapnousTokenizer(PreTrainedTokenizer):
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        merges_file: Optional[str] = None,
+        unk_token: str = "<|endoftext|>",
+        bos_token: str = "<|startoftext|>",
+        eos_token: str = "<|endoftext|>",
+        pad_token: str = "<|pad|>",
+        vision_start_token: str = "<|vision_start|>",
+        vision_end_token: str = "<|vision_end|>",
+        image_token: str = "<|image|>",
+        video_token: str = "<|video|>",
+        add_prefix_space: bool = False,
+        **kwargs
+    ):
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self.merges_file = merges_file
+        self.add_prefix_space = add_prefix_space
+
+        self.special_tokens = {
+            "unk_token": unk_token,
+            "bos_token": bos_token,
+            "eos_token": eos_token,
+            "pad_token": pad_token,
+            "vision_start_token": vision_start_token,
+            "vision_end_token": vision_end_token,
+            "image_token": image_token,
+            "video_token": video_token,
+        }
+
+        with Path(vocab_file).open(encoding="utf-8") as f:
+            self.encoder = json.load(f)
+        self.decoder = {v: k for k, v in self.encoder.items()}
+
+        if merges_file:
+            with Path(merges_file).open(encoding="utf-8") as f:
+                bpe_merges = f.read().strip().split('\n')[1:]
+            bpe_merges = [tuple(merge.split()) for merge in bpe_merges]
+            self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        else:
+            self.bpe_ranks = {}
+
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?\d+| ?[^\s\w\d]+|\s+(?!\S)|\s+""")
+
+    def bpe(self, token: str) -> str:
+        if token in self.special_tokens.values():
+            return token
+
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    if word[j + 1] == second:
+                        new_word.append(first + second)
+                        i = j + 2
+                    else:
+                        new_word.append(word[j])
+                        i = j + 1
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+            word = tuple(new_word)
+            if len(word) == 1:
+                break
+            pairs = get_pairs(word)
+        return ' '.join(word)
+
+    def _tokenize(self, text: str) -> List[str]:
+        if self.add_prefix_space:
+            text = ' ' + text
+
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[ord(b)] for b in token)
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.decoder.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        text = ''.join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors='replace')
+        return text
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str, str]:
+        if not filename_prefix:
+            filename_prefix = ""
+
+        vocab_file = Path(save_directory) / f"{filename_prefix}vocab.json"
+        merge_file = Path(save_directory) / f"{filename_prefix}merges.txt"
+
+        with vocab_file.open('w', encoding='utf-8') as f:
+            json.dump(self.encoder, f, ensure_ascii=False)
+
+        if self.merges_file:
+            with merge_file.open('w', encoding='utf-8') as f:
+                for merge in self.bpe_ranks:
+                    f.write(f"{merge[0]} {merge[1]}\n")
+            return str(vocab_file), str(merge_file)
+
+        return str(vocab_file)
+
+    def prepare_for_vision(self, text: str) -> str:
+        """Prepare text for vision tasks by adding special tokens."""
+        return f"{self.vision_start_token}{text}{self.vision_end_token}"
+
+    def prepare_for_image(self, text: str) -> str:
+        """Prepare text for image tasks."""
+        return f"{self.image_token}{text}"
+
+    def prepare_for_video(self, text: str) -> str:
+        """Prepare text for video tasks."""
+        return f"{self.video_token}{text}"
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self.encoder.copy()
+
+# Register the tokenizer
+AutoTokenizer.register(SapnousTokenizer, "sapnous")