Upload folder using huggingface_hub

README.md

@@ -1,262 +1,287 @@
----
-language: en
-license: apache-2.0
-library_name: transformers
-
-pipeline_tag: text-classification
-task_categories:
-  - text-classification
-
-model_type: sproto
-base_model: microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext
-
-datasets:
-  - mimic-iv
-
-metrics:
-  - auroc
-  - pr-auc
-
-tags:
-  - text-classification
-  - multi-label-classification
-  - long-tail-learning
-  - medical
-  - clinical-nlp
-  - interpretability
-  - prototypical-networks
-  - ehr
----
-
-# S-Proto: Sparse Prototypical Networks for Long-Tail Clinical Diagnosis Prediction
-
-**Published at ECML PKDD 2024 (CORE A)**  
-*Boosting Long-Tail Data Classification with Sparse Prototypical Networks*
-
-Alexei Figueroa*, Jens-Michalis Papaioannou*, et al.  
-DATEXIS, Berliner Hochschule für Technik, Feinstein Institutes, TU Munich, Leibniz University Hannover  
-(* equal contribution)
-
-![S-Proto](overview.png)
-
-This repository provides **S-Proto**, a sparse and interpretable prototypical network for extreme multi-label diagnosis prediction from clinical text. The model is designed to address the long-tail distribution of clinical diagnoses while preserving faithful, prototype-based explanations.
-
-## Interactive Demo
-
-You can explore the model's predictions and interpretability features through our interactive web demo:
-**[https://s-proto.demo.datexis.com/](https://s-proto.demo.datexis.com/)**
-
-S-Proto was introduced in the paper:
-
-**[Boosting Long-Tail Data Classification with Sparse Prototypical Networks](https://ecmlpkdd-storage.s3.eu-central-1.amazonaws.com/preprints/2024/lncs14947/lncs14947435.pdf)**
-European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML PKDD 2024, CORE A)  
-Alexei Figueroa*, Jens-Michalis Papaioannou*, et al.  
-DATEXIS, Berliner Hochschule für Technik, Feinstein Institutes, TU Munich, Leibniz University Hannover  
-(* equal contribution)
-
-## Overview
-
-Clinical outcome prediction from Electronic Health Records is characterized by extreme label imbalance. A small number of diagnoses account for most patients, while the majority of diagnoses appear rarely. Standard transformer classifiers tend to perform well on frequent diagnoses but degrade sharply in the long tail.
-
-S-Proto addresses this problem by extending prototypical networks with:
-
-- Multiple prototypes per diagnosis  
-- Sparse winner-takes-all activation  
-- Prototype-level interpretability  
-- Efficient training despite increased representational capacity  
-
-The model achieves state-of-the-art performance on MIMIC-IV diagnosis prediction, with particularly strong gains in PR-AUC for rare diagnoses, and transfers successfully to unseen clinical datasets.
-
-## Model Architecture
-
-S-Proto builds on **PubMedBERT** as the text encoder and introduces a sparse prototypical layer on top.
-
-For each diagnosis label, the model learns multiple sub-networks, each consisting of:
-
-- A label-specific attention vector  
-- A prototype vector representing a prototypical patient  
-
-Given an input clinical note:
-
-1. The note is encoded using PubMedBERT  
-2. Token embeddings are projected into a latent space  
-3. Each diagnosis activates multiple candidate sub-networks  
-4. A winner-takes-all mechanism selects the single most relevant sub-network per diagnosis  
-5. Only the winning prototype contributes to the prediction and receives gradient updates  
-
-This allows S-Proto to model heterogeneous disease phenotypes while remaining sparse and efficient.
-
-## Intended Use
-
-This model is intended for:
-
-- Clinical diagnosis prediction from admission notes  
-- Research on long-tail learning in healthcare NLP  
-- Interpretable clinical decision support systems  
-- Analysis of disease phenotypes via learned prototypes  
-
-This model is **not intended for direct clinical deployment** without external validation, auditing, and regulatory approval.
-
-## Inference Example
-
-```python
-from transformers import AutoTokenizer, AutoModel
-import torch
-
-tokenizer = AutoTokenizer.from_pretrained("DATEXIS/sproto")
-model = AutoModel.from_pretrained("DATEXIS/sproto", trust_remote_code=True)
-model.eval()
-
-text_input = [
-    "CHIEF COMPLAINT: Right Carotid Artery Stenosis. "
-    "PRESENT ILLNESS: Ms. ___ is a ___ year old woman with hyperlipidemia, "
-    "cirrhosis with esophageal varices, alcoholism, COPD, left eye blindness, "
-    "and right carotid stenosis status post right carotid endarterectomy."
-]
-
-inputs = tokenizer(
-    text_input,
-    padding=True,
-    truncation=True,
-    max_length=512,
-    return_tensors="pt"
-)
-
-tokens = [tokenizer.convert_ids_to_tokens(ids) for ids in inputs["input_ids"]]
-
-with torch.no_grad():
-    output = model(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
-        token_type_ids=inputs.get("token_type_ids"),
-        tokens=tokens
-    )
-
-logits = output["logits"]
-max_indices = output["max_indices"]
-metadata = output["metadata"]
-
-print("Inference successful")
-print("Logits shape:", logits.shape)
-print("Max indices:", max_indices)
-print("Metadata:", metadata)
-```
-
-## Outputs
-
-The model returns a dictionary with the following entries:
-
-- **logits**  
-  Prediction scores per diagnosis label.
-
-- **max_indices**  
-  Index of the winning prototype sub-network per diagnosis, corresponding to the selected prototype.
-
-- **metadata**  
-  Additional information useful for analysis and interpretability.
-
-
-![Output Example](output_example.png)
-
-## Explainability
-
-S-Proto provides built-in faithful explanations through its prototypical structure:
-
-- Attention vectors highlight clinically relevant tokens  
-- Prototype distances reflect similarity to prototypical patients  
-- Multiple prototypes per diagnosis capture disease subtypes and cohorts  
-- Faithfulness metrics remain comparable to ProtoPatient despite higher capacity  
-
-Qualitative evaluation with medical professionals confirms that learned prototypes often correspond to clinically meaningful phenotypes.
-
-## Training
-
-First, clone the repository:
-
-```bash
-git clone https://github.com/DATEXIS/sproto.git
-cd sproto
-```
-
-Set up the environment using Poetry:
-
-```bash
-poetry install
-```
-
-Activate the virtual environment:
-
-```bash
-poetry env activate
-```
-
-Once the environment is active, you can start training by running the train command with the desired arguments.
-
-Example:
-
-```bash
-train \
-  --batch_size 3 \
-  --pretrained_model microsoft/biomednlp-pubmedbert-base-uncased-abstract-fulltext \
-  --pretrained_model_path path_to_pretrained_model.ckpt \
-  --model_type MULTI_PROTO \
-  --train_file training_data.csv \
-  --val_file validation_data.csv \
-  --test_file test_data.csv \
-  --save_dir ../experiments/ \
-  --gpus 1 \
-  --check_val_every_n_epoch 2 \
-  --num_warmup_steps 0 \
-  --num_training_steps 50 \
-  --max_length 512 \
-  --lr_features 0.000005 \
-  --lr_prototypes 0.001 \
-  --lr_others 0.001 \
-  --num_val_samples None \
-  --use_attention True \
-  --reduce_hidden_size 256 \
-  --all_labels_path all_labels.pcl \
-  --seed 42 \
-  --label_column labels \
-  --metric_opt auroc_macro \
-  --train_files [] \
-  --val_files [] \
-  --only_test True \
-  --model_name 5p \
-  --store_metadata False \
-  --num_prototypes_per_class 5
-```
-
-## Citation
-
-```bibtex
-@inproceedings{figueroa2024sproto,
-  title={Boosting Long-Tail Data Classification with Sparse Prototypical Networks},
-  author={Figueroa, Alexei and Papaioannou, Jens-Michalis and Fallon, Conor and Bekiaridou, Alexandra and Bressem, Keno and Zanos, Stavros and Gers, Felix and Nejdl, Wolfgang and Löser, Alexander},
-  booktitle={Proceedings of the European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML PKDD)},
-  year={2024}
-}
-```
-
-## License
-
-This model and its associated code are released under the Apache License 2.0.
-
-The model was trained on the MIMIC-IV dataset, which is subject to restricted access. No training data is included or redistributed with this repository.
-The data were accessed under a data use agreement. No patient-identifiable information is shared.
-
-Use of this model must comply with all applicable data governance and ethical guidelines.
-
-### Limitations
-
-- Extremely rare diagnoses remain challenging  
-- Clinical dataset biases may be reflected in predictions  
-- Winner-takes-all selection is fixed and not learned dynamically  
-- Not validated for real-world clinical deployment  
-
-### Ethical Considerations
-
-- The model processes sensitive clinical text  
-- Predictions should always be reviewed by qualified professionals  
-- Outputs should not be used as sole evidence for clinical decisions  
+---
+language: en
+license: apache-2.0
+library_name: transformers
+
+pipeline_tag: text-classification
+task_categories:
+  - text-classification
+
+model_type: sproto
+base_model: microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext
+
+datasets:
+  - mimic-iv
+
+metrics:
+  - auroc
+  - pr-auc
+
+tags:
+  - text-classification
+  - multi-label-classification
+  - long-tail-learning
+  - medical
+  - clinical-nlp
+  - interpretability
+  - prototypical-networks
+  - ehr
+---
+
+# S-Proto: Sparse Prototypical Networks for Long-Tail Clinical Diagnosis Prediction
+
+**Published at ECML PKDD 2024 (CORE A)**  
+*Boosting Long-Tail Data Classification with Sparse Prototypical Networks*
+
+Alexei Figueroa*, Jens-Michalis Papaioannou*, et al.  
+DATEXIS, Berliner Hochschule für Technik, Feinstein Institutes, TU Munich, Leibniz University Hannover  
+(* equal contribution)
+
+![S-Proto](overview.png)
+
+This repository provides **S-Proto**, a sparse and interpretable prototypical network for extreme multi-label diagnosis prediction from clinical text. The model is designed to address the long-tail distribution of clinical diagnoses while preserving faithful, prototype-based explanations.
+
+## Interactive Demo
+
+You can explore the model's predictions and interpretability features through our interactive web demo:
+**[https://s-proto.demo.datexis.com/](https://s-proto.demo.datexis.com/)**
+
+S-Proto was introduced in the paper:
+
+**[Boosting Long-Tail Data Classification with Sparse Prototypical Networks](https://ecmlpkdd-storage.s3.eu-central-1.amazonaws.com/preprints/2024/lncs14947/lncs14947435.pdf)**
+European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML PKDD 2024, CORE A)  
+Alexei Figueroa*, Jens-Michalis Papaioannou*, et al.  
+DATEXIS, Berliner Hochschule für Technik, Feinstein Institutes, TU Munich, Leibniz University Hannover  
+(* equal contribution)
+
+## Overview
+
+Clinical outcome prediction from Electronic Health Records is characterized by extreme label imbalance. A small number of diagnoses account for most patients, while the majority of diagnoses appear rarely. Standard transformer classifiers tend to perform well on frequent diagnoses but degrade sharply in the long tail.
+
+S-Proto addresses this problem by extending prototypical networks with:
+
+- Multiple prototypes per diagnosis  
+- Sparse winner-takes-all activation  
+- Prototype-level interpretability  
+- Efficient training despite increased representational capacity  
+
+The model achieves state-of-the-art performance on MIMIC-IV diagnosis prediction, with particularly strong gains in PR-AUC for rare diagnoses, and transfers successfully to unseen clinical datasets.
+
+## Model Architecture
+
+S-Proto builds on **PubMedBERT** as the text encoder and introduces a sparse prototypical layer on top.
+
+For each diagnosis label, the model learns multiple sub-networks, each consisting of:
+
+- A label-specific attention vector  
+- A prototype vector representing a prototypical patient  
+
+Given an input clinical note:
+
+1. The note is encoded using PubMedBERT  
+2. Token embeddings are projected into a latent space  
+3. Each diagnosis activates multiple candidate sub-networks  
+4. A winner-takes-all mechanism selects the single most relevant sub-network per diagnosis  
+5. Only the winning prototype contributes to the prediction and receives gradient updates  
+
+This allows S-Proto to model heterogeneous disease phenotypes while remaining sparse and efficient.
+
+## Intended Use
+
+This model is intended for:
+
+- Clinical diagnosis prediction from admission notes  
+- Research on long-tail learning in healthcare NLP  
+- Interpretable clinical decision support systems  
+- Analysis of disease phenotypes via learned prototypes  
+
+This model is **not intended for direct clinical deployment** without external validation, auditing, and regulatory approval.
+
+## Requirements
+
+The model depends on the base `sproto` package (which contains `MultiProtoModule`) and specific versions of its dependencies. Version mismatches — especially in `torchmetrics` and `pytorch-lightning` — will cause `AttributeError` or import failures.
+
+```bash
+pip install torch>=1.12.1 \
+            transformers>=4.25.1 \
+            torchmetrics>=0.10.1 \
+            pytorch-lightning==1.9
+```
+
+| Package | Required version | Reason |
+|---------|-----------------|--------|
+| `torch` | `>= 1.12.1` | Minimum version for `nn.PairwiseDistance` and `torch.einsum` patterns used in the prototype layer |
+| `transformers` | `>= 4.25.1` | Minimum version with `trust_remote_code` + `auto_map` support for custom model loading |
+| `torchmetrics` | `>= 0.10.1` | `MultilabelAveragePrecision` was added in 0.10; older versions raise `AttributeError` on load |
+| `pytorch-lightning` | `== 1.9` | `MultiProtoModule` is a `pl.LightningModule`; the exact API (e.g. `validation_epoch_end`) changed in 2.x |
+| `sproto` | bundled | The `sproto/` package is included in this HF repo and downloaded automatically with `trust_remote_code=True` — no separate install needed |
+
+## Inference Example
+
+```python
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained("DATEXIS/sproto")
+model = AutoModel.from_pretrained("DATEXIS/sproto", trust_remote_code=True)
+model.eval()
+
+text_input = [
+    "CHIEF COMPLAINT: Right Carotid Artery Stenosis. "
+    "PRESENT ILLNESS: Ms. ___ is a ___ year old woman with hyperlipidemia, "
+    "cirrhosis with esophageal varices, alcoholism, COPD, left eye blindness, "
+    "and right carotid stenosis status post right carotid endarterectomy."
+]
+
+inputs = tokenizer(
+    text_input,
+    padding=True,
+    truncation=True,
+    max_length=512,
+    return_tensors="pt"
+)
+
+tokens = [tokenizer.convert_ids_to_tokens(ids) for ids in inputs["input_ids"]]
+
+with torch.no_grad():
+    output = model(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        token_type_ids=inputs.get("token_type_ids"),
+        tokens=tokens
+    )
+
+logits = output["logits"]
+max_indices = output["max_indices"]
+metadata = output["metadata"]
+
+print("Inference successful")
+print("Logits shape:", logits.shape)
+print("Max indices:", max_indices)
+print("Metadata:", metadata)
+```
+
+> **Note:** `tokens` (the list of token strings per sample) is **required** when `use_attention=True`
+> (which is the default). The attention mechanism uses the actual token strings to mask clinical
+> section headers (`[CLS]`, `[SEP]`, `"chief complaint :"`, etc.) before computing
+> token-to-prototype attention. Omitting `tokens` will raise a `ValueError`.
+> Obtain them with `tokenizer.convert_ids_to_tokens(input_ids[i])` as shown above.
+
+## Outputs
+
+The model returns a dictionary with the following entries:
+
+- **logits**  
+  Prediction scores per diagnosis label.
+
+- **max_indices**  
+  Index of the winning prototype sub-network per diagnosis, corresponding to the selected prototype.
+
+- **metadata**  
+  Additional information useful for analysis and interpretability.
+
+
+![Output Example](output_example.png)
+
+## Explainability
+
+S-Proto provides built-in faithful explanations through its prototypical structure:
+
+- Attention vectors highlight clinically relevant tokens  
+- Prototype distances reflect similarity to prototypical patients  
+- Multiple prototypes per diagnosis capture disease subtypes and cohorts  
+- Faithfulness metrics remain comparable to ProtoPatient despite higher capacity  
+
+Qualitative evaluation with medical professionals confirms that learned prototypes often correspond to clinically meaningful phenotypes.
+
+## Training
+
+First, clone the repository:
+
+```bash
+git clone https://github.com/DATEXIS/sproto.git
+cd sproto
+```
+
+Set up the environment using Poetry:
+
+```bash
+poetry install
+```
+
+Activate the virtual environment:
+
+```bash
+poetry env activate
+```
+
+Once the environment is active, you can start training by running the train command with the desired arguments.
+
+Example:
+
+```bash
+train \
+  --batch_size 3 \
+  --pretrained_model microsoft/biomednlp-pubmedbert-base-uncased-abstract-fulltext \
+  --pretrained_model_path path_to_pretrained_model.ckpt \
+  --model_type MULTI_PROTO \
+  --train_file training_data.csv \
+  --val_file validation_data.csv \
+  --test_file test_data.csv \
+  --save_dir ../experiments/ \
+  --gpus 1 \
+  --check_val_every_n_epoch 2 \
+  --num_warmup_steps 0 \
+  --num_training_steps 50 \
+  --max_length 512 \
+  --lr_features 0.000005 \
+  --lr_prototypes 0.001 \
+  --lr_others 0.001 \
+  --num_val_samples None \
+  --use_attention True \
+  --reduce_hidden_size 256 \
+  --all_labels_path all_labels.pcl \
+  --seed 42 \
+  --label_column labels \
+  --metric_opt auroc_macro \
+  --train_files [] \
+  --val_files [] \
+  --only_test True \
+  --model_name 5p \
+  --store_metadata False \
+  --num_prototypes_per_class 5
+```
+
+## Citation
+
+```bibtex
+@inproceedings{figueroa2024sproto,
+  title={Boosting Long-Tail Data Classification with Sparse Prototypical Networks},
+  author={Figueroa, Alexei and Papaioannou, Jens-Michalis and Fallon, Conor and Bekiaridou, Alexandra and Bressem, Keno and Zanos, Stavros and Gers, Felix and Nejdl, Wolfgang and Löser, Alexander},
+  booktitle={Proceedings of the European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML PKDD)},
+  year={2024}
+}
+```
+
+## License
+
+This model and its associated code are released under the Apache License 2.0.
+
+The model was trained on the MIMIC-IV dataset, which is subject to restricted access. No training data is included or redistributed with this repository.
+The data were accessed under a data use agreement. No patient-identifiable information is shared.
+
+Use of this model must comply with all applicable data governance and ethical guidelines.
+
+### Limitations
+
+- Extremely rare diagnoses remain challenging  
+- Clinical dataset biases may be reflected in predictions  
+- Winner-takes-all selection is fixed and not learned dynamically  
+- Not validated for real-world clinical deployment  
+
+### Ethical Considerations
+
+- The model processes sensitive clinical text  
+- Predictions should always be reviewed by qualified professionals  
+- Outputs should not be used as sole evidence for clinical decisions  
 - Care must be taken to avoid reinforcing existing healthcare biases

__init__.py

@@ -0,0 +1,8 @@
+from .configuration_sproto import SprotoConfig
+from .modeling_sproto import SprotoModel, SprotoOutput
+from transformers import AutoConfig, AutoModel
+
+AutoConfig.register("sproto", SprotoConfig)
+AutoModel.register(SprotoConfig, SprotoModel)
+
+__all__ = ["SprotoConfig", "SprotoModel", "SprotoOutput"]

base_bert_config.json

@@ -0,0 +1,17 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "model_type": "bert",
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "max_position_embeddings": 512,
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "type_vocab_size": 2,
+  "vocab_size": 30522
+}

config.json

@@ -1,33 +1,27 @@
-{
-  "attention_vector_path": null,
-  "auto_map": {
-    "AutoConfig": "configuration_sproto.SprotoConfig",
-    "AutoModel": "modeling_sproto.SprotoModel"
-  },
-  "batch_size": 21,
-  "dot_product": false,
-  "eval_buckets": null,
-  "final_layer": false,
-  "label_order_path": "/pvc/shared/continual/data/icd_10_all_labels_admission_mimiciv_dia.pcl",
-  "loss": "BCE",
-  "lr_features": 5e-06,
-  "lr_others": 0.001,
-  "lr_prototypes": 0.001,
-  "model_type": "sproto",
-  "normalize": null,
-  "num_classes": 1643,
-  "num_prototypes_per_class": 5,
-  "num_training_steps": 5000,
-  "num_warmup_steps": 5000,
-  "pretrained_model": "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
-  "prototype_vector_path": null,
-  "reduce_hidden_size": 256,
-  "save_dir": "/pvc/shared/continual/experiments/mimiciv/icd10_clinical-continual-5p-test",
-  "seed": 28,
-  "transformers_version": "4.25.1",
-  "use_attention": true,
-  "use_cuda": true,
-  "use_global_attention": false,
-  "use_prototype_loss": false,
-  "use_sigmoid": false
-}
+{
+  "auto_map": {
+    "AutoConfig": "configuration_sproto.SprotoConfig",
+    "AutoModel": "modeling_sproto.SprotoModel"
+  },
+  "attention_probs_dropout_prob": 0.1,
+  "dot_product": false,
+  "final_layer": false,
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "label_order_path": null,
+  "loss": "BCE",
+  "max_position_embeddings": 512,
+  "model_type": "sproto",
+  "normalize": null,
+  "num_classes": 1643,
+  "num_prototypes_per_class": 5,
+  "pretrained_model": "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
+  "reduce_hidden_size": 256,
+  "seed": 7,
+  "transformers_version": "4.25.1",
+  "use_attention": true,
+  "use_global_attention": false,
+  "use_prototype_loss": false,
+  "use_sigmoid": false,
+  "vocab_size": 28996
+}

configuration_sproto.py

@@ -1,61 +1,62 @@
-from transformers.configuration_utils import PretrainedConfig
-
-class SprotoConfig(PretrainedConfig):
-    model_type = "sproto"
-
-    def __init__(
-        self,
-        pretrained_model=None,
-        num_classes=None,
-        label_order_path=None,
-        use_sigmoid=False,
-        use_cuda=True,
-        lr_prototypes=5e-2,
-        lr_features=2e-6,
-        lr_others=2e-2,
-        num_training_steps=5000,
-        num_warmup_steps=1000,
-        loss="BCE",
-        save_dir="output",
-        use_attention=True,
-        use_global_attention=False,
-        dot_product=False,
-        normalize=None,
-        final_layer=False,
-        reduce_hidden_size=None,
-        use_prototype_loss=False,
-        prototype_vector_path=None,
-        attention_vector_path=None,
-        eval_buckets=None,
-        seed=7,
-        num_prototypes_per_class=1,
-        batch_size=10,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.pretrained_model = pretrained_model
-        self.num_classes = num_classes
-        self.label_order_path = label_order_path
-        self.use_sigmoid = use_sigmoid
-        self.use_cuda = use_cuda
-        self.lr_prototypes = lr_prototypes
-        self.lr_features = lr_features
-        self.lr_others = lr_others
-        self.num_training_steps = num_training_steps
-        self.num_warmup_steps = num_warmup_steps
-        self.loss = loss
-        self.save_dir = save_dir
-        self.use_attention = use_attention
-        self.use_global_attention = use_global_attention
-        self.dot_product = dot_product
-        self.normalize = normalize
-        self.final_layer = final_layer
-        self.reduce_hidden_size = reduce_hidden_size
-        self.use_prototype_loss = use_prototype_loss
-        self.prototype_vector_path = prototype_vector_path
-        self.attention_vector_path = attention_vector_path
-        self.eval_buckets = eval_buckets
-        self.seed = seed
-        self.num_prototypes_per_class = num_prototypes_per_class
-        self.batch_size = batch_size
+from transformers import PretrainedConfig
+
+
+class SprotoConfig(PretrainedConfig):
+    model_type = "sproto"
+
+    def __init__(
+        self,
+        pretrained_model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
+        num_classes=55,
+        label_order_path=None,
+        use_attention=True,
+        use_global_attention=False,
+        dot_product=False,
+        normalize=None,
+        final_layer=False,
+        reduce_hidden_size=None,
+        num_prototypes_per_class=1,
+        loss="BCE",
+        use_prototype_loss=False,
+        use_sigmoid=False,
+        seed=7,
+        vocab_size=28996,
+        hidden_size=768,
+        max_position_embeddings=512,
+        attention_probs_dropout_prob=0.1,
+        hidden_dropout_prob=0.1,
+        **kwargs,
+    ):
+        # Bug fix: the checkpoint serialises the num_prototypes_per_class buffer as a list
+        # (one float per class). MultiProtoModule expects a scalar int for the uniform case.
+        # Collapse the list → scalar, asserting all values are identical.
+        if isinstance(num_prototypes_per_class, (list, tuple)):
+            unique_vals = set(int(v) for v in num_prototypes_per_class)
+            assert len(unique_vals) == 1, (
+                "Non-uniform num_prototypes_per_class cannot be represented as a scalar "
+                "config value. Got: {}".format(unique_vals)
+            )
+            num_prototypes_per_class = unique_vals.pop()
+
+        self.pretrained_model = pretrained_model
+        self.num_classes = num_classes
+        self.label_order_path = label_order_path
+        self.use_attention = use_attention
+        self.use_global_attention = use_global_attention
+        self.dot_product = dot_product
+        self.normalize = normalize
+        self.final_layer = final_layer
+        self.reduce_hidden_size = reduce_hidden_size
+        self.num_prototypes_per_class = num_prototypes_per_class
+        self.loss = loss
+        self.use_prototype_loss = use_prototype_loss
+        self.use_sigmoid = use_sigmoid
+        self.seed = seed
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.max_position_embeddings = max_position_embeddings
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.hidden_dropout_prob = hidden_dropout_prob
+
+        super().__init__(**kwargs)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ef1e86215368bfcbb723cb3a28d2c343927f154decc09527cce2093326a07fd2
-size 455575332
+oid sha256:62960c454029f57c3a2dbfde86e34ad0ae11a6860b3233de8cbb3577a22b7d86
+size 455576732

modeling_sproto.py

@@ -1,77 +1,127 @@
-from transformers import PreTrainedModel
-from sproto.model.multi_proto import MultiProtoModule
-from .configuration_sproto import SprotoConfig
-
-class SprotoModel(PreTrainedModel):
-    config_class = SprotoConfig
-    base_model_prefix = "sproto"
-
-    def __init__(self, config: SprotoConfig):
-        super().__init__(config)
-
-        self.module = MultiProtoModule(
-            pretrained_model=config.pretrained_model,
-            num_classes=config.num_classes,
-            label_order_path=config.label_order_path,
-            use_sigmoid=config.use_sigmoid,
-            use_cuda=config.use_cuda,
-            lr_prototypes=config.lr_prototypes,
-            lr_features=config.lr_features,
-            lr_others=config.lr_others,
-            num_training_steps=config.num_training_steps,
-            num_warmup_steps=config.num_warmup_steps,
-            loss=config.loss,
-            save_dir=config.save_dir,
-            use_attention=config.use_attention,
-            use_global_attention=config.use_global_attention,
-            dot_product=config.dot_product,
-            normalize=config.normalize,
-            final_layer=config.final_layer,
-            reduce_hidden_size=config.reduce_hidden_size,
-            use_prototype_loss=config.use_prototype_loss,
-            prototype_vector_path=config.prototype_vector_path,
-            attention_vector_path=config.attention_vector_path,
-            eval_buckets=config.eval_buckets,
-            seed=config.seed,
-            num_prototypes_per_class=config.num_prototypes_per_class,
-            batch_size=config.batch_size,
-        )
-        
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        if isinstance(module, (MultiProtoModule)):
-            # MultiProtoModule handles its own initialization or is loaded from checkpoint
-            return
-        # Add other initializations if standard layers are used directly in SprotoModel
-        pass
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        token_type_ids=None,
-        targets=None,
-        tokens=None,
-        sample_ids=None,
-        **kwargs,
-    ):
-
-        batch = {
-            "input_ids": input_ids,
-            "attention_masks": attention_mask,
-            "token_type_ids": token_type_ids,
-            "targets": targets,
-            "tokens": tokens,
-            "sample_ids": sample_ids,
-        }
-
-        logits, max_indices, metadata = self.module(batch)
-
-        return {
-            "logits": logits,
-            "max_indices": max_indices,
-            "metadata": metadata,
-        }
+import torch
+from dataclasses import dataclass
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from sproto.model.multi_proto import MultiProtoModule
+
+from .configuration_sproto import SprotoConfig
+
+
+@dataclass
+class SprotoOutput(ModelOutput):
+    logits: torch.Tensor = None
+    max_indices: torch.Tensor = None
+    metadata: dict = None
+
+    def __contains__(self, key):
+        return hasattr(self, key)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+
+class SprotoModel(PreTrainedModel):
+    config_class = SprotoConfig
+    base_model_prefix = "sproto"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # HF's from_pretrained (on newer PyTorch / Transformers) can set a meta-device
+        # context during weight loading. MultiProtoModule.__init__ internally calls
+        # AutoModel.from_pretrained() for the BERT backbone, which conflicts with that
+        # context. We escape it by forcing CPU initialization when on torch >= 2.0
+        # (torch.device context manager was added in 2.0); on older torch the conflict
+        # does not occur and we construct normally.
+        _torch_version = tuple(int(x) for x in torch.__version__.split(".")[:2] if x.isdigit())
+        if _torch_version >= (2, 0):
+            with torch.device("cpu"):
+                self.module = MultiProtoModule(
+                    pretrained_model=config.pretrained_model,
+                    num_classes=config.num_classes,
+                    label_order_path=config.label_order_path,
+                    use_attention=config.use_attention,
+                    use_global_attention=config.use_global_attention,
+                    dot_product=config.dot_product,
+                    normalize=config.normalize,
+                    final_layer=config.final_layer,
+                    reduce_hidden_size=config.reduce_hidden_size,
+                    num_prototypes_per_class=config.num_prototypes_per_class,
+                    loss=config.loss,
+                    use_prototype_loss=config.use_prototype_loss,
+                    use_sigmoid=config.use_sigmoid,
+                    seed=config.seed,
+                    # use_cuda=False: HF handles device placement via .to(device);
+                    # manual .cuda() calls in the base model are bypassed here.
+                    use_cuda=False,
+                )
+        else:
+            self.module = MultiProtoModule(
+                pretrained_model=config.pretrained_model,
+                num_classes=config.num_classes,
+                label_order_path=config.label_order_path,
+                use_attention=config.use_attention,
+                use_global_attention=config.use_global_attention,
+                dot_product=config.dot_product,
+                normalize=config.normalize,
+                final_layer=config.final_layer,
+                reduce_hidden_size=config.reduce_hidden_size,
+                num_prototypes_per_class=config.num_prototypes_per_class,
+                loss=config.loss,
+                use_prototype_loss=config.use_prototype_loss,
+                use_sigmoid=config.use_sigmoid,
+                seed=config.seed,
+                use_cuda=False,
+            )
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        token_type_ids=None,
+        targets=None,
+        tokens=None,
+        sample_ids=None,
+    ):
+        # tokens MUST be provided when use_attention=True.
+        # attention_mask_from_tokens() relies on real token strings to zero out clinical
+        # section headers ([CLS], [SEP], "chief complaint :", etc.) before computing
+        # token-prototype attention. A fake default would silently produce wrong logits.
+        if tokens is None and self.config.use_attention:
+            raise ValueError(
+                "tokens (list-of-lists of token strings per sample) must be provided "
+                "when use_attention=True. Obtain them via:\n"
+                "  tokenizer.convert_ids_to_tokens(input_ids[i])\n"
+                "for each sample i in the batch, or pass the full batch at once with:\n"
+                "  [tokenizer.convert_ids_to_tokens(ids) for ids in input_ids]"
+            )
+
+        batch = {
+            "input_ids": input_ids,
+            "attention_masks": attention_mask,
+            "token_type_ids": token_type_ids if token_type_ids is not None else torch.zeros_like(input_ids),
+            "targets": targets if targets is not None else torch.zeros(input_ids.shape[0], self.config.num_classes),
+            "tokens": tokens if tokens is not None else [["[PAD]"] * input_ids.shape[1]] * input_ids.shape[0],
+            "sample_ids": sample_ids if sample_ids is not None else [f"sample_{i}" for i in range(input_ids.shape[0])],
+        }
+
+        logits, max_indices, metadata = self.module(batch)
+
+        return SprotoOutput(
+            logits=logits,
+            max_indices=max_indices,
+            metadata=metadata,
+        )
+
+    def get_embeddings(self, input_ids, attention_mask, token_type_ids=None):
+        bert_output = self.module.bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids if token_type_ids is not None else torch.zeros_like(input_ids),
+        )
+        return bert_output.last_hidden_state
+
+    def _init_weights(self, module):
+        pass

pytorch_model.bin

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

tokenizer_config.json

@@ -1,16 +1,15 @@
-{
-  "cls_token": "[CLS]",
-  "do_basic_tokenize": true,
-  "do_lower_case": true,
-  "mask_token": "[MASK]",
-  "model_max_length": 1000000000000000019884624838656,
-  "name_or_path": "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
-  "never_split": null,
-  "pad_token": "[PAD]",
-  "sep_token": "[SEP]",
-  "special_tokens_map_file": null,
-  "strip_accents": null,
-  "tokenize_chinese_chars": true,
-  "tokenizer_class": "BertTokenizer",
-  "unk_token": "[UNK]"
-}
+{
+  "do_lower_case": true,
+  "tokenizer_class": "BertTokenizer",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]",
+  "pad_token_id": 0,
+  "sep_token_id": 102,
+  "cls_token_id": 101,
+  "mask_token_id": 103,
+  "unk_token_id": 100,
+  "model_max_length": 512
+}