Add distance messure

esqa/cli.py

@@ -5,6 +5,7 @@
 
 import click
 
+from esqa.distance import load_rankings, compare_rankings
 from esqa.save import RankingSaver
 from esqa.validation_config import load
 from esqa.validator import Validator
@@ -47,6 +48,19 @@ def check(config, index):
 def save(config, index):
     runner = RankingSaver()
     results = runner.run(config=load(config), index_name=index)
+    print(_dump(list(results.values())))
+
+
+@main.command()
+@click.option("-r", "--ranking", type=str, help="ranking file")
+@click.option("-c", "--config", type=str, help="configuration file")
+@click.option("-t", "--threshold", type=float, help="threshold", default=0.7)
+@click.option("--index", type=str, help="target index name", required=True)
+def ranking(ranking, config, threshold, index):
+    runner = RankingSaver()
+    rankings = runner.run(config=load(config), index_name=index)
+    compared_rankings = load_rankings(ranking)
+    results = compare_rankings(rankings, compared_rankings, threshold)
     print(_dump(results))
 
 

esqa/distance.py

@@ -0,0 +1,48 @@
+import dataclasses
+import json
+import rbo
+from typing import Dict, List
+
+from esqa.save import Ranking
+
+
+@dataclasses.dataclass
+class FailedRanking:
+    name: str
+    similarity: float
+    ranking_pair: List[tuple]
+
+
+def load_rankings(path: str) -> Dict:
+    with open(path) as f:
+        rankings = json.load(f)
+    results = {}
+    for ranking in rankings:
+        results[ranking["name"]] = Ranking(ranking["name"], ranking["query"], ranking["ranking"])
+    return results
+
+
+def _extract(ranking: Ranking) -> List[str]:
+    return [e["id"] for e in ranking.ranking]
+
+
+def _compare(ranking_a, ranking_b):
+    return rbo.rbo.RankingSimilarity(ranking_a, ranking_b).rbo()
+
+
+def _generate(ranking_a: Ranking, ranking_b: Ranking, similarity: float):
+    return FailedRanking(
+        name=ranking_a.name,
+        similarity=similarity,
+        ranking_pair=list(zip(_extract(ranking_a), _extract(ranking_b)))
+    )
+
+
+def compare_rankings(rankings_a: Dict[str, Ranking], rankings_b: Dict[str, Ranking], threshold: float) -> List[FailedRanking]:
+    results = []
+    for ranking_name in rankings_a:
+        similarity = _compare(_extract(rankings_a[ranking_name]), _extract(rankings_b[ranking_name]))
+        if similarity > threshold:
+            continue
+        results.append(_generate(rankings_a[ranking_name], rankings_b[ranking_name], similarity))
+    return results

esqa/save.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import List
+from typing import List, Dict
 
 from elasticsearch import Elasticsearch
 
@@ -11,7 +11,6 @@
 class Ranking:
     name: str
     query: dict
-    asserts: List[EsAssert]
     ranking: List[dict]
 
 
@@ -21,20 +20,20 @@ class RankingSaver:
     def __init__(self):
         self.client = Elasticsearch([ELASTICSEARCH_URL])
 
-    def run(self, config: Configuration, index_name: str):
-        results = []
+    def run(self, config: Configuration, index_name: str) -> Dict[str, Ranking]:
+        results = {}
         for case in config.cases:
-            results.append(self._get(case, index_name))
+            ranking = self._get(case, index_name)
+            results[ranking.name] = ranking
         return results
 
-    def _get(self, case: Case, index_name: str):
+    def _get(self, case: Case, index_name: str) -> Ranking:
         search_results = self.client.search(body=case.query, index=index_name)
         return self._format(search_results, case)
 
     def _format(self, search_results: dict, case: Case) -> Ranking:
         return Ranking(
             case.name,
             case.query,
-            case.asserts,
             [{"id": candidate["_id"], "source": candidate["_source"]} for i, candidate in enumerate(search_results["hits"]["hits"])]
         )

pyproject.toml

@@ -9,6 +9,7 @@ readme = "README.md"
 python = "^3.7"
 click = ">=8.0.0"
 elasticsearch = "7.10.1"
+rbo = "^0.1.2"
 
 [tool.poetry.dev-dependencies]
 toml = "^0.10.0"