Update model to new HF format

README.md

@@ -17,9 +17,17 @@ paper: https://arxiv.org/abs/2510.19410
 
 # ToMMeR-pythia-2.8b_L5_R64
 
+
+[![Paper](https://img.shields.io/badge/Paper-Arxiv-red)](https://arxiv.org/abs/2510.19410)
+[![All Models](https://img.shields.io/badge/🤗%20Hugging%20Face%20Models-blue)](https://huggingface.co/llm2ner)
+[![GitHub](https://img.shields.io/badge/GitHub-Code-blue)](https://github.com/VictorMorand/llm2ner)
+
+
 ToMMeR is a lightweight probing model extracting emergent mention detection capabilities from early layers representations of any LLM backbone, achieving high Zero Shot recall across a wide set of 13 NER benchmarks.
 
-## Checkpoint Details
+## Model Details
+
+This model can be plugged at layer 5 of `EleutherAI/pythia-2.8b`, with a computational overhead not greater than an additional attention head.
 
 | Property  | Value |
 |-----------|-------|
@@ -31,28 +39,69 @@ ToMMeR is a lightweight probing model extracting emergent mention detection capa
 # Usage
 
 ## Installation
-
-Our code can be installed with pip+git, Please visit the [repository](https://github.com/VictorMorand/llm2ner) for more details.
+To use ToMMeR, you need to install its codebase first.
 
 ```bash
 pip install git+https://github.com/VictorMorand/llm2ner.git
 ```
 
+
+## Raw inference
+By default, ToMMeR outputs span probabilities, but we also propose built-in options for decoding entities.
+
+- Inputs:
+  - tokens (batch, seq): tokens to process,
+  - model: LLM to extract representation from.
+- Outputs: (batch, seq, seq) matrix (masked outside valid spans)
+
+```python
+from xpm_torch.huggingface import TorchHFHub
+from llm2ner import ToMMeR, utils
+
+tommer: ToMMeR = TorchHFHub.from_pretrained("llm2ner/ToMMeR-pythia-2.8b_L5_R64")
+# load Backbone llm, optionnally cut the unused layer to save GPU space.
+llm = utils.load_llm( tommer.llm_name, cut_to_layer=tommer.layer,)
+tommer.to(llm.device)
+
+#### Raw Inference
+text = ["Large language models are awesome"]
+print(f"Input text: {text[0]}")
+
+#tokenize in shape (1, seq_len)
+tokens = llm.tokenizer(text, return_tensors="pt")["input_ids"].to(llm.device)
+# Output raw scores
+output = tommer.forward(tokens, llm) # (batch_size, seq_len, seq_len)
+print(f"Raw Output shape: {output.shape}")
+
+#use given decoding strategy to infer entities
+entities = tommer.infer_entities(tokens=tokens, model=llm, threshold=0.5, decoding_strategy="greedy")
+str_entities = [ llm.tokenizer.decode(tokens[0,b:e+1]) for b, e in entities[0]]
+print(f"Predicted entities: {str_entities}")
+
+>>>INFO:root:Cut LlamaModel with 16 layers to 7 layers
+>>> Input text: Large language models are awesome
+>>> Raw Output shape: torch.Size([1, 6, 6])
+>>> Predicted entities: ['Large language models']
+```
+
+
 ## Fancy Outputs
 
+We also provide inference and plotting utils in `llm2ner.plotting`.
+
 ```python
-import llm2ner
-from llm2ner import ToMMeR
+from xpm_torch.huggingface import TorchHFHub
+from llm2ner import ToMMeR, utils, plotting
 
-tommer = ToMMeR.from_pretrained("llm2ner/ToMMeR-pythia-2.8b_L5_R64")
+tommer: ToMMeR = TorchHFHub.from_pretrained("llm2ner/ToMMeR-pythia-2.8b_L5_R64")
 # load Backbone llm, optionnally cut the unused layer to save GPU space.
-llm = llm2ner.utils.load_llm( tommer.llm_name, cut_to_layer=tommer.layer,) 
+llm = utils.load_llm( tommer.llm_name, cut_to_layer=tommer.layer,)
 tommer.to(llm.device)
 
 text = "Large language models are awesome. While trained on language modeling, they exhibit emergent Zero Shot abilities that make them suitable for a wide range of tasks, including Named Entity Recognition (NER). "
 
 #fancy interactive output
-outputs = llm2ner.plotting.demo_inference( text, tommer, llm,
+outputs = plotting.demo_inference( text, tommer, llm,
     decoding_strategy="threshold",  # or "greedy" for flat segmentation
     threshold=0.5, # default 50%
     show_attn=True,
@@ -89,7 +138,7 @@ outputs = llm2ner.plotting.demo_inference( text, tommer, llm,
 <span style="background: lightblue; top: 57px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
 </span>
 </span>
-are awesome . While trained on 
+are awesome . While trained on
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     language
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
@@ -105,7 +154,7 @@ are awesome . While trained on
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
 </span>
 </span>
-, they exhibit 
+, they exhibit
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     emergent
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
@@ -121,18 +170,18 @@ are awesome . While trained on
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
 </span>
 </span>
-that make them suitable for a wide range of 
+that make them suitable for a wide range of
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     tasks
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
-</span> 
+</span>
 <span style="background: lightblue; top: 40px; height: 4px; border-top-left-radius: 3px; border-bottom-left-radius: 3px; left: -1px; width: calc(100% + 2px); position: absolute;">
     <span style="background: lightblue; z-index: 10; color: #000; top: -0.5em; padding: 2px 3px; position: absolute; font-size: 0.6em; font-weight: bold; line-height: 1; border-radius: 3px">
         PRED
     </span>
 </span>
 </span>
-, including 
+, including
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     Named
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
@@ -145,7 +194,7 @@ that make them suitable for a wide range of
 </span>
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     Entity
-    
+
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
 </span>
 </span>
@@ -154,7 +203,7 @@ that make them suitable for a wide range of
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
 </span>
 </span>
-( 
+(
 <span style="font-weight: bold; display: inline-block; position: relative; height: 60px;">
     NER
 <span style="background: lightblue; top: 40px; height: 4px; left: -1px; width: calc(100% + 2px); position: absolute;">
@@ -168,42 +217,6 @@ that make them suitable for a wide range of
 ) . </div></span>
 </div>
 
-
-## Raw inference
-By default, ToMMeR outputs span probabilities, but we also propose built-in options for decoding entities.
-
-- Inputs:
-  - tokens (batch, seq): tokens to process, 
-  - model: LLM to extract representation from.
-- Outputs: (batch, seq, seq) matrix (masked outside valid spans)
-
-```python
-
-tommer = ToMMeR.from_pretrained("llm2ner/ToMMeR-pythia-2.8b_L5_R64")
-# load Backbone llm, optionnally cut the unused layer to save GPU space.
-llm = llm2ner.utils.load_llm( tommer.llm_name, cut_to_layer=tommer.layer,) 
-tommer.to(llm.device)
-
-#### Raw Inference
-text = ["Large language models are awesome"]
-print(f"Input text: {text[0]}")
-
-#tokenize in shape (1, seq_len)
-tokens = model.tokenizer(text, return_tensors="pt")["input_ids"].to(device)
-# Output raw scores
-output = tommer.forward(tokens, model) # (batch_size, seq_len, seq_len)
-print(f"Raw Output shape: {output.shape}")
-
-#use given decoding strategy to infer entities
-entities = tommer.infer_entities(tokens=tokens, model=model, threshold=0.5, decoding_strategy="greedy")
-str_entities = [ model.tokenizer.decode(tokens[0,b:e+1]) for b, e in entities[0]]
-print(f"Predicted entities: {str_entities}")
-
->>> Input text: Large language models are awesome
->>> Raw Output shape: torch.Size([1, 6, 6])
->>> Predicted entities: ['Large language models']
-```
-
 Please visit the [repository](https://github.com/VictorMorand/llm2ner) for more details and a demo notebook.
 
 ## Evaluation Results
@@ -225,18 +238,17 @@ Please visit the [repository](https://github.com/VictorMorand/llm2ner) for more
 | Ontonotes           |      0.2296 |   0.6734 | 0.3424 |       42193 |
 | Aggregated          |      0.2121 |   0.8771 | 0.3415 |      353250 |
 | Mean                |      0.2633 |   0.8198 | 0.3904 |      353250 |
-
 ## Citation
 If using this model or the approach, please cite the associated paper:
 ```
 @misc{morand2025tommerefficiententity,
-      title={ToMMeR -- Efficient Entity Mention Detection from Large Language Models}, 
+      title={ToMMeR -- Efficient Entity Mention Detection from Large Language Models},
       author={Victor Morand and Nadi Tomeh and Josiane Mothe and Benjamin Piwowarski},
       year={2025},
       eprint={2510.19410},
       archivePrefix={arXiv},
       primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2510.19410}, 
+      url={https://arxiv.org/abs/2510.19410},
 }
 ```
 

definition.json

@@ -1 +0,0 @@
-{"objects": [{"id": 140521468942352, "module": "llm2ner.models.tommer", "type": "ToMMeR", "typename": "llm2ner.models.tommer.ToMMeR", "identifier": "8e3a62b0403a4172fcb037667ae94cd59bccde77e6f06ae1f0756664fd6a35db", "fields": {"llm_name": "EleutherAI/pythia-2.8b", "layer": 5, "rank": 64, "causal_mask": true, "sliding_window": 25, "use_cosine": true, "normalize_scores": ""}}, {"id": 140521600014832, "module": "llm2ner.xpmModel", "type": "xpmTorchHubModule.Loader", "typename": "llm2ner.xpmModel.xpmTorchHubModule.Loader", "identifier": "9f137eae73d32a9be6ff608a808bbe7cdc731525ff20a1095cdbdb68ff059e56", "fields": {"model": {"type": "python", "value": 140521468942352}, "parameters": {"type": "path.serialized", "value": "parameters", "is_folder": false}}}], "data": [{"type": "python", "value": 140521468942352}, [{"type": "python", "value": 140521600014832}]]}

experimaestro.json

@@ -0,0 +1 @@
+[{"id": 5605317632, "module": "llm2ner.models.tommer", "type": "ToMMeR", "typename": "llm2ner.models.tommer.ToMMeR", "identifier": "c3bc1cd395a94210f0aac837c568ae710f731328e8711bcf00ea64fc43578279", "fields": {"llm_name": "EleutherAI/pythia-2.8b", "layer": 5, "rank": 64, "causal_mask": true, "sliding_window": 25, "use_cosine": true, "normalize_scores": ""}}, {"id": 5605317536, "module": "xpm_torch.module", "type": "SimpleModuleLoader", "typename": "xpm_torch.module.SimpleModuleLoader", "identifier": "821960ec1ccf4291eab475ab8e989519810584d3ce0982d1403bf8262eae5509", "fields": {"value": {"type": "python", "value": 5605317632}, "settings": null, "path": {"type": "path.serialized", "value": "model.safetensors", "is_folder": false}}}]

parameters → model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:35db9a84b09b1c9f5e5ee39cb23e215b6f580aa6a839adb1597c581b9bb4d0db
-size 1323770
+oid sha256:5d7ee3e91558d3fc68125685fdc23bb55df3adf76513ef59f7757b8756c63c4b
+size 1321376