Update README.md

README.md

@@ -35,52 +35,30 @@ BioHiCL aligns:
 
 ---
 
-## 🚀 Usage — Text Similarity
+## 🚀 Usage — Evaluation on BEIR Benchmark
 
 ```python
-from transformers import AutoTokenizer, AutoModel
-import torch
-import torch.nn.functional as F
+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
 
-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")
-    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)
+dataset = "scifact"
+url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip"
 
-
-
-```python
-from transformers import AutoTokenizer, AutoModel
-import torch
-import torch.nn.functional as F
+data_path = util.download_and_unzip(url, "datasets")
+corpus, queries, qrels = GenericDataLoader(data_path).load(split="test")
 
 model_name = "LunaLan07/BioHiCL-base"
+model = SentenceBERT(model_name)
+
+retriever = DenseRetrievalExactSearch(model, batch_size=16)
+results = retriever.search(corpus, queries, top_k=10, score_function="cos_sim")
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+ndcg, _map, recall, precision = EvaluateRetrieval.evaluate(
+    qrels, results, k_values=[1, 3, 5, 10]
+)
 
-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)