Initial ZeusMM: tiny-pretrained weights, tokenizer, config, custom code

config.json

@@ -1,4 +1,5 @@
 {
+  "_name_or_path": "ZeusMM",
   "architectures": [
     "ZeusForCausalLM"
   ],

special_tokens_map.json

@@ -1,78 +1,42 @@
 {
   "additional_special_tokens": [
-    {
-      "content": "<|system|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<|user|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<|assistant|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<image>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "</image>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<audio>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "</audio>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<kb>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "</kb>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<|end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    }
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<image>",
+    "</image>",
+    "<audio>",
+    "</audio>",
+    "<kb>",
+    "</kb>",
+    "<|end|>"
   ],
-  "bos_token": "<|endoftext|>",
-  "eos_token": "<|endoftext|>",
-  "pad_token": "<|endoftext|>",
-  "unk_token": "<|endoftext|>"
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
 }

zeus_mm.py

@@ -514,15 +514,20 @@ class ZeusForCausalLM(PreTrainedModel):
     def forward(
         self,
         input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,   # [B,T_new]
+        attention_mask: Optional[torch.LongTensor] = None,
         labels: Optional[torch.LongTensor] = None,
-        role_ids: Optional[torch.LongTensor] = None,         # [B,T_total] (0/1/2)
+        role_ids: Optional[torch.LongTensor] = None,
         past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
         use_cache: Optional[bool] = None,
 
+        # HF Generation adds these — accept & ignore
+        return_dict: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+
         # ---- Raw inputs for backends OR precomputed memories ----
         # Vision
-        pixel_values: Optional[torch.FloatTensor] = None,         # [B,3,H,W] normalized
+        pixel_values: Optional[torch.FloatTensor] = None,         # [B,3,H,W]
         image_memory: Optional[torch.FloatTensor] = None,         # [B,Li,D]
         # Audio
         input_values: Optional[torch.FloatTensor] = None,         # [B,T_audio]
@@ -532,11 +537,15 @@ class ZeusForCausalLM(PreTrainedModel):
         retr_input_ids: Optional[torch.LongTensor] = None,        # [B,Nr]
         retr_attention_mask: Optional[torch.LongTensor] = None,   # [B,Nr]
         retr_memory: Optional[torch.FloatTensor] = None,          # [B,Lr,D]
-
-        # Pre-assembled (advanced): concat memory & mask + pooled summary
+        # Pre-assembled
         memory_mask: Optional[torch.LongTensor] = None,           # [B,Lm]
         media_summary: Optional[torch.FloatTensor] = None,        # [B,D]
+
+        # future-proof
+        **unused,
     ):
+
+
         B, T = input_ids.shape
         x = self.embed_tokens(input_ids)
         x = self.drop(x)