Hard negative mining for Retriever fine-tuning

nemo_curator/modules/semantic_dedup/clusteringmodel.py

@@ -15,6 +15,7 @@
 
 import logging
 import os
+import pdb
 import shutil
 import time
 from typing import Optional, Union
@@ -124,7 +125,9 @@ def __call__(self, embeddings_dataset: DocumentDataset):
             )
 
         with performance_report_if_with_ts_suffix(self.profile_dir, "clustering-model"):
-            embeddings_df = embeddings_df[[self.id_col, self.embedding_column]]
+
+            # embeddings_df = embeddings_df[[self.id_col, self.embedding_col]]
+
             embeddings_df = embeddings_df.repartition(
                 partition_size=self.clustering_input_partition_size
             )

tutorials/nemo-retriever-synthetic-data-generation/config/config.py

@@ -92,3 +92,30 @@ class RetrieverEvalSDGConfig(BaseConfig):
     Question:
     {question}
     """
+
+
+@dataclass
+class RetrieverHardNegativeMiningConfig(BaseConfig):
+    """
+    Configuration for mining hard negatives
+
+    Attributes:
+
+    """
+
+    model_name: str = "nvdev/nvidia/llama-3.2-nv-embedqa-1b-v2"
+    model_type: str = "nvidia"
+    base_url: str = "https://integrate.api.nvidia.com/v1"
+    api_key: str = None
+    truncate: str = "END"
+    hard_negatives_to_mine: int = 4
+    hard_neg_mining_algorithm: str = "topk_abs"
+    max_hardness_threshold: float = 0.75
+    min_hardness_threshold: float = 0
+    query_prefix: str = ""
+    passage_prefix: str = ""
+    percpos: float = 0.95
+    min_cluster_size: int = 50
+    max_number_clusters: int = 200
+    cluster_output_dir: str = "/tmp/clusters"
+    logger_output_dir: str = "/tmp/logs"

tutorials/nemo-retriever-synthetic-data-generation/config/hard-negative-mining-config.yaml

@@ -0,0 +1,19 @@
+# HARD-NEGATIVE MINING parameters
+
+# base_url: "https://integrate.api.nvidia.com/v1"
+model_name: "sentence-transformers/all-MiniLM-L6-v2"
+model_type: "hf"
+query_prefix: "query:"
+passage_prefix: "passage:"
+api_key: "your api key here"
+truncate: "END"
+hard_negatives_to_mine: 4
+hard_neg_mining_algorithm: "topk_percpos"
+percpos: 0.95
+
+
+# SEMANTIC CLUSTERING parameters
+min_cluster_size: 100
+max_number_clusters: 200
+cluster_output_dir: "/workspace/hnm/clusters"
+logger_output_dir: "/workspace/hnm/logs"

tutorials/nemo-retriever-synthetic-data-generation/main.py

@@ -25,7 +25,6 @@
 from dask.distributed import progress
 from retriever_evalset_generator import RetrieverEvalSetGenerator
 
-from config.config import RetrieverEvalSDGConfig
 from nemo_curator import AsyncOpenAIClient, ScoreFilter, Sequential, get_client
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.filters import (
@@ -36,6 +35,11 @@
 from nemo_curator.modules.filter import Score, ScoreFilter
 from nemo_curator.utils.file_utils import get_all_files_paths_under
 
+config = importlib.import_module(
+    "tutorials.nemo-retriever-synthetic-data-generation.config.config"
+)
+RetrieverEvalSDGConfig = config.RetrieverEvalSDGConfig
+
 
 def get_pipeline(args: Any) -> Any:
 
@@ -53,6 +57,7 @@ def get_pipeline(args: Any) -> Any:
         sdg_pipeline = None
 
     filters = []
+
     if args.pipeline_type == "filter":
         if cfg.easiness_filter:
             filters.append(

tutorials/nemo-retriever-synthetic-data-generation/mine_hard_negatives.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import glob
+import importlib
+import os
+import pdb
+import shutil
+import time
+from pathlib import Path
+from typing import Any, List
+
+from retriever_hardnegative_miner import HardNegativeMiner
+from tqdm.dask import TqdmCallback
+
+from config.config import RetrieverHardNegativeMiningConfig
+from nemo_curator.datasets import DocumentDataset
+from nemo_curator.utils.distributed_utils import get_client
+from nemo_curator.utils.file_utils import get_all_files_paths_under
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        default="",
+        help="Input dir path containing annotated data files in jsonl format",
+    )
+    parser.add_argument(
+        "--hard-negative-mining-config",
+        type=str,
+        default="",
+        help="Configuration yaml file path containing config for hard negative mining",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="",
+        help="Output file containing hard negatives",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=None,
+        help="The API key to use for the synthetic data generation LLM client.",
+    )
+    parser.add_argument(
+        "--api-timeout",
+        type=int,
+        default=120,
+        help="The timeout value for API calls in seconds.",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.input_dir):
+        raise ValueError("Input directory not found")
+
+    if os.path.exists(args.output_dir):
+        raise ValueError("Output dir exists already, use a new file name!")
+
+    if args.input_dir:
+        input_files = get_all_files_paths_under(args.input_dir, keep_extensions="part")
+        input_dataset = DocumentDataset.read_json(input_files)
+    else:
+        raise ValueError("provide input file path")
+
+    if args.hard_negative_mining_config:
+        cfg = RetrieverHardNegativeMiningConfig.from_yaml(
+            args.hard_negative_mining_config
+        )
+
+    else:
+        raise ValueError("provide config for hard negative mining")
+    if args.api_key:
+        cfg.api_key = args.api_key
+
+    mine_hard_negatives = HardNegativeMiner(cfg)
+    print("Mining hard negatives ...")
+    st_time = time.time()
+    mined_dataset = mine_hard_negatives(input_dataset)
+
+    print("Time taken = {:.2f} s".format(time.time() - st_time))
+    print("Saving data in jsonl format ...")
+    mined_dataset.df.to_json(
+        os.path.join(args.output_dir), lines=True, orient="records"
+    )
+
+
+if __name__ == "__main__":
+    dask_client = get_client(cluster_type="gpu")
+    main()

tutorials/nemo-retriever-synthetic-data-generation/repartition.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import glob
+import importlib
+import os
+import pdb
+import shutil
+import time
+from typing import Any, List
+
+from retriever_hardnegative_miner import HardNegativeMiner
+from tqdm.dask import TqdmCallback
+
+from config.config import RetrieverHardNegativeMiningConfig
+from nemo_curator.datasets import DocumentDataset
+from nemo_curator.utils.distributed_utils import get_client
+from nemo_curator.utils.file_utils import get_all_files_paths_under
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        default="",
+        help="Input dir path containing annotated data files in jsonl format",
+    )
+    parser.add_argument(
+        "--hard-negative-mining-config",
+        type=str,
+        default="",
+        help="Configuration yaml file path containing config for hard negative mining",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="",
+        help="Output file containing clustered dataset",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=None,
+        help="The API key to use for the synthetic data generation LLM client.",
+    )
+    parser.add_argument(
+        "--api-timeout",
+        type=int,
+        default=120,
+        help="The timeout value for API calls in seconds.",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.input_dir):
+        raise ValueError("Input directory not found")
+
+    if os.path.exists(args.output_dir):
+        raise ValueError("Output dir exists already, use a new file name!")
+
+    if args.input_dir:
+        input_files = get_all_files_paths_under(args.input_dir, keep_extensions="jsonl")
+        input_dataset = DocumentDataset.read_json(input_files)
+    else:
+        raise ValueError("provide input file path")
+    if args.hard_negative_mining_config:
+        cfg = RetrieverHardNegativeMiningConfig.from_yaml(
+            args.hard_negative_mining_config
+        )
+    else:
+        raise ValueError("provide config for hard negative mining")
+    if args.api_key:
+        cfg.api_key = args.api_key
+
+    st_time = time.time()
+    miner = HardNegativeMiner(cfg)
+    clustered_dataset = miner.repartition_semantic_similarity(input_dataset)
+    clustered_dataset.persist()
+
+    # saving clustered dataset
+    print("saving clustered dataset")
+    clustered_dataset.df.to_json(os.path.join(args.output_dir, "clustered_dataset"))
+    print("Time taken to cluster data = {:.2f} s".format(time.time() - st_time))
+
+
+if __name__ == "__main__":
+    dask_client = get_client(cluster_type="gpu")
+    main()

tutorials/nemo-retriever-synthetic-data-generation/retriever_evalset_generator.py

@@ -24,8 +24,6 @@
 from tqdm import tqdm
 
 tqdm.pandas()
-
-# from tqdm.dask import TqdmCallback
 import importlib
 
 import dask.array as da
@@ -145,7 +143,6 @@ def _process_on_partition(self, df: pd.DataFrame) -> pd.DataFrame:
 
         df = df.explode("qa_pairs").reset_index(drop=True)
         df["question"] = df["qa_pairs"].apply(lambda x: x["question"])
-
         df["question-id"] = df["question"].apply(self._get_random_hash)
         df["answer"] = df["qa_pairs"].apply(lambda x: x["answer"])
         df["score"] = df["question"].apply(lambda x: 1)