Update README.md

README.md

@@ -10,6 +10,12 @@ tags:
 license: apache-2.0
 language:
 - en
+metrics:
+- precision
+- recall
+- f1
+base_model:
+- sentence-transformers/all-mpnet-base-v2
 ---
 
 # Hierarchy-Transformers/HiT-MPNet-WordNetNoun
@@ -20,15 +26,23 @@ A **Hi**erarchy **T**ransformer Encoder (HiT) model that explicitly encodes enti
 
 <!-- Provide a longer summary of what this model is. -->
 
-HiT-MPNet-WordNetNoun is a HiT model trained on WordNet's noun hierarchy with random negative sampling.
+HiT-MPNet-WordNetNoun is a HiT model trained on WordNet's subsumption (hypernym) hierarchy of noun entities.
 
 - **Developed by:** [Yuan He](https://www.yuanhe.wiki/), Zhangdie Yuan, Jiaoyan Chen, and Ian Horrocks
 - **Model type:** Hierarchy Transformer Encoder (HiT)
 - **License:** Apache license 2.0
-- **Hierarchy**: WordNet (Noun)
-- **Training Dataset**: Download `wordnet.zip` from [Datasets for HiTs on Zenodo](https://zenodo.org/doi/10.5281/zenodo.10511042)
+- **Hierarchy**: WordNet's subsumption (hypernym) hierarchy of noun entities.
+- **Training Dataset**: Download `wordnet-mixed.zip` from [Datasets for HiTs on Zenodo](https://zenodo.org/doi/10.5281/zenodo.10511042)
 - **Pre-trained model:** [sentence-transformers/all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2)
-- **Training Objectives**: Jointly optimised on *hyperbolic clustering* and *hyperbolic centripetal* losses
+- **Training Objectives**: Jointly optimised on *Hyperbolic Clustering* and *Hyperbolic Centripetal* losses (see definitions in the [paper](https://arxiv.org/abs/2401.11374))
+
+### Model Versions
+
+| **Version** | **Model Revision** | **Note** |
+|------------|---------|----------|
+|v1.0 (Random Negatives)| `main` or `v1-random-negatives`| The variant trained on random negatives, as detailed in the [paper](https://arxiv.org/abs/2401.11374).|
+|v1.0 (Hard Negatives)| `v1-hard-negatives` | The variant trained on hard negatives, as detailed in the [paper](https://arxiv.org/abs/2401.11374). |
+
 
 ### Model Sources
 
@@ -58,7 +72,12 @@ gpu_id = 0
 device = get_torch_device(gpu_id)
 
 # load the model
-model = HierarchyTransformer.load_pretrained('Hierarchy-Transformers/HiT-MPNet-WordNetNoun', device)
+revision = "main"  # change for a different version
+model = HierarchyTransformer.from_pretrained(
+  model_name_or_path='Hierarchy-Transformers/HiT-MPNet-WordNetNoun',
+  revision=revision
+  device=device
+)
 
 # entity names to be encoded.
 entity_names = ["computer", "personal computer", "fruit", "berry"]
@@ -86,7 +105,8 @@ parent_norms = model.manifold.dist0(parent_entity_embeddings)
 subsumption_scores = - (dists + centri_score_weight * (parent_norms - child_norms))
 ```
 
-Training and evaluation scripts are available at [GitHub](https://github.com/KRR-Oxford/HierarchyTransformers).
+Training and evaluation scripts are available at [GitHub](https://github.com/KRR-Oxford/HierarchyTransformers/tree/main/scripts). See `scripts/evaluate.py` for how we determine the hyperparameters on the validation set for subsumption prediction.
+
 Technical details are presented in the [paper](https://arxiv.org/abs/2401.11374).
 
 
@@ -105,7 +125,7 @@ HierarchyTransformer(
 
 Preprint on arxiv: https://arxiv.org/abs/2401.11374.
 
-*Yuan He, Zhangdie Yuan, Jiaoyan Chen, Ian Horrocks.* **Language Models as Hierarchy Encoders.** arXiv preprint arXiv:2401.11374 (2024).
+*Yuan He, Zhangdie Yuan, Jiaoyan Chen, Ian Horrocks.* **Language Models as Hierarchy Encoders.** To Appear at NeurIPS 2024.
 
 ```
 @article{he2024language,
@@ -119,4 +139,4 @@ Preprint on arxiv: https://arxiv.org/abs/2401.11374.
 
 ## Model Card Contact
 
-For any queries or feedback, please contact Yuan He (yuan.he@cs.ox.ac.uk).
+For any queries or feedback, please contact Yuan He (`yuan.he(at)cs.ox.ac.uk`).