Direct upload

bucket_memory_model.py

@@ -243,22 +243,26 @@ class BucketMemoryTransformerLayer(nn.Module):
         return x
 
 
-
 # Updated model with HuggingFace compatibility
 class BucketMemoryModel(PreTrainedModel):
-    config_class = BucketMemoryConfig  # Add this line
+    config_class = BucketMemoryConfig
     base_model_prefix = "bucket-memory-model2"
     def __init__(self, config, adapter_kwargs=None):
         super().__init__(config)
         self.d_model = config.d_model
         self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
+
+        # TAPE-style dynamic position encoding
+        self.tape_position_encoder = nn.Sequential(
+            nn.Linear(config.d_model, config.d_model),
+            nn.ReLU(),
+            nn.Linear(config.d_model, config.d_model)
+        )
+
         self.pos_encoding = nn.Parameter(torch.zeros(1, config.max_seq_length, config.d_model))
         self._init_positional_encoding(config.max_seq_length, config.d_model)
 
-        # Use config.num_attention_heads if available, otherwise calculate
-        num_heads = getattr(config, 'num_attention_heads', config.d_model // 64)
-        num_heads = max(1, num_heads)  # Ensure at least 1 head
-
+        num_heads = max(1, getattr(config, 'num_attention_heads', config.d_model // 64))
         self.layers = nn.ModuleList([
             BucketMemoryTransformerLayer(
                 d_model=config.d_model,
@@ -287,10 +291,11 @@ class BucketMemoryModel(PreTrainedModel):
     def forward(self, input_ids, attention_mask=None, labels=None):
         batch_size, seq_len = input_ids.size()
         x = self.token_embedding(input_ids) * np.sqrt(self.d_model)
+        tape_pos = self.tape_position_encoder(x)
+        x = x + tape_pos
         x = x + self.pos_encoding[:, :seq_len]
         x = self.dropout(x)
 
-        # Process through transformer layers
         for layer in self.layers:
             x = layer(x, attention_mask)
 
@@ -303,6 +308,22 @@ class BucketMemoryModel(PreTrainedModel):
             return type('ModelOutput', (), {'loss': loss, 'logits': logits})
         return logits
 
+    def generate(self, input_ids, max_length=50):
+        generated_tokens = input_ids
+        for _ in range(max_length):
+            logits = self.forward(generated_tokens)
+            # Handle both cases: when logits is a tensor or when it's a ModelOutput object
+            if hasattr(logits, 'logits'):
+                next_token_logits = logits.logits[:, -1, :]
+            else:
+                next_token_logits = logits[:, -1, :]
+
+            next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)
+            generated_tokens = torch.cat((generated_tokens, next_token_id), dim=1)
+
+            if next_token_id.item() == self.config.eos_token_id:
+                break
+        return generated_tokens
 AutoConfig.register("bucket-memory-model3", BucketMemoryConfig)
 AutoModel.register(BucketMemoryConfig, BucketMemoryModel)
 BucketMemoryConfig.register_for_auto_class()

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc67a1cb0105fc1bcabf585d7b9d60aed7c787486008f53d73e90a12e5a71762
-size 410490444
+oid sha256:eac09ba151cbcef373c847cabad46af2a3cde38d98403438e949c4b88f8ae061
+size 412592100