Sync latest GeneMamba docs and next-token training updates

scripts/push_to_hub.py

@@ -12,6 +12,7 @@ Requirements:
 import os
 import shutil
 import argparse
+import json
 from pathlib import Path
 from huggingface_hub import HfApi
 
@@ -29,6 +30,47 @@ def collect_local_files(root: Path):
     return files
 
 
+def normalize_config_for_hf(config_path: Path):
+    with config_path.open("r", encoding="utf-8") as f:
+        config = json.load(f)
+
+    if "d_model" in config and "hidden_size" not in config:
+        config["hidden_size"] = config["d_model"]
+    if "mamba_layer" in config and "num_hidden_layers" not in config:
+        config["num_hidden_layers"] = config["mamba_layer"]
+
+    legacy_checkpoint_config = ("d_model" in config) or ("mamba_layer" in config)
+
+    config["model_type"] = "genemamba"
+    config.setdefault("architectures", ["GeneMambaModel"])
+    config.setdefault("max_position_embeddings", 2048)
+    config.setdefault("intermediate_size", 2048)
+    config.setdefault("hidden_dropout_prob", 0.1)
+    config.setdefault("initializer_range", 0.02)
+    if legacy_checkpoint_config and config.get("mamba_mode") == "gate":
+        config["mamba_mode"] = "mean"
+    else:
+        config.setdefault("mamba_mode", "mean")
+    config.setdefault("embedding_pooling", "mean")
+    config.setdefault("num_labels", 2)
+    config.setdefault("pad_token_id", 1)
+    config.setdefault("bos_token_id", 0)
+    config.setdefault("eos_token_id", 2)
+    config.setdefault("use_cache", True)
+    config.setdefault("torch_dtype", "float32")
+    config.setdefault("transformers_version", "4.40.2")
+    config["auto_map"] = {
+        "AutoConfig": "configuration_genemamba.GeneMambaConfig",
+        "AutoModel": "modeling_genemamba.GeneMambaModel",
+        "AutoModelForMaskedLM": "modeling_genemamba.GeneMambaForMaskedLM",
+        "AutoModelForSequenceClassification": "modeling_genemamba.GeneMambaForSequenceClassification",
+    }
+
+    with config_path.open("w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2)
+        f.write("\n")
+
+
 def main():
     project_root = Path(__file__).resolve().parent.parent
 
@@ -118,18 +160,23 @@ def main():
             "modeling_genemamba.py",
             "configuration_genemamba.py",
             "modeling_outputs.py",
+            "README.md",
         ]
         
-        print("  - Copying model definition files...")
+        print("  - Syncing model definition files...")
         for file in model_files:
             src = script_dir / file
             dst = model_path / file
-            if src.exists() and not dst.exists():
-                shutil.copy(src, dst)
-                print(f"    ✓ Copied {file}")
-            elif dst.exists():
-                print(f"    ✓ {file} already exists")
+            if not src.exists():
+                print(f"    ✗ Missing source file: {file}")
+                return 1
+            shutil.copy(src, dst)
+            print(f"    ✓ Synced {file}")
         
+        config_path = model_path / "config.json"
+        normalize_config_for_hf(config_path)
+        print("  - Normalized config.json for custom AutoModel loading")
+
         print("✓ Model files prepared")
         
     except Exception as e: