Fix Title

README.md

@@ -30,11 +30,7 @@ Before using this metric, you need to install the dependencies:
 pip install -U sentence-transformers nltk
 ```
 
-Sem-nCG takes three mandatory arguments:
- - `predictions` - List of predictions
- - `references` - List of references
- - `documents` - List of input documents
-
+#### Python Usage
 ```python
 from evaluate import load
 predictions = [
@@ -55,6 +51,14 @@ mean_score, scores = metric.compute(predictions=predictions, references=referenc
 print(f"Mean SemnCG: {mean_score}")
 ```
 
+First step is to initialize the metric as `metric = load("nbansal/semncg", model_name=model_name)` where `model_name` is 
+the sentence embedding model. The default value is `all-MiniLM-L6-v2`.
+
+To `compute` the Sem-nCG scores, you need to provide three mandatory arguments:
+ - `predictions` - List of predictions
+ - `references` - List of references
+ - `documents` - List of input documents
+
 Sem-nCG also accepts several optional arguments:
  - `tokenize_sentences (bool)`: Flag to indicate whether to tokenize the sentences in the input documents. Default: True
  - `pre_compute_embeddings (bool)`: Flag to indicate whether to pre-compute embeddings for all sentences. Default=False

encoder_models.py

@@ -125,5 +125,3 @@ def get_sbert_encoder(model_name: str) -> SentenceTransformer:
         raise RuntimeError(str(err)) from None
 
     return encoder
-
-

semncg.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 """Sem-NCG metric"""
 
-from dataclasses import dataclass
-import evaluate
-import datasets
-import re
 import statistics
-from typing import Dict, List, Tuple, Union
+from dataclasses import dataclass
+from typing import List, Tuple, Union
 
+import datasets
+import evaluate
 import nltk
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
@@ -27,7 +26,8 @@ from tqdm import tqdm
 
 from .encoder_models import get_sbert_encoder, get_encoder
 from .type_aliases import DEVICE_TYPE, NDArray, DOCUMENT_TYPE
-from .utils import get_gpu, prep_sentences, flatten_list, slice_embeddings, is_nested_list_of_type, tokenize_and_prep_document
+from .utils import get_gpu, flatten_list, slice_embeddings, is_nested_list_of_type, \
+    tokenize_and_prep_document
 
 _CITATION = """\
 @inproceedings{akter-etal-2022-revisiting,
@@ -128,8 +128,6 @@ Examples:
 """
 
 
-
-
 @dataclass
 class RankedGains:
     """
@@ -154,7 +152,7 @@ class RankedGains:
     k: int
     ncg: float
 
-    
+
 def compute_cosine_similarity(doc_embeds: NDArray, ref_embeds: NDArray) -> List[float]:
     """
    Compute cosine similarity scores between each document embedding and reference embeddings.
@@ -333,7 +331,7 @@ def _validate_input_format(
 
 
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
-class SemnCG(evaluate.Metric):
+class SemNCG(evaluate.Metric):
     """
     SemnCG (Semantic Normalized Cumulative Gain) Metric.
 
@@ -454,7 +452,7 @@ class SemnCG(evaluate.Metric):
 
             # This is only done for debug case
             sent_tokenized_documents = documents
-            
+
             # Compute All Embeddings
             all_sentences = flatten_list(documents) + flatten_list(references) + flatten_list(predictions)
             embeddings = encoder.encode(all_sentences)
@@ -467,7 +465,7 @@ class SemnCG(evaluate.Metric):
             doc_embeddings = slice_embeddings(embeddings, document_sentences_count)
             ref_embeddings = slice_embeddings(embeddings[sum(document_sentences_count):], reference_sentences_count)
             pred_embeddings = slice_embeddings(
-                embeddings[sum(document_sentences_count+reference_sentences_count):], prediction_sentences_count
+                embeddings[sum(document_sentences_count + reference_sentences_count):], prediction_sentences_count
             )
 
             iterable_obj = zip(pred_embeddings, ref_embeddings, doc_embeddings)
@@ -495,7 +493,7 @@ class SemnCG(evaluate.Metric):
                 doc_embeddings = doc
                 ref_embeddings = ref
                 pred_embeddings = pred
-                
+
                 doc_sentences = sent_tokenized_documents[idx]
 
             # Compute Pair-Wise Cosine Similarity

tests.py

@@ -1,16 +1,27 @@
-import statistics
 import unittest
-from unittest.mock import patch, MagicMock
 
 import numpy as np
 import torch
-from numpy.testing import assert_almost_equal
 from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
 
 from .encoder_models import SBertEncoder, get_encoder, get_sbert_encoder
-from .semncg import RankedGains, compute_cosine_similarity, compute_gain, score_ncg, compute_ncg, _validate_input_format, SemnCG
-from .utils import get_gpu, slice_embeddings, is_nested_list_of_type, flatten_list, prep_sentences, tokenize_and_prep_document
+from .semncg import (
+    RankedGains,
+    compute_cosine_similarity,
+    compute_gain,
+    score_ncg,
+    compute_ncg,
+    _validate_input_format,
+    SemNCG
+)
+from .utils import (
+    get_gpu,
+    slice_embeddings,
+    is_nested_list_of_type,
+    flatten_list,
+    prep_sentences,
+    tokenize_and_prep_document
+)
 
 
 class TestUtils(unittest.TestCase):
@@ -116,7 +127,7 @@ class TestUtils(unittest.TestCase):
         doc_embeddings = slice_embeddings(all_embeddings, document_sentences_count)
         ref_embeddings = slice_embeddings(all_embeddings[sum(document_sentences_count):], reference_sentences_count)
         pred_embeddings = slice_embeddings(
-            all_embeddings[sum(document_sentences_count+reference_sentences_count):], pred_sentences_count
+            all_embeddings[sum(document_sentences_count + reference_sentences_count):], pred_sentences_count
         )
 
         self.assertTrue(doc_embeddings, expected_doc_embeddings)
@@ -350,7 +361,7 @@ class TestValidateInputFormat(unittest.TestCase):
 class TestSemnCG(unittest.TestCase):
     def setUp(self):
         self.model_name = "stsb-distilbert-base"
-        self.metric = SemnCG(self.model_name)
+        self.metric = SemNCG(self.model_name)
 
     def _basic_assertion(self, result, debug: bool = False):
         self.assertIsInstance(result, tuple)

type_aliases.py

@@ -1,9 +1,7 @@
-
-from typing import List, Union, Tuple
+from typing import List, Union
 
 from numpy.typing import NDArray
 
-
 NumSentencesType = Union[List[int], List[List[int]]]
 EmbeddingSlicesType = Union[List[NDArray], List[List[NDArray]]]
 DEVICE_TYPE = Union[bool, str, int, List[Union[str, int]]]

utils.py

@@ -1,11 +1,9 @@
-
 import string
-from typing import List, Tuple, Union
+from typing import List, Union
 
 import nltk
-import numpy as np
-from numpy.typing import NDArray
 import torch
+from numpy.typing import NDArray
 
 from .type_aliases import DEVICE_TYPE, ENCODER_DEVICE_TYPE, NumSentencesType, EmbeddingSlicesType
 
@@ -204,7 +202,8 @@ def is_nested_list_of_type(lst_obj, element_type, depth: int) -> bool:
     if depth == 0:
         return isinstance(lst_obj, element_type)
     elif depth > 0:
-        return isinstance(lst_obj, list) and all(is_nested_list_of_type(item, element_type, depth - 1) for item in lst_obj)
+        return isinstance(lst_obj, list) and all(
+            is_nested_list_of_type(item, element_type, depth - 1) for item in lst_obj)
     else:
         raise ValueError("Depth can't be negative")