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-MiniLM-L12-v2
 ---
 
 # Hierarchy-Transformers/HiT-MiniLM-L12-WordNetNoun
@@ -20,16 +26,25 @@ A **Hi**erarchy **T**ransformer Encoder (HiT) model that explicitly encodes enti
 
 <!-- Provide a longer summary of what this model is. -->
 
-HiT-MiniLM-L12-WordNet is a HiT model trained on WordNet's noun hierarchy with random negative sampling.
+HiT-MiniLM-L12-WordNet 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)
+- **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-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2)
 - **Training Objectives**: Jointly optimised on *hyperbolic clustering* and *hyperbolic centripetal* losses
 
+### Model Versions
+
+
+| **Version** | **Model Revision** | **Note** |
+|------------|---------|----------|
+|v1.0 (Random Negatives)| `main` or `v1-random-negative`| The variant trained on random negatives, as detailed in the [paper](https://arxiv.org/abs/2401.11374).|
+|v1.0 (Hard Negatives)| `v1-hard-negative` | The variant trained on hard negatives, as detailed in the [paper](https://arxiv.org/abs/2401.11374). |
+
+
 ### Model Sources
 
 <!-- Provide the basic links for the model. -->
@@ -58,7 +73,12 @@ gpu_id = 0
 device = get_torch_device(gpu_id)
 
 # load the model
-model = HierarchyTransformer.load_pretrained('Hierarchy-Transformers/HiT-MiniLM-L12-WordNetNoun', device)
+revision = "main"  # change for a different version
+model = HierarchyTransformer.from_pretrained(
+  model_name_or_path='Hierarchy-Transformers/HiT-MiniLM-L12-WordNetNoun',
+  revision=revision
+  device=device
+)
 
 # entity names to be encoded.
 entity_names = ["computer", "personal computer", "fruit", "berry"]