Adding ndcg calculations.

o19s · Dec 6, 2024 · 8208f25 · 8208f25
1 parent b4c29a2
commit 8208f25
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 20 deletions.
diff --git a/...-quality-evaluation-plugin/src/main/java/org/opensearch/eval/metrics/DcgSearchMetric.java b/...-quality-evaluation-plugin/src/main/java/org/opensearch/eval/metrics/DcgSearchMetric.java
@@ -26,6 +26,10 @@ public String getName() {
 
     @Override
     public double calculate() {
+        return calculateDcg(relevanceScores);
+    }
+
+    protected double calculateDcg(final List<Double> relevanceScores) {
 
         double dcg = 0.0;
         for(int i = 1; i <= relevanceScores.size(); i++) {

diff --git a/...quality-evaluation-plugin/src/main/java/org/opensearch/eval/metrics/NdcgSearchMetric.java b/...quality-evaluation-plugin/src/main/java/org/opensearch/eval/metrics/NdcgSearchMetric.java
@@ -12,11 +12,8 @@
 
 public class NdcgSearchMetric extends DcgSearchMetric {
 
-    private final List<Double> idealRelevanceScores;
-
-    public NdcgSearchMetric(final int k, final List<Double> relevanceScores, final List<Double> idealRelevanceScores) {
+    public NdcgSearchMetric(final int k, final List<Double> relevanceScores) {
         super(k, relevanceScores);
-        this.idealRelevanceScores = idealRelevanceScores;
     }
 
     @Override
@@ -27,17 +24,11 @@ public String getName() {
     @Override
     public double calculate() {
 
-        double dcg = super.calculate();
-
-        double idcg = 0.0;
-        for(int i = 0; i < idealRelevanceScores.size(); i++) {
-            double relevance = idealRelevanceScores.get(i);
-            idcg += relevance / Math.log(i + 2); // Add 2 to avoid log(1) = 0
-        }
+        // Make the ideal relevance scores by sorting the relevance scores largest to smallest.
+        relevanceScores.sort(Double::compare);
 
-        if(idcg == 0) {
-            return 0;
-        }
+        double dcg = super.calculate();
+        double idcg = super.calculateDcg(relevanceScores);
 
         return dcg / idcg;
 

diff --git a/...evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java b/...evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java
@@ -18,6 +18,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.eval.SearchQualityEvaluationPlugin;
 import org.opensearch.eval.metrics.DcgSearchMetric;
+import org.opensearch.eval.metrics.NdcgSearchMetric;
 import org.opensearch.eval.metrics.SearchMetric;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.search.SearchHit;
@@ -97,28 +98,24 @@ public void onResponse(final SearchResponse searchResponse) {
 
                             final List<String> orderedDocumentIds = new ArrayList<>();
 
-                           // LOGGER.info("Found {} results for query {}", searchResponse.getHits().getTotalHits().value, userQuery);
-
                             for (final SearchHit hit : searchResponse.getHits().getHits()) {
 
                                 final Map<String, Object> sourceAsMap = hit.getSourceAsMap();
                                 final String documentId = sourceAsMap.get(idField).toString();
 
-                                //LOGGER.info("     -- query {}, documentId: {}", userQuery, documentId);
                                 orderedDocumentIds.add(documentId);
 
                             }
 
-                            //LOGGER.info("Number of documents: " + orderedDocumentIds.size());
-
                             try {
 
                                 // TODO: If no hits are returned, there's no need to get the relevance scores.
                                 final List<Double> relevanceScores = getRelevanceScores(judgmentsId, userQuery, orderedDocumentIds, k);
 
                                 final SearchMetric dcgSearchMetric = new DcgSearchMetric(k, relevanceScores);
+                                final SearchMetric ndcgSearchmetric = new NdcgSearchMetric(k, relevanceScores);
+
                                 // TODO: Add these metrics in, too.
-                                //final SearchMetric ndcgSearchmetric = new NdcgSearchMetric(k, relevanceScores, idealRelevanceScores);
                                 //final SearchMetric precisionSearchMetric = new PrecisionSearchMetric(k, relevanceScores);
 
                                 //LOGGER.info("size list for query {}: {}", userQuery, relevanceScores.size());