Upload magic-bert-50m-classification model files

README.md

@@ -0,0 +1,324 @@
+---
+language:
+  - en
+license: apache-2.0
+library_name: transformers
+tags:
+  - binary-analysis
+  - file-type-detection
+  - byte-level
+  - classification
+  - mime-type
+  - security
+pipeline_tag: text-classification
+base_model: magic-bert-50m-mlm
+model-index:
+  - name: magic-bert-50m-classification
+    results:
+      - task:
+          type: text-classification
+          name: File Type Classification
+        metrics:
+          - name: Probing Accuracy
+            type: accuracy
+            value: 89.7
+          - name: Silhouette Score
+            type: silhouette
+            value: 0.55
+          - name: F1 (Weighted)
+            type: f1
+            value: 0.886
+---
+
+# Magic-BERT 50M Classification
+
+A BERT-style transformer model fine-tuned for binary file type classification. This model classifies binary files into 106 MIME types based on their content structure.
+
+## 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
+
+This model extends magic-bert-50m-mlm with contrastive learning fine-tuning to produce embeddings optimized for file type discrimination. It uses a projection head and classifier trained with supervised contrastive loss.
+
+| Property | Value |
+|----------|-------|
+| Parameters | 59M (+ 0.4M classifier head) |
+| Hidden Size | 512 |
+| Projection Dimension | 256 |
+| Number of Classes | 106 MIME types |
+| Base Model | magic-bert-50m-mlm |
+
+### 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.
+
+## Intended Uses
+
+**Primary use cases:**
+- File type classification from binary content
+- MIME type detection without relying on file extensions
+- Embedding-based file similarity search
+- Security analysis and malware triage
+
+**Example tasks:**
+- Identifying file types in network traffic
+- Classifying files with missing or incorrect extensions
+- Building file type indexes for large archives
+
+## 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. Classification embeddings 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. This model 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.
+
+### Similarity Search
+The embedding space (256-dimensional, L2-normalized) enables similarity search across file collections: "find files structurally similar to this sample" for malware clustering, duplicate detection, or content-based retrieval.
+
+## MLM vs Classification: Two-Phase Training
+
+This is the **Phase 2 (Classification)** model built on Magic-BERT. The training pipeline has two phases:
+
+| Phase | Model | Task | Purpose |
+|-------|-------|------|---------|
+| Phase 1 | magic-bert-50m-mlm | Masked Language Modeling | Learn byte-level patterns and file structure |
+| **Phase 2** | **This model** | Contrastive Learning | Optimize embeddings for file type discrimination |
+
+### Two-Phase Training
+
+| Phase | Steps | Learning Rate | Objective |
+|-------|-------|---------------|-----------|
+| 1: MLM Pre-training | 100,000 | 1e-4 | Masked Language Modeling |
+| 2: Contrastive Fine-tuning | 50,000 | 1e-6 | Supervised Contrastive Loss |
+
+**Phase 2 specifics:**
+- Frozen: Embeddings + first 4 transformer layers
+- Learning rate: 100x lower than Phase 1
+- Objective: Pull same-MIME-type samples together, push different types apart
+
+## Evaluation Results
+
+### Classification Performance
+
+| Metric | Value |
+|--------|-------|
+| Linear Probe Accuracy | 89.7% |
+| F1 (Macro) | 0.787 |
+| F1 (Weighted) | 0.886 |
+
+### Embedding Quality
+
+| Metric | Value |
+|--------|-------|
+| Silhouette Score | 0.55 |
+| Separation Ratio | 3.60 |
+| Intra-class Distance | 12.6 |
+| Inter-class Distance | 45.2 |
+
+### MLM Capability (Retained)
+
+| Metric | Value |
+|--------|-------|
+| Fill-mask Top-1 | 41.8% |
+| Perplexity | 1.32 |
+
+This model retains moderate fill-mask capability, making it suitable for hybrid tasks that need both classification and byte prediction.
+
+## Supported MIME Types (106 Classes)
+
+The model classifies files into 106 MIME types across these categories:
+
+| Category | Count | Examples |
+|----------|-------|----------|
+| application/ | 41 | PDF, ZIP, GZIP, Office docs, executables |
+| text/ | 24 | Python, C, Java, HTML, XML, shell scripts |
+| image/ | 18 | PNG, JPEG, GIF, WebP, TIFF, PSD |
+| video/ | 9 | MP4, WebM, MKV, AVI, MOV |
+| audio/ | 8 | MP3, FLAC, WAV, OGG, M4A |
+| font/ | 3 | SFNT, WOFF, WOFF2 |
+| other | 3 | biosig/atf, inode/x-empty, message/rfc822 |
+
+<details>
+<summary>Click to expand full MIME type list</summary>
+
+**application/** (41 types):
+- application/SIMH-tape-data, application/encrypted, application/gzip
+- application/javascript, application/json, application/msword
+- application/mxf, application/octet-stream, application/pdf
+- application/pgp-keys, application/postscript
+- application/vnd.microsoft.portable-executable, application/vnd.ms-excel
+- application/vnd.ms-opentype, application/vnd.ms-powerpoint
+- application/vnd.oasis.opendocument.spreadsheet
+- application/vnd.openxmlformats-officedocument.* (3 variants)
+- application/vnd.rn-realmedia, application/vnd.wordperfect
+- application/wasm, application/x-7z-compressed, application/x-archive
+- application/x-bzip2, application/x-coff, application/x-dbf
+- application/x-dosexec, application/x-executable
+- application/x-gettext-translation, application/x-ms-ne-executable
+- application/x-ndjson, application/x-object, application/x-ole-storage
+- application/x-sharedlib, application/x-shockwave-flash
+- application/x-tar, application/x-wine-extension-ini
+- application/zip, application/zlib, application/zstd
+
+**text/** (24 types):
+- text/csv, text/html, text/plain, text/rtf, text/troff
+- text/x-Algol68, text/x-asm, text/x-c, text/x-c++
+- text/x-diff, text/x-file, text/x-fortran, text/x-java
+- text/x-m4, text/x-makefile, text/x-msdos-batch, text/x-perl
+- text/x-php, text/x-po, text/x-ruby, text/x-script.python
+- text/x-shellscript, text/x-tex, text/xml
+
+**image/** (18 types):
+- image/bmp, image/fits, image/gif, image/heif, image/jpeg
+- image/png, image/svg+xml, image/tiff, image/vnd.adobe.photoshop
+- image/vnd.microsoft.icon, image/webp, image/x-eps, image/x-exr
+- image/x-jp2-codestream, image/x-portable-bitmap
+- image/x-portable-greymap, image/x-tga, image/x-xpixmap
+
+**video/** (9 types):
+- video/3gpp, video/mp4, video/mpeg, video/quicktime, video/webm
+- video/x-ivf, video/x-matroska, video/x-ms-asf, video/x-msvideo
+
+**audio/** (8 types):
+- audio/amr, audio/flac, audio/mpeg, audio/ogg, audio/x-ape
+- audio/x-hx-aac-adts, audio/x-m4a, audio/x-wav
+
+**font/** (3 types):
+- font/sfnt, font/woff, font/woff2
+
+**other** (3 types):
+- biosig/atf, inode/x-empty, message/rfc822
+
+</details>
+
+## How to Use
+
+```python
+from transformers import AutoTokenizer
+from safetensors.torch import load_file
+import torch
+import json
+
+# Load tokenizer and MIME mapping
+tokenizer = AutoTokenizer.from_pretrained("path/to/magic-bert-50m-classification")
+with open("path/to/magic-bert-50m-classification/mime_type_mapping.json") as f:
+    mime_mapping = json.load(f)
+id_to_mime = {int(k): v for k, v in mime_mapping.items()}
+
+# Load model
+from modeling_magic_bert import MagicBERTForSequenceClassification
+from configuration_magic_bert import MagicBERTConfig
+
+config = MagicBERTConfig.from_pretrained("path/to/magic-bert-50m-classification")
+model = MagicBERTForSequenceClassification(config)
+
+# Load base model weights
+state_dict = load_file("path/to/magic-bert-50m-classification/model.safetensors")
+model.load_state_dict(state_dict, strict=False)
+
+# Load contrastive head weights
+contrastive_dict = load_file("path/to/magic-bert-50m-classification/contrastive_head.safetensors")
+model.projection.load_state_dict({k.replace("projection.", ""): v for k, v in contrastive_dict.items() if "projection" in k})
+model.classifier.load_state_dict({k.replace("classifier.", ""): v for k, v in contrastive_dict.items() if "classifier" in k})
+
+model.eval()
+
+# Classify a file
+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")
+
+inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+
+with torch.no_grad():
+    outputs = model(**inputs)
+    predicted_id = outputs.logits.argmax(-1).item()
+
+print(f"Predicted MIME type: {id_to_mime[predicted_id]}")
+```
+
+### Getting Embeddings for Similarity Search
+
+```python
+# Get normalized projection embeddings
+with torch.no_grad():
+    embeddings = model.get_embeddings(inputs["input_ids"], inputs["attention_mask"])
+    # embeddings shape: [batch_size, 256], L2 normalized
+
+# Compute cosine similarity between files
+similarity = torch.mm(embeddings1, embeddings2.T)
+```
+
+## Limitations
+
+1. **Position bias:** Best performance when content starts at position 0. Accuracy degrades for content at higher offsets.
+
+2. **Class imbalance:** Performance varies by file type. Common formats (PDF, PNG, ZIP) perform better than rare formats.
+
+3. **Ambiguous types:** Some file types share similar structure (e.g., ZIP-based formats like DOCX, XLSX, JAR), which can cause confusion.
+
+4. **Encrypted content:** Cannot classify encrypted or compressed content that lacks recognizable patterns.
+
+## Architecture: Absolute vs Rotary Position Embeddings
+
+This model uses **absolute position embeddings**, where each position (0-511) has a learned embedding vector. An alternative is **Rotary Position Embeddings (RoPE)**, used by the RoFormer variant.
+
+| Metric | Magic-BERT (this) | RoFormer |
+|--------|-------------------|----------|
+| Classification Accuracy | 89.7% | **93.7%** |
+| Silhouette Score | 0.55 | **0.663** |
+| F1 (Weighted) | 0.886 | **0.933** |
+| Fill-mask Retention | **41.8%** | 14.5% |
+| Parameters | 59M | **42M** |
+
+Magic-BERT retains better fill-mask capability after classification fine-tuning, making it suitable when both tasks are needed. For pure classification, consider the RoFormer variant.
+
+## Model Selection Guide
+
+| Use Case | Recommended Model | Reason |
+|----------|-------------------|--------|
+| Classification + fill-mask | **This model** | Retains 41.8% fill-mask capability |
+| Fill-mask / byte prediction | magic-bert-50m-mlm | Best perplexity (1.05) |
+| Research baseline | magic-bert-50m-mlm | Established BERT architecture |
+| **Production classification** | **magic-bert-50m-roformer-classification** | Highest accuracy (93.7%), efficient (42M params) |
+
+## Related Models
+
+- **magic-bert-50m-mlm**: Base model before classification fine-tuning
+- **magic-bert-50m-roformer-mlm**: RoFormer variant with rotary position embeddings
+- **magic-bert-50m-roformer-classification**: RoFormer variant with higher classification accuracy (93.7%, 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,27 @@
+{
+  "model_type": "magic-bert",
+  "architectures": [
+    "MagicBERTForSequenceClassification"
+  ],
+  "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",
+    "AutoModelForSequenceClassification": "modeling_magic_bert.MagicBERTForSequenceClassification"
+  },
+  "num_labels": 106,
+  "contrastive_projection_dim": 256
+}

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

contrastive_head.safetensors

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

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:a3923cd4384639bde231f53f2b40822cc71fdc920d43bf4b97a5b6edafad3d2c
+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]"
+}