Upload magic-bert-50m-classification model files

README.md

@@ -205,32 +205,13 @@ The model classifies files into 106 MIME types across these categories:
 ## How to Use
 
 ```python
-from transformers import AutoTokenizer
-from safetensors.torch import load_file
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 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 = AutoModelForSequenceClassification.from_pretrained(
+    "mjbommar/magic-bert-50m-classification", trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained("mjbommar/magic-bert-50m-classification")
 
 model.eval()
 
@@ -240,28 +221,42 @@ with open("example.pdf", "rb") as f:
 
 # 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()
+    confidence = torch.softmax(outputs.logits, dim=-1).max().item()
 
-print(f"Predicted MIME type: {id_to_mime[predicted_id]}")
+print(f"Predicted class: {predicted_id}")
+print(f"Confidence: {confidence:.2%}")
 ```
 
 ### Getting Embeddings for Similarity Search
 
 ```python
-# Get normalized projection embeddings
+# Get normalized embeddings (256-dim, L2-normalized)
 with torch.no_grad():
     embeddings = model.get_embeddings(inputs["input_ids"], inputs["attention_mask"])
-    # embeddings shape: [batch_size, 256], L2 normalized
+    # embeddings shape: [batch_size, 256]
 
 # Compute cosine similarity between files
 similarity = torch.mm(embeddings1, embeddings2.T)
 ```
 
+### Loading MIME Type Labels
+
+```python
+from huggingface_hub import hf_hub_download
+import json
+
+mime_path = hf_hub_download("mjbommar/magic-bert-50m-classification", "mime_type_mapping.json")
+with open(mime_path) as f:
+    id_to_mime = {int(k): v for k, v in json.load(f).items()}
+
+print(f"Predicted MIME type: {id_to_mime[predicted_id]}")
+```
+
 ## Limitations
 
 1. **Position bias:** Best performance when content starts at position 0. Accuracy degrades for content at higher offsets.

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a3923cd4384639bde231f53f2b40822cc71fdc920d43bf4b97a5b6edafad3d2c
-size 236291992
+oid sha256:93b0ae348f5fac2a4eac22b83b2540fd21bb0b45e4d308c26d9f849ebc1ebd22
+size 170737312