Upload folder using huggingface_hub

README.md

@@ -0,0 +1,6 @@
+---
+{}
+---
+
+    # EnhancAR
+    

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_enhancar import EnhancARConfig
+from .modeling_enhancar import EnhancARModel
+
+__all__ = ["EnhancARConfig", "EnhancARModel"]

config.json

@@ -0,0 +1,47 @@
+{
+  "model_config": {
+    "hidden_size": 256,
+    "intermediate_size": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "num_key_value_heads": 8,
+    "use_mamba_kernels": true,
+    "mamba_d_state": 16,
+    "mamba_d_conv": 4,
+    "mamba_expand": 2,
+    "mamba_dt_rank": "auto",
+    "mamba_conv_bias": true,
+    "mamba_proj_bias": false,
+    "output_router_logits": true,
+    "use_cache": false,
+    "_attn_implementation": "flash_attention_2",
+    "vocab_size": 16,
+    "pad_token_id": 6,
+    "bos_token_id": 9,
+    "eos_token_id": 7
+  },
+  "alphabet": [
+    "G",
+    "A",
+    "T",
+    "C",
+    "N",
+    "-",
+    "!",
+    "*",
+    "/",
+    "@",
+    "[",
+    "]",
+    "{",
+    "}"
+  ],
+  "model_type": "enhancar",
+  "architectures": [
+    "EnhancARModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_enhancar.EnhancARConfig",
+    "AutoModel": "modeling_enhancar.EnhancARModel"
+  }
+}

configuration_enhancar.py

@@ -0,0 +1,6 @@
+from transformers import PretrainedConfig
+class EnhancARConfig(PretrainedConfig):
+    model_type = "enhancar"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7095d07d96c0ff232a2b65755274ecddf36269752e696937b3ac3911c7e1db4a
+size 681263120

modeling_enhancar.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel, AutoModel
+from configuration_enhancar import EnhancARConfig
+
+class EnhancARDNAModel(nn.Module):
+    def __init__(self, jamba_base, d_model, vocab_size):
+        super().__init__()
+        self.embedder = jamba_base.model 
+        self.lm_head = nn.Linear(d_model, vocab_size)
+        self.seq_embedding = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, input_ids, labels=None):
+        inputs_embeds = self.seq_embedding(input_ids)
+        outputs = self.embedder(inputs_embeds=inputs_embeds)
+        hidden_states = outputs["last_hidden_state"]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)), 
+                shift_labels.view(-1)
+            )
+
+        return {
+            "loss": loss,
+            "logits": logits,
+            "representation": hidden_states
+        }
+
+# ── The Hugging Face PreTrainedModel Wrapper ──────────────────────────────
+
+class EnhancARModel(PreTrainedModel):
+    config_class = EnhancARConfig
+    base_model_prefix = "model"
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: EnhancARConfig):
+        super().__init__(config)
+        
+        hf_cfg = AutoConfig.from_pretrained("ai21labs/Jamba-v0.1", trust_remote_code=True)
+        merged = {**hf_cfg.to_dict(), **config.to_dict()}
+        hf_cfg = type(hf_cfg).from_dict(merged)
+        
+        base_jamba = AutoModelForCausalLM.from_config(hf_cfg, trust_remote_code=True)
+        
+        self.model = EnhancARDNAModel(
+            base_jamba, 
+            d_model=config.hidden_size, 
+            vocab_size=config.vocab_size
+        )
+        self.post_init()
+
+    def forward(self, input_ids, labels=None, **kwargs):
+        return self.model(input_ids=input_ids, labels=labels)

tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "alphabet": [
+    "G",
+    "A",
+    "T",
+    "C",
+    "N",
+    "-",
+    "!",
+    "*",
+    "/",
+    "@",
+    "[",
+    "]",
+    "{",
+    "}"
+  ]
+}