Upload magic-bert-50m-mlm model files

README.md

@@ -0,0 +1,261 @@
+---
+language:
+  - en
+license: apache-2.0
+library_name: transformers
+tags:
+  - binary-analysis
+  - file-type-detection
+  - byte-level
+  - fill-mask
+  - mlm
+  - magic-bytes
+  - security
+pipeline_tag: fill-mask
+model-index:
+  - name: magic-bert-50m-mlm
+    results:
+      - task:
+          type: fill-mask
+          name: Masked Language Modeling
+        metrics:
+          - name: Perplexity
+            type: perplexity
+            value: 1.05
+          - name: Fill-mask Top-1 Accuracy
+            type: accuracy
+            value: 58.9
+          - name: Fill-mask Top-5 Accuracy
+            type: accuracy
+            value: 73.5
+          - name: Probing Classification Accuracy
+            type: accuracy
+            value: 87.0
+---
+
+# Magic-BERT 50M MLM
+
+A BERT-style transformer model trained for binary file understanding using masked language modeling (MLM). This model learns byte-level patterns in binary files, including magic bytes, headers, and structural patterns across 106 file types.
+
+## Why Not Just Use libmagic?
+
+For intact files starting at byte 0, libmagic works well. But libmagic matches *signatures at fixed offsets*. Magic-BERT learns *structural patterns* throughout the file, enabling use cases where you don't have clean file boundaries:
+
+- **Network streams**: Classifying packet payloads mid-connection, before headers arrive
+- **Disk forensics**: Identifying file types during carving, when scanning raw disk images without filesystem metadata
+- **Fragment analysis**: Working with partial files, slack space, or corrupted data
+- **Adversarial contexts**: Detecting file types when magic bytes are stripped, spoofed, or deliberately misleading
+
+## Model Description
+
+Magic-BERT uses a custom BERT architecture with absolute position embeddings, trained on binary file data using a byte-level BPE tokenizer. The MLM objective teaches the model to predict masked bytes given surrounding context, which implicitly learns file format structure.
+
+| Property | Value |
+|----------|-------|
+| Parameters | 59M |
+| Hidden Size | 512 |
+| Layers | 8 |
+| Attention Heads | 8 |
+| Max Sequence Length | 512 tokens |
+| Vocabulary Size | 32,768 (byte-level BPE) |
+| Position Encoding | Absolute (learned embeddings) |
+
+### Tokenizer
+
+The tokenizer uses the Binary BPE methodology introduced in [Bommarito (2025)](https://arxiv.org/abs/2511.17573). The original Binary BPE tokenizers (available at [mjbommar/binary-tokenizer-001-64k](https://huggingface.co/mjbommar/binary-tokenizer-001-64k)) were trained exclusively on executable binaries (ELF, PE, Mach-O). This tokenizer uses the same BPE training approach but was trained on a diverse corpus spanning 106 file types including documents, images, audio/video, archives, and source code.
+
+## Intended Uses
+
+**Primary use cases:**
+- Fill-mask: Predicting missing bytes in binary files
+- Magic byte and file signature recognition
+- Feature extraction for downstream classification
+- Research on binary file structure
+
+**Example tasks:**
+- Completing partial file headers
+- Identifying file type from structure
+- Anomaly detection in binary data
+
+## Detailed Use Cases
+
+### Network Traffic Analysis
+When inspecting packet payloads, you often see file data mid-stream—TCP reassembly may give you bytes 1500-3000 of a PDF before you ever see byte 0. Traditional signature matching fails here. Magic-BERT's structural understanding can identify file types from interior content.
+
+### Disk Forensics & File Carving
+During disk image analysis, you scan raw bytes looking for file boundaries. Tools like Scalpel rely on header/footer signatures, but many files lack clear footers. Magic-BERT can score byte ranges for file type probability, helping identify carved fragments or validate carving results.
+
+### Incident Response
+Malware often strips or modifies magic bytes to evade detection. Polyglot files (valid as multiple types) exploit signature-based tools. Learning structural patterns provides a second opinion that doesn't rely solely on the first few bytes.
+
+### Embedded Content Detection
+Files within files (email attachments, archive contents, OLE streams) may appear at arbitrary offsets. Embeddings from Magic-BERT enable similarity search: "find all chunks that look structurally like JPEG data" regardless of where they appear.
+
+## Training
+
+### Data
+Trained on a diverse corpus of binary files spanning 106 MIME types, including:
+- Documents (PDF, Office formats, OpenDocument)
+- Images (PNG, JPEG, GIF, WebP, TIFF)
+- Audio/Video (MP3, MP4, WebM, FLAC)
+- Archives (ZIP, GZIP, 7z, TAR)
+- Executables (ELF, PE, Mach-O)
+- And 90+ additional formats
+
+### Procedure
+
+| Phase | Steps | Learning Rate | Batch Size | Objective |
+|-------|-------|---------------|------------|-----------|
+| MLM Pre-training | 100,000 | 1e-4 | 240 | Masked LM (15% masking) |
+
+**Data augmentation:** 50% of samples use random byte offset to reduce position bias.
+
+## Evaluation Results
+
+### Perplexity by Region
+| Region | Perplexity |
+|--------|------------|
+| Magic Bytes (0-9) | 1.07 |
+| Header (10-49) | 1.06 |
+| Body (50+) | 1.05 |
+| **Overall** | **1.05** |
+
+### Fill-Mask Accuracy
+| Metric | Value |
+|--------|-------|
+| Top-1 Accuracy | 58.9% |
+| Top-5 Accuracy | 73.5% |
+| Mean Reciprocal Rank | 0.67 |
+
+### Representation Quality
+| Metric | Value |
+|--------|-------|
+| Linear Probe Accuracy | 87.0% |
+| Silhouette Score | 0.39 |
+| Separation Ratio | 2.78 |
+
+## Architecture: Absolute vs Rotary Position Embeddings
+
+This model uses **absolute position embeddings**, where each position (0-511) has a learned embedding vector added to the token embedding. This is the original BERT approach.
+
+An alternative is **Rotary Position Embeddings (RoPE)**, used by the RoFormer variant. RoPE encodes relative position through rotation matrices applied to query and key vectors in attention, rather than learning absolute position vectors.
+
+**Key finding from our experiments:** Both approaches show similar position bias (~47-48% accuracy drop at offset 1000). Position bias is primarily a data distribution issue (files naturally start at offset 0) rather than an architecture limitation.
+
+| Aspect | Magic-BERT (this) | RoFormer |
+|--------|-------------------|----------|
+| Position Encoding | Absolute (learned) | RoPE (rotary) |
+| Parameters | 59M | 42.3M |
+| Perplexity | **1.05** | 1.13 |
+| Fill-mask Top-1 | 58.9% | **61.8%** |
+| Probing Accuracy | **87.0%** | 85.0% |
+
+Magic-BERT achieves slightly better perplexity and probing accuracy, while RoFormer achieves better fill-mask accuracy with fewer parameters.
+
+## MLM vs Classification: Two-Phase Training
+
+This is the **Phase 1 (MLM)** model. The training pipeline has two phases:
+
+| Phase | Model | Task | Purpose |
+|-------|-------|------|---------|
+| **Phase 1** | **This model** | Masked Language Modeling | Learn byte-level patterns and file structure |
+| Phase 2 | magic-bert-50m-classification | Contrastive Learning | Optimize embeddings for file type discrimination |
+
+**When to use each:**
+- Use **this model (MLM)** for: fill-mask tasks, research, or as a base for custom fine-tuning
+- Use **classification model** for: file type detection, similarity search, production classification
+
+## How to Use
+
+```python
+from transformers import AutoTokenizer
+from safetensors.torch import load_file
+import torch
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("path/to/magic-bert-50m-mlm")
+
+# For custom MagicBERT architecture, load directly
+from modeling_magic_bert import MagicBERTForMaskedLM
+from configuration_magic_bert import MagicBERTConfig
+
+config = MagicBERTConfig.from_pretrained("path/to/magic-bert-50m-mlm")
+model = MagicBERTForMaskedLM(config)
+state_dict = load_file("path/to/magic-bert-50m-mlm/model.safetensors")
+model.load_state_dict(state_dict)
+
+# Fill-mask example
+with open("example.pdf", "rb") as f:
+    data = f.read(512)
+
+# Decode bytes to string using latin-1 (preserves all byte values 0-255)
+text = data.decode("latin-1")
+
+# Tokenize and mask
+inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+mask_pos = 0  # Mask first token
+inputs["input_ids"][0, mask_pos] = tokenizer.mask_token_id
+
+# Predict
+with torch.no_grad():
+    outputs = model(**inputs)
+    predictions = outputs.logits[0, mask_pos].topk(5)
+
+print("Top-5 predictions:", tokenizer.convert_ids_to_tokens(predictions.indices))
+```
+
+### Getting Embeddings
+
+```python
+# Get CLS embeddings for downstream tasks
+with torch.no_grad():
+    embeddings = model.get_embeddings(inputs["input_ids"], inputs["attention_mask"])
+    # embeddings shape: [batch_size, 512]
+```
+
+## Limitations
+
+1. **Position bias:** The model performs best when file content starts at position 0. Accuracy drops ~48% when content starts at offset 1000. This reflects training data distribution, not architectural limitations.
+
+2. **Sequence length:** Limited to 512 tokens. Longer files require truncation or chunking.
+
+3. **Text files:** Lower performance on high-entropy or highly variable content (e.g., encrypted data, random bytes).
+
+4. **Domain specificity:** Trained on common file formats; may not generalize to rare or proprietary formats.
+
+## Model Selection Guide
+
+| Use Case | Recommended Model | Reason |
+|----------|-------------------|--------|
+| Fill-mask / byte prediction | **This model** | Best perplexity (1.05) |
+| Research baseline | **This model** | Established BERT architecture |
+| Classification + fill-mask | magic-bert-50m-classification | Retains 41.8% fill-mask capability |
+| **Production classification** | **magic-bert-50m-roformer-classification** | Highest accuracy (93.7%), efficient (42M params) |
+
+## Related Models
+
+- **magic-bert-50m-classification**: Same architecture fine-tuned for classification (89.7% accuracy)
+- **magic-bert-50m-roformer-mlm**: RoFormer variant with rotary position embeddings
+- **magic-bert-50m-roformer-classification**: RoFormer variant fine-tuned for classification (93.7% accuracy, recommended for production)
+
+## Related Work
+
+This model builds on the Binary BPE tokenization approach:
+
+- **Binary BPE Paper**: [Bommarito (2025)](https://arxiv.org/abs/2511.17573) introduced byte-level BPE tokenization for binary analysis, demonstrating 2-3x compression over raw bytes for executable content.
+- **Binary BPE Tokenizers**: Pre-trained tokenizers for executables are available at [mjbommar/binary-tokenizer-001-64k](https://huggingface.co/mjbommar/binary-tokenizer-001-64k).
+
+**Key difference**: The original Binary BPE work focused on executable binaries (ELF, PE, Mach-O). Magic-BERT extends this to general file type understanding across 106 diverse formats, using a tokenizer trained on the broader dataset.
+
+## Citation
+
+A paper describing Magic-BERT, the training methodology, and the dataset is forthcoming.
+
+```bibtex
+@article{bommarito2025binarybpe,
+  title={Binary BPE: A Family of Cross-Platform Tokenizers for Binary Analysis},
+  author={Bommarito, Michael J., II},
+  journal={arXiv preprint arXiv:2511.17573},
+  year={2025}
+}
+```

config.json

@@ -0,0 +1,24 @@
+{
+  "model_type": "magic-bert",
+  "architectures": [
+    "MagicBERTForMaskedLM"
+  ],
+  "vocab_size": 32768,
+  "hidden_size": 512,
+  "num_hidden_layers": 8,
+  "num_attention_heads": 8,
+  "intermediate_size": 2048,
+  "hidden_dropout_prob": 0.1,
+  "attention_probs_dropout_prob": 0.1,
+  "max_position_embeddings": 512,
+  "pad_token_id": 2,
+  "hidden_act": "gelu",
+  "layer_norm_eps": 1e-12,
+  "torch_dtype": "float32",
+  "transformers_version": "4.57.0",
+  "auto_map": {
+    "AutoConfig": "configuration_magic_bert.MagicBERTConfig",
+    "AutoModel": "modeling_magic_bert.MagicBERTModel",
+    "AutoModelForMaskedLM": "modeling_magic_bert.MagicBERTForMaskedLM"
+  }
+}

configuration_magic_bert.py

@@ -0,0 +1,40 @@
+"""MagicBERT configuration for HuggingFace transformers."""
+
+from transformers import PretrainedConfig
+
+
+class MagicBERTConfig(PretrainedConfig):
+    """Configuration class for MagicBERT model.
+
+    MagicBERT is a BERT-style transformer model designed for binary file
+    type classification. It uses a byte-level BPE tokenizer with a 32K vocabulary.
+    """
+
+    model_type = "magic-bert"
+
+    def __init__(
+        self,
+        vocab_size: int = 32768,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 8,
+        num_attention_heads: int = 8,
+        intermediate_size: int = 2048,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 512,
+        pad_token_id: int = 2,
+        hidden_act: str = "gelu",
+        layer_norm_eps: float = 1e-12,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_act = hidden_act
+        self.layer_norm_eps = layer_norm_eps

mime_type_mapping.json

@@ -0,0 +1,108 @@
+{
+  "0": "application/SIMH-tape-data",
+  "1": "application/encrypted",
+  "2": "application/gzip",
+  "3": "application/javascript",
+  "4": "application/json",
+  "5": "application/msword",
+  "6": "application/mxf",
+  "7": "application/octet-stream",
+  "8": "application/pdf",
+  "9": "application/pgp-keys",
+  "10": "application/postscript",
+  "11": "application/vnd.microsoft.portable-executable",
+  "12": "application/vnd.ms-excel",
+  "13": "application/vnd.ms-opentype",
+  "14": "application/vnd.ms-powerpoint",
+  "15": "application/vnd.oasis.opendocument.spreadsheet",
+  "16": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+  "17": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+  "18": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+  "19": "application/vnd.rn-realmedia",
+  "20": "application/vnd.wordperfect",
+  "21": "application/wasm",
+  "22": "application/x-7z-compressed",
+  "23": "application/x-archive",
+  "24": "application/x-bzip2",
+  "25": "application/x-coff",
+  "26": "application/x-dbf",
+  "27": "application/x-dosexec",
+  "28": "application/x-executable",
+  "29": "application/x-gettext-translation",
+  "30": "application/x-ms-ne-executable",
+  "31": "application/x-ndjson",
+  "32": "application/x-object",
+  "33": "application/x-ole-storage",
+  "34": "application/x-sharedlib",
+  "35": "application/x-shockwave-flash",
+  "36": "application/x-tar",
+  "37": "application/x-wine-extension-ini",
+  "38": "application/zip",
+  "39": "application/zlib",
+  "40": "application/zstd",
+  "41": "audio/amr",
+  "42": "audio/flac",
+  "43": "audio/mpeg",
+  "44": "audio/ogg",
+  "45": "audio/x-ape",
+  "46": "audio/x-hx-aac-adts",
+  "47": "audio/x-m4a",
+  "48": "audio/x-wav",
+  "49": "biosig/atf",
+  "50": "font/sfnt",
+  "51": "font/woff",
+  "52": "font/woff2",
+  "53": "image/bmp",
+  "54": "image/fits",
+  "55": "image/gif",
+  "56": "image/heif",
+  "57": "image/jpeg",
+  "58": "image/png",
+  "59": "image/svg+xml",
+  "60": "image/tiff",
+  "61": "image/vnd.adobe.photoshop",
+  "62": "image/vnd.microsoft.icon",
+  "63": "image/webp",
+  "64": "image/x-eps",
+  "65": "image/x-exr",
+  "66": "image/x-jp2-codestream",
+  "67": "image/x-portable-bitmap",
+  "68": "image/x-portable-greymap",
+  "69": "image/x-tga",
+  "70": "image/x-xpixmap",
+  "71": "inode/x-empty",
+  "72": "message/rfc822",
+  "73": "text/csv",
+  "74": "text/html",
+  "75": "text/plain",
+  "76": "text/rtf",
+  "77": "text/troff",
+  "78": "text/x-Algol68",
+  "79": "text/x-asm",
+  "80": "text/x-c",
+  "81": "text/x-c++",
+  "82": "text/x-diff",
+  "83": "text/x-file",
+  "84": "text/x-fortran",
+  "85": "text/x-java",
+  "86": "text/x-m4",
+  "87": "text/x-makefile",
+  "88": "text/x-msdos-batch",
+  "89": "text/x-perl",
+  "90": "text/x-php",
+  "91": "text/x-po",
+  "92": "text/x-ruby",
+  "93": "text/x-script.python",
+  "94": "text/x-shellscript",
+  "95": "text/x-tex",
+  "96": "text/xml",
+  "97": "video/3gpp",
+  "98": "video/mp4",
+  "99": "video/mpeg",
+  "100": "video/quicktime",
+  "101": "video/webm",
+  "102": "video/x-ivf",
+  "103": "video/x-matroska",
+  "104": "video/x-ms-asf",
+  "105": "video/x-msvideo"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:494f918a228fcd32d8967cb5decebafdf3b2d0e9a34601e6f9771387e0080d1f
+size 236291992

modeling_magic_bert.py

@@ -0,0 +1,346 @@
+"""MagicBERT model implementation for HuggingFace transformers.
+
+This module provides HuggingFace-compatible implementations of MagicBERT,
+a BERT-style model trained for binary file type understanding.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+    BaseModelOutput,
+)
+
+try:
+    from .configuration_magic_bert import MagicBERTConfig
+except ImportError:
+    from configuration_magic_bert import MagicBERTConfig
+
+
+class MagicBERTEmbeddings(nn.Module):
+    """MagicBERT embeddings: token + position embeddings."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__()
+        self.token_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.register_buffer(
+            "position_ids",
+            torch.arange(config.max_position_embeddings).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_length = input_ids.shape
+        token_embeds = self.token_embeddings(input_ids)
+        position_ids = self.position_ids[:, :seq_length]
+        position_embeds = self.position_embeddings(position_ids)
+        embeddings = token_embeds + position_embeds
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class MagicBERTAttention(nn.Module):
+    """Multi-head self-attention."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(new_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        query_layer = self.transpose_for_scores(self.query(hidden_states))
+        key_layer = self.transpose_for_scores(self.key(hidden_states))
+        value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, None, None, :]
+            attention_scores = attention_scores + (1.0 - attention_mask) * -10000.0
+
+        attention_probs = F.softmax(attention_scores, dim=-1)
+        attention_probs = self.dropout(attention_probs)
+        context = torch.matmul(attention_probs, value_layer)
+        context = context.permute(0, 2, 1, 3).contiguous()
+        new_shape = context.size()[:-2] + (self.all_head_size,)
+        context = context.view(new_shape)
+        return context
+
+
+class MagicBERTLayer(nn.Module):
+    """Single transformer layer."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__()
+        self.attention = MagicBERTAttention(config)
+        self.attention_output = nn.Linear(config.hidden_size, config.hidden_size)
+        self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.attention_dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.intermediate = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.output = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.output_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.output_dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # Self-attention with residual
+        attention_output = self.attention(hidden_states, attention_mask)
+        attention_output = self.attention_output(attention_output)
+        attention_output = self.attention_dropout(attention_output)
+        attention_output = self.attention_norm(hidden_states + attention_output)
+
+        # Feed-forward with residual
+        intermediate_output = self.intermediate(attention_output)
+        intermediate_output = F.gelu(intermediate_output)
+        layer_output = self.output(intermediate_output)
+        layer_output = self.output_dropout(layer_output)
+        layer_output = self.output_norm(attention_output + layer_output)
+        return layer_output
+
+
+class MagicBERTEncoder(nn.Module):
+    """Stack of transformer layers."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [MagicBERTLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask)
+        return hidden_states
+
+
+class MagicBERTPreTrainedModel(PreTrainedModel):
+    """Base class for MagicBERT models."""
+
+    config_class = MagicBERTConfig
+    base_model_prefix = "magic_bert"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class MagicBERTModel(MagicBERTPreTrainedModel):
+    """MagicBERT base model outputting raw hidden states."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__(config)
+        self.config = config
+        self.embeddings = MagicBERTEmbeddings(config)
+        self.encoder = MagicBERTEncoder(config)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,  # Ignored, for tokenizer compatibility
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], BaseModelOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        hidden_states = self.embeddings(input_ids)
+        sequence_output = self.encoder(hidden_states, attention_mask)
+        pooled_output = sequence_output[:, 0, :]
+
+        if not return_dict:
+            return (sequence_output, pooled_output)
+
+        return BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class MagicBERTForMaskedLM(MagicBERTPreTrainedModel):
+    """MagicBERT for masked language modeling (fill-mask task)."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__(config)
+        self.config = config
+        self.embeddings = MagicBERTEmbeddings(config)
+        self.encoder = MagicBERTEncoder(config)
+        self.mlm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,  # Ignored, for tokenizer compatibility
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor, ...], MaskedLMOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        hidden_states = self.embeddings(input_ids)
+        sequence_output = self.encoder(hidden_states, attention_mask)
+        logits = self.mlm_head(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def get_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        pooling: str = "cls",
+    ) -> torch.Tensor:
+        """Get embeddings for downstream tasks.
+
+        Args:
+            input_ids: Input token IDs
+            attention_mask: Attention mask
+            pooling: Pooling strategy ("cls" or "mean")
+
+        Returns:
+            Pooled embeddings [batch_size, hidden_size]
+        """
+        hidden_states = self.embeddings(input_ids)
+        sequence_output = self.encoder(hidden_states, attention_mask)
+
+        if pooling == "cls":
+            return sequence_output[:, 0, :]
+        elif pooling == "mean":
+            if attention_mask is not None:
+                mask = attention_mask.unsqueeze(-1).float()
+                return (sequence_output * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
+            return sequence_output.mean(dim=1)
+        else:
+            raise ValueError(f"Unknown pooling: {pooling}")
+
+
+class MagicBERTForSequenceClassification(MagicBERTPreTrainedModel):
+    """MagicBERT for sequence classification (file type classification)."""
+
+    def __init__(self, config: MagicBERTConfig):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = getattr(config, "num_labels", 106)
+
+        self.embeddings = MagicBERTEmbeddings(config)
+        self.encoder = MagicBERTEncoder(config)
+
+        # Projection head (for contrastive learning compatibility)
+        projection_dim = getattr(config, "contrastive_projection_dim", 256)
+        self.projection = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size),
+            nn.ReLU(),
+            nn.Linear(config.hidden_size, projection_dim),
+        )
+        self.classifier = nn.Linear(projection_dim, self.num_labels)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,  # Ignored, for tokenizer compatibility
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        hidden_states = self.embeddings(input_ids)
+        sequence_output = self.encoder(hidden_states, attention_mask)
+        pooled_output = sequence_output[:, 0, :]
+
+        projections = self.projection(pooled_output)
+        projections = F.normalize(projections, p=2, dim=1)
+        logits = self.classifier(projections)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def get_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Get normalized projection embeddings for similarity search."""
+        hidden_states = self.embeddings(input_ids)
+        sequence_output = self.encoder(hidden_states, attention_mask)
+        pooled_output = sequence_output[:, 0, :]
+        projections = self.projection(pooled_output)
+        return F.normalize(projections, p=2, dim=1)

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "mask_token": "[MASK]",
+  "cls_token": "[CLS]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}