Upload model.py with huggingface_hub

model.py

@@ -10,7 +10,7 @@ import bulletchess
 import numpy as np
 
 
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, PreTrainedModel
 
 class Gating(nn.Module):
     def __init__(self, features_shape, additive=True, init_value=None):
@@ -280,31 +280,46 @@ class BT4Config(PretrainedConfig):
         self.smol_gen_sz = smol_gen_sz
         self.smol_activation = smol_activation
 
-class BT4(nn.Module):
+class BT4(PreTrainedModel):
     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)
+        # Initialize PreTrainedModel with config
+        if config is None:
+            config = BT4Config(
+                embedding_size=embedding_size,
+                embedding_dense_sz=embedding_dense_sz,
+                encoder_layers=encoder_layers,
+                encoder_d_model=encoder_d_model,
+                encoder_heads=encoder_heads,
+                encoder_dff=encoder_dff,
+                dropout_rate=dropout_rate,
+                pol_embedding_size=pol_embedding_size,
+                policy_d_model=policy_d_model,
+                val_embedding_size=val_embedding_size,
+                use_smolgen=use_smolgen,
+                smol_hidden_channels=smol_hidden_channels,
+                smol_hidden_sz=smol_hidden_sz,
+                smol_gen_sz=smol_gen_sz,
+                smol_activation=smol_activation,
+            )
+        super(BT4, self).__init__(config)
+        
+        # Use config values (config is now guaranteed to exist)
+        embedding_size = config.embedding_size
+        embedding_dense_sz = config.embedding_dense_sz
+        encoder_layers = config.encoder_layers
+        encoder_d_model = config.encoder_d_model
+        encoder_heads = config.encoder_heads
+        encoder_dff = config.encoder_dff
+        dropout_rate = config.dropout_rate
+        pol_embedding_size = config.pol_embedding_size
+        policy_d_model = config.policy_d_model
+        val_embedding_size = config.val_embedding_size
+        use_smolgen = config.use_smolgen
+        smol_hidden_channels = config.smol_hidden_channels
+        smol_hidden_sz = config.smol_hidden_sz
+        smol_gen_sz = config.smol_gen_sz
+        smol_activation = config.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