Upload magic-bert-50m-roformer-classification model files

README.md

@@ -0,0 +1,343 @@
+---
+language:
+  - en
+license: apache-2.0
+library_name: transformers
+tags:
+  - binary-analysis
+  - file-type-detection
+  - byte-level
+  - classification
+  - mime-type
+  - roformer
+  - rope
+  - security
+pipeline_tag: text-classification
+base_model: magic-bert-50m-roformer-mlm
+model-index:
+  - name: magic-bert-50m-roformer-classification
+    results:
+      - task:
+          type: text-classification
+          name: File Type Classification
+        metrics:
+          - name: Probing Accuracy
+            type: accuracy
+            value: 93.7
+          - name: Silhouette Score
+            type: silhouette
+            value: 0.663
+          - name: F1 (Weighted)
+            type: f1
+            value: 0.933
+---
+
+# Magic-BERT 50M RoFormer Classification
+
+A RoFormer-based transformer model fine-tuned for binary file type classification. This model achieves 93.7% classification accuracy across 106 MIME types, making it the **recommended choice for production file type detection**.
+
+## 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-roformer-mlm with contrastive learning fine-tuning. It uses Rotary Position Embeddings (RoPE) and produces highly discriminative embeddings for file type classification.
+
+| Property | Value |
+|----------|-------|
+| Parameters | 42.0M (+ 0.45M classifier head) |
+| Hidden Size | 512 |
+| Projection Dimension | 256 |
+| Number of Classes | 106 MIME types |
+| Base Model | magic-bert-50m-roformer-mlm |
+| Position Encoding | RoPE (Rotary Position 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.
+
+## Intended Uses
+
+**Primary use cases:**
+- Production file type classification
+- MIME type detection from binary content
+- Embedding-based file similarity search
+- Security analysis and content filtering
+
+This is the recommended model for file classification tasks due to its combination of high accuracy (93.7%) and parameter efficiency (42M parameters).
+
+## 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.
+
+## Architecture: RoPE vs Absolute Position Embeddings
+
+This model uses **Rotary Position Embeddings (RoPE)**, which encode position through rotation matrices in attention. This differs from the Magic-BERT variant which uses absolute position embeddings.
+
+| Metric | RoFormer (this) | Magic-BERT |
+|--------|-----------------|------------|
+| Classification Accuracy | **93.7%** | 89.7% |
+| Silhouette Score | **0.663** | 0.55 |
+| F1 (Weighted) | **0.933** | 0.886 |
+| Parameters | **42.5M** | 59M |
+| Fill-mask Retention | 14.5% | **41.8%** |
+
+This model achieves higher classification accuracy with fewer parameters, making it the preferred choice for production deployment when only classification is needed.
+
+## MLM vs Classification: Two-Phase Training
+
+This is the **Phase 2 (Classification)** model built on RoFormer. The training pipeline has two phases:
+
+| Phase | Model | Task | Purpose |
+|-------|-------|------|---------|
+| Phase 1 | magic-bert-50m-roformer-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
+- Result: Significantly improved embedding quality for classification
+
+## Evaluation Results
+
+### Classification Performance
+
+| Metric | Value |
+|--------|-------|
+| Linear Probe Accuracy | **93.7%** |
+| F1 (Macro) | 0.829 |
+| F1 (Weighted) | 0.933 |
+
+### Embedding Quality
+
+| Metric | Value |
+|--------|-------|
+| Silhouette Score | **0.663** |
+| Separation Ratio | 4.00 |
+| Intra-class Distance | 7.24 |
+| Inter-class Distance | 28.98 |
+
+The silhouette score of 0.663 indicates well-separated clusters, suitable for embedding-based retrieval and similarity search.
+
+### Phase 1 → Phase 2 Improvement
+
+| Metric | Phase 1 | Phase 2 | Change |
+|--------|---------|---------|--------|
+| Probing Accuracy | 85.0% | 93.7% | +8.7% |
+| Silhouette Score | 0.328 | 0.663 | +102% |
+| Separation Ratio | 2.65 | 4.00 | +51% |
+
+## Supported MIME Types (106 Classes)
+
+The model classifies files into 106 MIME types across these categories:
+
+| Category | Count | Examples | Typical Accuracy |
+|----------|-------|----------|------------------|
+| application/ | 41 | PDF, ZIP, GZIP, Office docs, executables | >90% |
+| text/ | 24 | Python, C, Java, HTML, XML, shell scripts | >80% |
+| image/ | 18 | PNG, JPEG, GIF, WebP, TIFF, PSD | >95% |
+| video/ | 9 | MP4, WebM, MKV, AVI, MOV | >90% |
+| audio/ | 8 | MP3, FLAC, WAV, OGG, M4A | >90% |
+| font/ | 3 | SFNT, WOFF, WOFF2 | >85% |
+| other | 3 | biosig/atf, inode/x-empty, message/rfc822 | varies |
+
+<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 RoFormerModel, AutoTokenizer
+from safetensors.torch import load_file
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+
+# Load tokenizer and MIME mapping
+tokenizer = AutoTokenizer.from_pretrained("path/to/magic-bert-50m-roformer-classification")
+with open("path/to/magic-bert-50m-roformer-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 base model
+base_model = RoFormerModel.from_pretrained("path/to/magic-bert-50m-roformer-classification")
+
+# Create classification head
+class ClassificationHead(nn.Module):
+    def __init__(self, hidden_size=512, projection_dim=256, num_classes=106):
+        super().__init__()
+        self.projection = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Linear(hidden_size, projection_dim),
+        )
+        self.classifier = nn.Linear(projection_dim, num_classes)
+
+    def forward(self, hidden_states):
+        pooled = hidden_states[:, 0, :]  # CLS token
+        projected = self.projection(pooled)
+        projected = F.normalize(projected, p=2, dim=1)
+        return self.classifier(projected), projected
+
+head = ClassificationHead()
+contrastive_dict = load_file("path/to/magic-bert-50m-roformer-classification/contrastive_head.safetensors")
+head.projection.load_state_dict({k.replace("projection.", ""): v for k, v in contrastive_dict.items() if "projection" in k})
+head.classifier.load_state_dict({k.replace("classifier.", ""): v for k, v in contrastive_dict.items() if "classifier" in k})
+
+base_model.eval()
+head.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 = base_model(**inputs)
+    logits, embeddings = head(outputs.last_hidden_state)
+    predicted_id = logits.argmax(-1).item()
+
+print(f"Predicted MIME type: {id_to_mime[predicted_id]}")
+print(f"Confidence: {F.softmax(logits, dim=-1).max().item():.2%}")
+```
+
+### Embedding-Based Similarity Search
+
+```python
+# Get normalized embeddings for similarity search
+with torch.no_grad():
+    outputs = base_model(**inputs)
+    _, embeddings = head(outputs.last_hidden_state)
+    # embeddings shape: [batch_size, 256], L2 normalized
+
+# Compute cosine similarity
+similarity = torch.mm(embeddings1, embeddings2.T)
+
+# Find most similar files
+top_k = similarity[0].topk(5)
+```
+
+## Limitations
+
+1. **MLM capability sacrificed:** Fill-mask accuracy drops to 14.5% after classification fine-tuning. Use the MLM variant if byte prediction is needed.
+
+2. **Position bias:** Still present (~46% accuracy drop at offset 1000), though less relevant for classification than for fill-mask tasks.
+
+3. **Ambiguous formats:** ZIP-based formats (DOCX, XLSX, JAR, APK) share similar structure and may be confused.
+
+4. **Rare types:** Lower accuracy on underrepresented file types in training data.
+
+## Model Selection Guide
+
+| Use Case | Recommended Model | Reason |
+|----------|-------------------|--------|
+| **Production classification** | **This model** | Highest accuracy (93.7%), efficient (42M params) |
+| Classification + fill-mask | magic-bert-50m-classification | Retains 41.8% fill-mask capability |
+| Fill-mask / byte prediction | magic-bert-50m-roformer-mlm | Optimized for MLM |
+| Research baseline | magic-bert-50m-mlm | Best perplexity (1.05) |
+
+## Related Models
+
+- **magic-bert-50m-roformer-mlm**: Base model before classification fine-tuning
+- **magic-bert-50m-mlm**: Absolute position embedding variant (MLM)
+- **magic-bert-50m-classification**: Magic-BERT variant that retains better fill-mask capability (89.7% accuracy)
+
+## 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,25 @@
+{
+  "architectures": [
+    "RoFormerForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "embedding_size": 512,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "roformer",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 8,
+  "pad_token_id": 2,
+  "rotary_value": false,
+  "transformers_version": "4.57.3",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 32768,
+  "num_labels": 106,
+  "problem_type": "single_label_classification"
+}

contrastive_head.safetensors

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

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:2a3c0be25fef5e6e5da5c470a8feec34d20ad3a7467bdb3fb742fd521310b639
+size 169324736

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