Upload model.py with huggingface_hub

model.py

@@ -9,6 +9,9 @@ from typing import Union, List, Optional
 import bulletchess
 import numpy as np
 
+
+from transformers import PretrainedConfig
+
 class Gating(nn.Module):
     def __init__(self, features_shape, additive=True, init_value=None):
         super(Gating, self).__init__()
@@ -237,10 +240,71 @@ class ValueHead(nn.Module):
         x = self.dense2(x)
         return x
 
+class BT4Config(PretrainedConfig):
+    """Configuration class for BT4 model."""
+    model_type = "bt4"
+    
+    def __init__(
+        self,
+        embedding_size=1024,
+        embedding_dense_sz=512,
+        encoder_layers=15,
+        encoder_d_model=1024,
+        encoder_heads=32,
+        encoder_dff=1536,
+        dropout_rate=0.0,
+        pol_embedding_size=1024,
+        policy_d_model=1024,
+        val_embedding_size=128,
+        use_smolgen=True,
+        smol_hidden_channels=32,
+        smol_hidden_sz=256,
+        smol_gen_sz=256,
+        smol_activation="swish",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.embedding_size = embedding_size
+        self.embedding_dense_sz = embedding_dense_sz
+        self.encoder_layers = encoder_layers
+        self.encoder_d_model = encoder_d_model
+        self.encoder_heads = encoder_heads
+        self.encoder_dff = encoder_dff
+        self.dropout_rate = dropout_rate
+        self.pol_embedding_size = pol_embedding_size
+        self.policy_d_model = policy_d_model
+        self.val_embedding_size = val_embedding_size
+        self.use_smolgen = use_smolgen
+        self.smol_hidden_channels = smol_hidden_channels
+        self.smol_hidden_sz = smol_hidden_sz
+        self.smol_gen_sz = smol_gen_sz
+        self.smol_activation = smol_activation
+
 class BT4(nn.Module):
-    def __init__(self, embedding_size=1024, embedding_dense_sz=512, encoder_layers=15, encoder_d_model=1024, encoder_heads=32, encoder_dff=1536, dropout_rate=0.0, pol_embedding_size=1024, policy_d_model=1024, val_embedding_size=128, default_activation=Mish(),
+    def __init__(self, config=None, embedding_size=1024, embedding_dense_sz=512, encoder_layers=15, encoder_d_model=1024, encoder_heads=32, encoder_dff=1536, dropout_rate=0.0, pol_embedding_size=1024, policy_d_model=1024, val_embedding_size=128, default_activation=Mish(),
                  use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
         super(BT4, self).__init__()
+        
+        # Store config if provided
+        self.config = config
+        
+        # If config is provided, use it to override parameters
+        if config is not None:
+            embedding_size = getattr(config, 'embedding_size', embedding_size)
+            embedding_dense_sz = getattr(config, 'embedding_dense_sz', embedding_dense_sz)
+            encoder_layers = getattr(config, 'encoder_layers', encoder_layers)
+            encoder_d_model = getattr(config, 'encoder_d_model', encoder_d_model)
+            encoder_heads = getattr(config, 'encoder_heads', encoder_heads)
+            encoder_dff = getattr(config, 'encoder_dff', encoder_dff)
+            dropout_rate = getattr(config, 'dropout_rate', dropout_rate)
+            pol_embedding_size = getattr(config, 'pol_embedding_size', pol_embedding_size)
+            policy_d_model = getattr(config, 'policy_d_model', policy_d_model)
+            val_embedding_size = getattr(config, 'val_embedding_size', val_embedding_size)
+            use_smolgen = getattr(config, 'use_smolgen', use_smolgen)
+            smol_hidden_channels = getattr(config, 'smol_hidden_channels', smol_hidden_channels)
+            smol_hidden_sz = getattr(config, 'smol_hidden_sz', smol_hidden_sz)
+            smol_gen_sz = getattr(config, 'smol_gen_sz', smol_gen_sz)
+            smol_activation = getattr(config, 'smol_activation', smol_activation)
         self.embedding_dense_sz = embedding_dense_sz
         # DeepNorm alpha used in embedding residual; default uses provided encoder_layers
         self.deepnorm_alpha = (2. * encoder_layers) ** -0.25
@@ -271,6 +335,53 @@ class BT4(nn.Module):
         self.activation = default_activation
         
         self.apply(self._init_weights)
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Load model from pretrained checkpoint (required by transformers)."""
+        from transformers import AutoConfig
+        
+        # Load config
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        
+        # Create model with config
+        model = cls(config=config)
+        
+        # Load weights if available
+        try:
+            from safetensors.torch import load_file
+            import os
+            
+            # Try safetensors first
+            safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors")
+            if os.path.exists(safetensors_path):
+                state_dict = load_file(safetensors_path)
+                model.load_state_dict(state_dict)
+            else:
+                # Fall back to pytorch format
+                pt_path = os.path.join(pretrained_model_name_or_path, "model.pt")
+                if os.path.exists(pt_path):
+                    checkpoint = torch.load(pt_path, map_location="cpu")
+                    if isinstance(checkpoint, dict):
+                        if "state_dict" in checkpoint:
+                            model.load_state_dict(checkpoint["state_dict"])
+                        elif "model" in checkpoint:
+                            model.load_state_dict(checkpoint["model"])
+                        else:
+                            model.load_state_dict(checkpoint)
+                    else:
+                        model.load_state_dict(checkpoint)
+        except Exception as e:
+            # If weights don't exist or fail to load, return model without weights
+            pass
+        
+        return model
+    
+    @classmethod
+    def register_for_auto_class(cls, auto_class):
+        """Register this class for auto class loading (required by transformers)."""
+        # This is a no-op for custom models with trust_remote_code=True
+        pass
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):