Update README.md

README.md

@@ -1,49 +1,51 @@
 # BioHiCL-base: Hierarchical Multi-Label Contrastive Biomedical Retriever
 
 ## 🔍 Overview
-BioHiCL-base is a biomedical dense retriever trained with hierarchical MeSH supervision to capture fine-grained semantic relationships between biomedical texts. 
+BioHiCL-base is a biomedical dense retriever trained with hierarchical MeSH supervision to capture fine-grained semantic relationships between biomedical texts.
 
-Unlike traditional dense retrievers trained with binary relevance signals, BioHiCL models semantic similarity using structured multi-label supervision derived from the MeSH ontology.
+Unlike traditional dense retrievers trained with binary relevance signals, BioHiCL models semantic similarity using structured multi-label supervision derived from the MeSH ontology, enabling it to capture partial semantic overlap between documents.
 
 ---
 
 ## 💡 Key Features
-- **Hierarchical supervision**: Uses MeSH ontology to model semantic relationships
-- **Multi-label similarity learning**: Captures partial semantic overlap between documents
-- **Contrastive + regression training**: Aligns embedding similarity with label similarity
-- **Efficient**: ~0.1B parameters, suitable for deployment on a single GPU
+- **Hierarchical supervision**: Leverages MeSH ontology to encode structured biomedical semantics  
+- **Multi-label similarity learning**: Captures graded semantic overlap beyond binary relevance  
+- **Contrastive + regression training**: Aligns embedding similarity with label similarity  
+- **Efficient**: ~0.1B parameters, suitable for deployment on a single GPU  
 
 ---
 
 ## 🧠 Model Details
-- **Model type**: Bi-encoder (dense retriever)
-- **Backbone**: BAAI/bge-base-en-v1.5
-- **Parameters**: ~0.1B
-- **Fine-tuning**: LoRA (merged into base model)
-- **Max input length**: 8192 tokens
+- **Model type**: Bi-encoder (dense retriever)  
+- **Backbone**: `BAAI/bge-base-en-v1.5`  
+- **Parameters**: ~0.1B  
+- **Fine-tuning**: LoRA (merged into base model)  
+- **Max input length**: 8192 tokens  
 
 ---
 
 ## ⚙️ How It Works
 BioHiCL aligns:
-- **Embedding similarity (SimE)**: cosine similarity between embeddings
-- **Label similarity (SimL)**: cosine similarity over weighted MeSH labels
+- **Embedding similarity (SimE)**: cosine similarity between embeddings  
+- **Label similarity (SimL)**: cosine similarity over weighted MeSH label vectors  
 
-Training objective:
-- MSE loss to align SimE with SimL
-- Hierarchical contrastive loss to separate unrelated documents
+### Training Objective
+- Mean Squared Error (MSE) loss to align SimE with SimL  
+- Hierarchical contrastive loss to separate unrelated documents and prevent embedding collapse  
 
 ---
 
-## 🚀 Usage - Text Similarity
+## 🚀 Usage — Text Similarity
 
 ```python
 from transformers import AutoTokenizer, AutoModel
 import torch
 import torch.nn.functional as F
 
-tokenizer = AutoTokenizer.from_pretrained("LunaLan07/BioHiCL-Large")
-model = AutoModel.from_pretrained("LunaLan07/BioHiCL-Large")
+model_name = "LunaLan07/BioHiCL-base"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
 
 def encode(texts):
     inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
@@ -60,41 +62,25 @@ print(similarity)
 
 
 
----
-
-## 🚀 Usage - Evaluation on BEIR Benchmark
-
 ```python
-from beir import util
-from beir.datasets.data_loader import GenericDataLoader
-from beir.retrieval.models import SentenceBERT
-from beir.retrieval.search.dense import DenseRetrievalExactSearch
-from beir.retrieval.evaluation import EvaluateRetrieval
-
-dataset = "scifact"
-url = ...
-data_path = util.download_and_unzip(url, "datasets")
-corpus, queries, qrels = GenericDataLoader(data_path).load(split="test")
-
-model_name = "LunaLan07/BioHiCL-Large"
-model = SentenceBERT(model_name)
-retriever = DenseRetrievalExactSearch(model, batch_size=16)
-top_k = 10  # top 10 documents per query
-results = retriever.search(corpus, queries, top_k=top_k, score_function="cos_sim")
-
-k_values = [1, 3, 5, 10]
-ndcg, _map, recall, precision = EvaluateRetrieval.evaluate(qrels, results, k_values=k_values)
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn.functional as F
 
----
+model_name = "LunaLan07/BioHiCL-base"
 
-## 📖 Citation
-If you use this model, please cite:
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
 
-```bibtex
-@article{lan2026biohicl,
-  title={BioHiCL: Hierarchical Multi-Label Contrastive Learning for Biomedical Retrieval with MeSH Labels},
-  author={Lan, Mengfei, Zheng, Lecheng, and Kilicoglu, Halil},
-  journal={ACL 2026},
-  year={2026}
-}
+def encode(texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
+    outputs = model(**inputs)
+    embeddings = outputs.last_hidden_state[:, 0]  # CLS token
+    return F.normalize(embeddings, p=2, dim=1)
 
+# Example
+query = encode(["What are treatments for COPD?"])
+doc = encode(["Chronic obstructive pulmonary disease is treated with bronchodilators."])
+
+similarity = (query @ doc.T).item()
+print(similarity)