resolve conflicts

main.py

@@ -1,71 +1,7 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-from jsonargparse import ActionConfigFile, ArgumentParser
-from loguru import logger
-
-from src.embedders import TextEmbedder
-from src.evaluators import EmbeddingEvaluator
-from src.utils.score_recorder import JsonScoreRecorder
-
-
-def main(
-    text_embedder: TextEmbedder,
-    evaluators: dict[str, EmbeddingEvaluator],
-    save_dir: str | None = None,
-    overwrite_cache: bool = False,
-):
-    logger.info(f"Start evaluating the following tasks\n{list(evaluators.keys())}")
-
-    if save_dir:
-        Path(save_dir).mkdir(parents=True, exist_ok=True)
-
-    score_recorder = JsonScoreRecorder(save_dir)
-
-    for eval_name, evaluator in evaluators.items():
-        logger.info(f"Evaluating {eval_name}")
-
-        cache_dir = None
-        if save_dir is not None:
-            cache_dir = Path(save_dir) / "cache" / eval_name
-
-        metrics = evaluator(text_embedder, cache_dir=cache_dir, overwrite_cache=overwrite_cache)
-        score_recorder.record_task_scores(
-            scores=metrics,
-            dataset_name=eval_name,
-            task_name=evaluator.__class__.__name__.replace("Evaluator", ""),
-        )
-
-        logger.info(f"Results for {eval_name}\n{json.dumps(metrics.as_dict(), indent=4, ensure_ascii=False)}")
-
-    logger.info(f"Saving result summary to {Path(save_dir) / 'summary.json'}")
-    score_recorder.record_summary()
-
+from jmteb.main import main, parse_args
 
 if __name__ == "__main__":
-    parser = ArgumentParser(parser_mode="jsonnet")
-
-    parser.add_subclass_arguments(TextEmbedder, nested_key="embedder", required=True)
-    parser.add_argument(
-        "--evaluators", type=dict[str, EmbeddingEvaluator], enable_path=True, default="src/configs/jmteb.jsonnet"
-    )
-    parser.add_argument("--config", action=ActionConfigFile, help="Path to the config file.")
-    parser.add_argument("--save_dir", type=str, default=None, help="Directory to save the outputs")
-    parser.add_argument("--overwrite_cache", type=bool, default=False, help="Overwrite the save_dir if it exists")
-    parser.add_argument("--eval_exclude", type=list[str], default=None, help="Evaluators to exclude.")
-
-    args = parser.parse_args()
-
-    if args.eval_exclude is not None:
-        evaluator_keys = list(args.evaluators.keys())
-        # remove evaluators in eval_exclude
-        for key in evaluator_keys:
-            if key in args.eval_exclude:
-                args.evaluators.pop(key)
-
-    args = parser.instantiate_classes(args)
+    args = parse_args()
 
     main(
         text_embedder=args.embedder,

pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 ]
 description = "The evaluation scripts for JMTEB (Japanese Massive Text Embedding Benchmark)"
 name = "JMTEB"
-packages = [{from = ".", include = "src"}]
+packages = [{from = "src", include = "jmteb"}]
 readme = "README.md"
 version = "1.0.0"
 

src/jmteb/__main__.py

@@ -0,0 +1,9 @@
+from jmteb.main import main, parse_args
+
+args = parse_args()
+main(
+    text_embedder=args.embedder,
+    evaluators=args.evaluators,
+    save_dir=args.save_dir,
+    overwrite_cache=args.overwrite_cache,
+)

src/jmteb/embedders/__init__.py

@@ -0,0 +1,2 @@
+from jmteb.embedders.base import TextEmbedder
+from jmteb.embedders.sbert_embedder import SentenceBertEmbedder

src/embedders/sbert_embedder.py → src/jmteb/embedders/sbert_embedder.py

@@ -3,7 +3,7 @@
 import numpy as np
 from sentence_transformers import SentenceTransformer
 
-from src.embedders.base import TextEmbedder
+from jmteb.embedders.base import TextEmbedder
 
 
 class SentenceBertEmbedder(TextEmbedder):

src/evaluators/base.py → src/jmteb/evaluators/base.py

@@ -6,7 +6,7 @@
 
 from pydantic.dataclasses import dataclass
 
-from src.embedders import TextEmbedder
+from jmteb.embedders import TextEmbedder
 
 
 @dataclass

src/evaluators/classification/evaluator.py → src/jmteb/evaluators/classification/evaluator.py

@@ -6,8 +6,8 @@
 from loguru import logger
 from sklearn.metrics import accuracy_score, f1_score
 
-from src.embedders.base import TextEmbedder
-from src.evaluators.base import EmbeddingEvaluator, EvaluationResults
+from jmteb.embedders.base import TextEmbedder
+from jmteb.evaluators.base import EmbeddingEvaluator, EvaluationResults
 
 from .classifiers import Classifier, KnnClassifier, LogRegClassifier
 from .data import ClassificationDataset

src/evaluators/clustering/evaluator.py → src/jmteb/evaluators/clustering/evaluator.py

@@ -13,8 +13,8 @@
 )
 from sklearn.metrics import homogeneity_completeness_v_measure
 
-from src.embedders.base import TextEmbedder
-from src.evaluators.base import EmbeddingEvaluator, EvaluationResults
+from jmteb.embedders.base import TextEmbedder
+from jmteb.evaluators.base import EmbeddingEvaluator, EvaluationResults
 
 from .data import ClusteringDataset
 

src/evaluators/pair_classification/evaluator.py → src/jmteb/evaluators/pair_classification/evaluator.py

@@ -3,8 +3,8 @@
 from os import PathLike
 from pathlib import Path
 
-from src.embedders.base import TextEmbedder
-from src.evaluators.base import EmbeddingEvaluator, EvaluationResults
+from jmteb.embedders.base import TextEmbedder
+from jmteb.evaluators.base import EmbeddingEvaluator, EvaluationResults
 
 from .data import PairClassificationDataset
 from .threshold_accuracy import ThresholdAccuracyMetric

src/evaluators/retrieval/evaluator.py → src/jmteb/evaluators/retrieval/evaluator.py

@@ -12,8 +12,8 @@
 from loguru import logger
 from torch import Tensor
 
-from src.embedders.base import TextEmbedder
-from src.evaluators.base import EmbeddingEvaluator, EvaluationResults
+from jmteb.embedders.base import TextEmbedder
+from jmteb.evaluators.base import EmbeddingEvaluator, EvaluationResults
 
 from .data import RetrievalDocDataset, RetrievalQueryDataset
 

src/evaluators/sts/evaluator.py → src/jmteb/evaluators/sts/evaluator.py

@@ -9,8 +9,8 @@
 from scipy.stats import pearsonr, spearmanr
 from torch import Tensor
 
-from src.embedders.base import TextEmbedder
-from src.evaluators.base import EmbeddingEvaluator, EvaluationResults
+from jmteb.embedders.base import TextEmbedder
+from jmteb.evaluators.base import EmbeddingEvaluator, EvaluationResults
 
 from .data import STSDataset
 

src/jmteb/main.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from jsonargparse import ActionConfigFile, ArgumentParser
+from loguru import logger
+
+from jmteb.embedders import TextEmbedder
+from jmteb.evaluators import EmbeddingEvaluator
+from jmteb.utils.score_recorder import JsonScoreRecorder
+
+
+def main(
+    text_embedder: TextEmbedder,
+    evaluators: dict[str, EmbeddingEvaluator],
+    save_dir: str | None = None,
+    overwrite_cache: bool = False,
+):
+    logger.info(f"Start evaluating the following tasks\n{list(evaluators.keys())}")
+
+    if save_dir:
+        Path(save_dir).mkdir(parents=True, exist_ok=True)
+
+    score_recorder = JsonScoreRecorder(save_dir)
+
+    for eval_name, evaluator in evaluators.items():
+        logger.info(f"Evaluating {eval_name}")
+
+        cache_dir = None
+        if save_dir is not None:
+            cache_dir = Path(save_dir) / "cache" / eval_name
+
+        metrics = evaluator(text_embedder, cache_dir=cache_dir, overwrite_cache=overwrite_cache)
+        score_recorder.record_task_scores(
+            scores=metrics,
+            dataset_name=eval_name,
+            task_name=evaluator.__class__.__name__.replace("Evaluator", ""),
+        )
+
+        logger.info(f"Results for {eval_name}\n{json.dumps(metrics.as_dict(), indent=4, ensure_ascii=False)}")
+
+    logger.info(f"Saving result summary to {Path(save_dir) / 'summary.json'}")
+    score_recorder.record_summary()
+
+
+def parse_args():
+    parser = ArgumentParser(parser_mode="jsonnet")
+
+    parser.add_subclass_arguments(TextEmbedder, nested_key="embedder", required=True)
+    parser.add_argument(
+        "--evaluators", type=dict[str, EmbeddingEvaluator], enable_path=True, default="src/configs/jmteb.jsonnet"
+    )
+    parser.add_argument("--config", action=ActionConfigFile, help="Path to the config file.")
+    parser.add_argument("--save_dir", type=str, default=None, help="Directory to save the outputs")
+    parser.add_argument("--overwrite_cache", type=bool, default=False, help="Overwrite the save_dir if it exists")
+    parser.add_argument("--eval_exclude", type=list[str], default=None, help="Evaluators to exclude.")
+
+    args = parser.parse_args()
+
+    if args.eval_exclude is not None:
+        evaluator_keys = list(args.evaluators.keys())
+        # remove evaluators in eval_exclude
+        for key in evaluator_keys:
+            if key in args.eval_exclude:
+                args.evaluators.pop(key)
+
+    args = parser.instantiate_classes(args)

src/utils/score_recorder.py → src/jmteb/utils/score_recorder.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import Any
 
-from src.evaluators import EvaluationResults
+from jmteb.evaluators import EvaluationResults
 
 
 class AbstractScoreRecorder(ABC):

tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 
+from jmteb.embedders.sbert_embedder import SentenceBertEmbedder
+
 
 def pytest_addoption(parser: pytest.Parser):
     parser.addoption("--runslow", action="store_true", default=False, help="run slow tests")
@@ -17,3 +19,8 @@ def pytest_collection_modifyitems(config: pytest.Config, items: pytest.Parser):
     for item in items:
         if "slow" in item.keywords:
             item.add_marker(skip_slow)
+
+
+@pytest.fixture(scope="module")
+def embedder(model_name_or_path: str = "prajjwal1/bert-tiny"):
+    return SentenceBertEmbedder(model_name_or_path=model_name_or_path)

tests/embedders/test_sbert.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from src.embedders.sbert_embedder import SentenceBertEmbedder
+from jmteb.embedders.sbert_embedder import SentenceBertEmbedder
 
 MODEL_NAME_OR_PATH = "prajjwal1/bert-tiny"
 OUTPUT_DIM = 128

tests/evaluator/conftest.py

@@ -1,6 +1,6 @@
 import pytest
 
-from src.embedders.sbert_embedder import SentenceBertEmbedder
+from jmteb.embedders.sbert_embedder import SentenceBertEmbedder
 
 
 @pytest.fixture(scope="module")

tests/evaluator/test_classification_evaluator.py

@@ -1,11 +1,11 @@
-from src.evaluators.classification import (
+from jmteb.evaluators.classification import (
     ClassificationDataset,
     ClassificationEvaluator,
     ClassificationInstance,
     KnnClassifier,
     LogRegClassifier,
 )
-from src.evaluators.classification.data import JsonlClassificationDataset
+from jmteb.evaluators.classification.data import JsonlClassificationDataset
 
 
 class DummyClassificationDataset(ClassificationDataset):

tests/evaluator/test_clustering_evaluator.py

@@ -1,9 +1,9 @@
-from src.evaluators.clustering import (
+from jmteb.evaluators.clustering import (
     ClusteringDataset,
     ClusteringEvaluator,
     ClusteringInstance,
 )
-from src.evaluators.clustering.data import JsonlClusteringDataset
+from jmteb.evaluators.clustering.data import JsonlClusteringDataset
 
 
 class DummyClusteringDataset(ClusteringDataset):

tests/evaluator/test_pair_classification_evaluator.py

@@ -1,9 +1,9 @@
-from src.evaluators.pair_classification import (
+from jmteb.evaluators.pair_classification import (
     PairClassificationDataset,
     PairClassificationEvaluator,
     PairClassificationInstance,
 )
-from src.evaluators.pair_classification.data import JsonlPairClassificationDataset
+from jmteb.evaluators.pair_classification.data import JsonlPairClassificationDataset
 
 
 class DummyBinaryDataset(PairClassificationDataset):

tests/evaluator/test_retrieval_evaluator.py

@@ -1,11 +1,11 @@
-from src.evaluators.retrieval import (
+from jmteb.evaluators.retrieval import (
     RetrievalDoc,
     RetrievalDocDataset,
     RetrievalEvaluator,
     RetrievalQuery,
     RetrievalQueryDataset,
 )
-from src.evaluators.retrieval.data import (
+from jmteb.evaluators.retrieval.data import (
     JsonlRetrievalDocDataset,
     JsonlRetrievalQueryDataset,
 )

tests/evaluator/test_sts_evaluator.py

@@ -1,5 +1,5 @@
-from src.evaluators.sts import STSDataset, STSEvaluator, STSInstance
-from src.evaluators.sts.data import JsonlSTSDataset
+from jmteb.evaluators.sts import STSDataset, STSEvaluator, STSInstance
+from jmteb.evaluators.sts.data import JsonlSTSDataset
 
 
 class DummySTSDataset(STSDataset):