hard-negatives 0.1.0__tar.gz

hard_negatives-0.1.0/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: hard-negatives
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: accelerate>=1.14.0
+Requires-Dist: datasets>=5.0.0
+Requires-Dist: flash-attn>=2.8.3.post1
+Requires-Dist: psutil>=7.2.2
+Requires-Dist: sentence-transformers>=5.6.0
+Requires-Dist: torch>=2.13.0
+Requires-Dist: voyager>=2.1.0
+
+## Usage
+```bash
+accelerate launch --multi_gpu main.py --input_path /home/henry/hard_negatives/test_pairs.parquet
+```
+
+```bibtex
+@software{shippole2026hardnegatives,
+  title        = {Hard Negatives: A Python Library for Mining Hard Negative Examples},
+  author       = {Shippole, Enrico and Chaffin, Antoine and Aarsen, Tom},
+  year         = {2026},
+  url          = {https://github.com/teraflop-ai/hard-negatives},
+  version      = {1.0.0},
+  note         = {Python package available on PyPI at \url{https://pypi.org/project/hard-negatives/}}
+}
+```

hard_negatives-0.1.0/README.md

@@ -0,0 +1,15 @@
+## Usage
+```bash
+accelerate launch --multi_gpu main.py --input_path /home/henry/hard_negatives/test_pairs.parquet
+```
+
+```bibtex
+@software{shippole2026hardnegatives,
+  title        = {Hard Negatives: A Python Library for Mining Hard Negative Examples},
+  author       = {Shippole, Enrico and Chaffin, Antoine and Aarsen, Tom},
+  year         = {2026},
+  url          = {https://github.com/teraflop-ai/hard-negatives},
+  version      = {1.0.0},
+  note         = {Python package available on PyPI at \url{https://pypi.org/project/hard-negatives/}}
+}
+```

hard_negatives-0.1.0/pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "hard-negatives"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "accelerate>=1.14.0",
+    "datasets>=5.0.0",
+    "flash-attn>=2.8.3.post1",
+    "psutil>=7.2.2",
+    "sentence-transformers>=5.6.0",
+    "torch>=2.13.0",
+    "voyager>=2.1.0",
+]

hard_negatives-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

hard_negatives-0.1.0/src/hard_negatives/args/voyager_args.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+
+import argparse
+
+DEFAULT_REPORT_OUTPUT_DIR = str(Path(__file__).resolve().parent / "output" / "filtering_reports")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Mine hard negatives from a Hugging Face dataset."
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        default="lightonai/DenseOn",
+        help="The sentence transformer model to use."
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        default=8192,
+        type=int,
+        help="The seqlen to use during encoding."
+    )
+    parser.add_argument(
+        "--input_path",
+        type=str,
+        required=True,
+        help="The input path to parquet files.",
+    )
+    parser.add_argument(
+        "--num_dim",
+        type=int,
+        default=768
+    )
+    parser.add_argument(
+        "--hnsw_m",
+        type=int,
+        default=64
+    )
+    parser.add_argument(
+        "--hnsw_ef",
+        type=int,
+        default=200,
+    )
+    parser.add_argument(
+        "--mini_batch",
+        type=int,
+        default=50,
+        help=f"Batch size for encoding (default: 50)",
+    )
+    parser.add_argument(
+        "--test",
+        action="store_true",
+        help="Run in test mode with limited data",
+    )
+    parser.add_argument(
+        "--num_negatives",
+        type=int,
+        default=2048,
+        help=f"Number of negative samples to keep per positive (default: 2048)",
+    )
+    parser.add_argument(
+        "--nvembed_threshold",
+        type=float,
+        default=0.95,
+        help=f"Threshold for filtering negatives: keep negatives with sim < threshold * positive_sim (default: 0.95)",
+    )
+    parser.add_argument(
+        "--max_negatives_filter",
+        type=int,
+        default=10,
+        help=f"Maximum number of negatives to keep after threshold filtering (default: 10)",
+    )
+    parser.add_argument(
+        "--report_output_dir",
+        type=str,
+        default=DEFAULT_REPORT_OUTPUT_DIR,
+        help="Directory to save filtering analysis reports",
+    )
+
+    return parser.parse_args()

hard_negatives-0.1.0/src/hard_negatives/data_loader.py

@@ -0,0 +1,30 @@
+from datasets import load_dataset
+
+def load_pair_dataset(path: str):
+    data = load_dataset("parquet", data_files=path, split="train")
+
+    if set(data.column_names) != {"query", "document"}:
+        raise ValueError("Dataset must contain exactly: query, document")
+
+    query_to_id = {}
+    document_to_id = {}
+    positives = {}
+
+    for query, document in zip(data["query"], data["document"]):
+        query_id = query_to_id.setdefault(query, len(query_to_id))
+        document_id = document_to_id.setdefault(
+            document, len(document_to_id)
+        )
+        positives.setdefault(query_id, set()).add(document_id)
+
+    queries = {query_id: query for query, query_id in query_to_id.items()}
+    documents = {
+        document_id: document
+        for document, document_id in document_to_id.items()
+    }
+    positives = {
+        query_id: sorted(document_ids)
+        for query_id, document_ids in positives.items()
+    }
+
+    return queries, documents, positives

hard_negatives-0.1.0/src/hard_negatives/mine_negatives.py

@@ -0,0 +1,314 @@
+import os
+from pathlib import Path
+
+import sys
+import json
+import numpy as np
+
+from accelerate import Accelerator
+from accelerate.utils import gather_object
+from datasets import Dataset, Features, Value
+from hard_negatives.data_loader import load_pair_dataset
+
+from voyager import Index, Space
+
+from hard_negatives.prepare_model import load_model
+
+
+PATH_TO_HUB_UPLOAD = "TeraflopAI/mined_example"
+
+GATHER_CHUNK_SIZE = 2000
+QUERY_BATCH_SIZE = 1000
+
+
+# Extra buffer added to k when querying the index, to leave room for positives/dupes
+K_BUFFER = 500
+
+
+
+def mine_negatives(args):
+    dataset_name = Path(args.input_path).stem
+
+    queries, unique_index_to_doc_text, positives = load_pair_dataset(
+        args.input_path
+    )
+
+    model = load_model(args.model_name, args.max_seq_len)
+
+    accelerator = Accelerator()
+    accelerator.prepare(model)
+
+    # build distributed index for all unique documents
+    unique_doc_indices = sorted(unique_index_to_doc_text.keys())
+    process_indices = unique_doc_indices[
+        accelerator.process_index :: accelerator.num_processes
+    ]
+
+    if accelerator.is_main_process:
+        index = Index(Space.Cosine, num_dimensions=args.num_dim, M=args.hnsw_m, ef_construction=args.hnsw_ef)
+        doc_id_to_embedding = {}
+
+    # Process documents in chunks
+    dataset = [unique_index_to_doc_text[i] for i in process_indices]
+    num_chunks = (len(dataset) + GATHER_CHUNK_SIZE - 1) // GATHER_CHUNK_SIZE
+
+    if accelerator.is_main_process:
+        print(f"Processing documents in {num_chunks} chunks...")
+
+    for chunk_idx in range(num_chunks):
+        start_idx = chunk_idx * GATHER_CHUNK_SIZE
+        end_idx = min((chunk_idx + 1) * GATHER_CHUNK_SIZE, len(dataset))
+
+        chunk_dataset = dataset[start_idx:end_idx]
+        chunk_indices = process_indices[start_idx:end_idx]
+
+        embeddings = model.encode(chunk_dataset, batch_size=args.mini_batch, show_progress_bar=True)
+        full_embeddings = gather_object(embeddings)
+        full_indices = gather_object(chunk_indices)
+
+        if accelerator.is_main_process:
+            index.add_items(full_embeddings, full_indices)
+
+            for embedding, indice in zip(full_embeddings, full_indices):
+                doc_id_to_embedding[indice] = embedding
+
+        # Free memory after each chunk
+        del full_embeddings
+        del full_indices
+        del embeddings
+
+        if accelerator.is_main_process:
+            print(f"Processed document chunk {chunk_idx + 1}/{num_chunks}")
+
+    accelerator.wait_for_everyone()
+
+    # Process queries in chunks
+    process_indices = list(range(len(queries)))[
+        accelerator.process_index :: accelerator.num_processes
+    ]
+
+    if accelerator.is_main_process:
+        query_id_to_embedding = {}
+
+    dataset = [queries[i] for i in process_indices]
+    num_chunks = (len(dataset) + GATHER_CHUNK_SIZE - 1) // GATHER_CHUNK_SIZE
+
+    if accelerator.is_main_process:
+        print(f"Processing queries in {num_chunks} chunks...")
+
+    for chunk_idx in range(num_chunks):
+        start_idx = chunk_idx * GATHER_CHUNK_SIZE
+        end_idx = min((chunk_idx + 1) * GATHER_CHUNK_SIZE, len(dataset))
+
+        chunk_dataset = dataset[start_idx:end_idx]
+        chunk_indices = process_indices[start_idx:end_idx]
+
+        embeddings = model.encode(chunk_dataset, batch_size=args.mini_batch, show_progress_bar=True)
+        full_embeddings = gather_object(embeddings)
+        full_indices = gather_object(chunk_indices)
+
+        if accelerator.is_main_process:
+            for embedding, indice in zip(full_embeddings, full_indices):
+                query_id_to_embedding[indice] = embedding
+
+        # Free memory after each chunk
+        del full_embeddings
+        del full_indices
+        del embeddings
+
+        if accelerator.is_main_process:
+            print(f"Processed query chunk {chunk_idx + 1}/{num_chunks}")
+
+    accelerator.wait_for_everyone()
+
+    # exit non-main processes after ALL processing is complete
+    if not accelerator.is_main_process:
+        print(f"Rank {accelerator.process_index} finished all processing, exiting")
+        sys.exit(0)
+
+    if accelerator.is_main_process:
+        # Prepare three datasets
+
+        # 1. Documents dataset (document_id, document_text)
+        documents_rows = []
+        for doc_id, doc_text in unique_index_to_doc_text.items():
+            documents_rows.append({"document_id": doc_id, "document": doc_text})
+
+        documents_features = Features({
+            "document_id": Value("int64"),
+            "document": Value("large_string"),
+        })
+        documents_dataset = Dataset.from_list(documents_rows, features=documents_features)
+        documents_dataset.push_to_hub(
+            PATH_TO_HUB_UPLOAD,
+            config_name="documents",
+            data_dir="documents",
+            split=dataset_name,
+        )
+        print(f"Pushed documents dataset with {len(documents_rows)} documents")
+
+        # 2. Queries dataset (query_id, query_text)
+        queries_rows = []
+        for query_id, query_text in queries.items():
+            queries_rows.append({"query_id": query_id, "query": query_text})
+
+        queries_features = Features({
+            "query_id": Value("int64"),
+            "query": Value("large_string"),
+        })
+        queries_dataset = Dataset.from_list(queries_rows, features=queries_features)
+        queries_dataset.push_to_hub(
+            PATH_TO_HUB_UPLOAD,
+            config_name="queries",
+            data_dir="queries",
+            split=dataset_name,
+        )
+        print(f"Pushed queries dataset with {len(queries_rows)} queries")
+
+        # 3. Scores dataset (query_id, document_ids, scores)
+        scores_rows = []
+
+        # Batch query the index
+        query_ids_list = list(positives.keys())
+        num_query_batches = (len(query_ids_list) + QUERY_BATCH_SIZE - 1) // QUERY_BATCH_SIZE
+
+        all_query_results = {}  # Store results: query_id -> (indexes, distances)
+
+        for batch_idx in range(num_query_batches):
+            start_idx = batch_idx * QUERY_BATCH_SIZE
+            end_idx = min((batch_idx + 1) * QUERY_BATCH_SIZE, len(query_ids_list))
+
+            batch_query_ids = query_ids_list[start_idx:end_idx]
+            query_embeddings_batch = [query_id_to_embedding[qid] for qid in batch_query_ids]
+
+            print(
+                f"Querying index for batch {batch_idx + 1}/{num_query_batches} ({len(batch_query_ids)} queries)..."
+            )
+            # Query with extra buffer to account for potential duplicates and positives
+            k_value = min(args.num_negatives + K_BUFFER, len(unique_index_to_doc_text))
+            batch_indexes, batch_distances = index.query(query_embeddings_batch, k=k_value)
+
+            # Store results
+            for i, qid in enumerate(batch_query_ids):
+                all_query_results[qid] = (batch_indexes[i], batch_distances[i])
+
+        print("Index querying complete")
+
+        # Process results for each query
+        for query_id in query_ids_list:
+            indexes, distances = all_query_results[query_id]
+            positives_list = positives[query_id]
+
+            for positive in positives_list:
+                document_ids = []
+                scores_list = []
+
+                # Calculate similarity for positive (first element)
+                positive_simi = np.dot(
+                    query_id_to_embedding[query_id], doc_id_to_embedding[positive]
+                ) / (
+                    np.linalg.norm(query_id_to_embedding[query_id])
+                    * np.linalg.norm(doc_id_to_embedding[positive])
+                )
+
+                # Add positive as first element
+                document_ids.append(positive)
+                scores_list.append(float(positive_simi))
+
+                # Add negatives with their distances converted to similarities
+                negatives = []
+                negative_scores = []
+                for ind, distance in zip(indexes, distances):
+                    if ind not in positives_list:
+                        negatives.append(ind)
+                        negative_scores.append(
+                            1 - distance
+                        )  # Convert distance to similarity
+                        if len(negatives) >= args.num_negatives:
+                            break
+
+                # Only create row if we have enough negatives
+                if len(negatives) < args.num_negatives:
+                    continue
+
+                # Add negative document IDs and scores
+                document_ids.extend(negatives)
+                scores_list.extend(negative_scores)
+
+                scores_rows.append(
+                    {
+                        "query_id": query_id,
+                        "document_ids": document_ids,
+                        "scores": scores_list,
+                    }
+                )
+
+        # Analyze filtering impact
+        print("\n" + "="*50)
+        print("Filtering Analysis")
+        print("="*50)
+        filtered_count = 0
+        total_negatives_before = 0
+        total_negatives_after = 0
+
+        for row in scores_rows:
+            positive_score = row["scores"][0]  # First score is the positive
+            negative_scores = row["scores"][1:]  # Rest are negatives
+            total_negatives_before += len(negative_scores)
+
+            # Apply threshold filter
+            threshold_value = args.nvembed_threshold * positive_score
+            filtered_negatives = [score for score in negative_scores if score < threshold_value]
+
+            # Apply max negatives filter if specified
+            filtered_negatives = filtered_negatives[:args.max_negatives_filter]
+
+            total_negatives_after += len(filtered_negatives)
+
+            # Count rows with at least max_negatives_filter negatives after filtering
+            if len(filtered_negatives) >= args.max_negatives_filter:
+                filtered_count += 1
+
+        # Prepare report data
+        report_data = {
+            "dataset_name": dataset_name,
+            "original_rows": len(scores_rows),
+            "rows_after_filtering": filtered_count,
+            "rows_dropped": len(scores_rows) - filtered_count,
+            "retention_rate": filtered_count / len(scores_rows) * 100 if len(scores_rows) > 0 else 0,
+            "nvembed_threshold": args.nvembed_threshold,
+            "min_negatives_required": args.max_negatives_filter,
+        }
+
+        # Save report to JSON
+        os.makedirs(args.report_output_dir, exist_ok=True)
+        report_path = os.path.join(args.report_output_dir, f"{dataset_name}_filtering_report.json")
+        with open(report_path, "w") as f:
+            json.dump(report_data, f, indent=2)
+
+        print(f"Original rows: {report_data['original_rows']}")
+        print(f"Rows after filtering (>= {args.max_negatives_filter} negatives): {report_data['rows_after_filtering']}")
+        print(f"Rows dropped: {report_data['rows_dropped']}")
+        print(f"Retention rate: {report_data['retention_rate']:.2f}%")
+        print(f"Threshold: {report_data['nvembed_threshold']}")
+        print(f"Min negatives required: {report_data['min_negatives_required'] if report_data['min_negatives_required'] else 'None (keep all)'}")
+        print(f"Report saved to: {report_path}")
+        print("="*50 + "\n")
+
+        scores_features = Features({
+            "query_id": Value("int64"),
+            "document_ids": [Value("int64")],
+            "scores": [Value("float64")],
+        })
+
+        scores_dataset = Dataset.from_list(scores_rows, features=scores_features)
+        scores_dataset.push_to_hub(
+            PATH_TO_HUB_UPLOAD,
+            config_name="scores",
+            data_dir="scores",
+            split=dataset_name,
+        )
+        print(f"Pushed scores dataset with {len(scores_rows)} query-document pairs")
+
+        print(f"All three datasets pushed successfully for {dataset_name}")
+

hard_negatives-0.1.0/src/hard_negatives/prepare_model.py

@@ -0,0 +1,13 @@
+import torch
+from sentence_transformers import SentenceTransformer
+
+def load_model(model_name: str, max_seq_len: int):
+    model = SentenceTransformer(
+        model_name,
+        model_kwargs={
+            "attn_implementation": "sdpa",
+            "dtype": torch.bfloat16,
+        },
+    )
+    model.max_seq_length = max_seq_len
+    return model

hard_negatives-0.1.0/src/hard_negatives.egg-info/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: hard-negatives
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: accelerate>=1.14.0
+Requires-Dist: datasets>=5.0.0
+Requires-Dist: flash-attn>=2.8.3.post1
+Requires-Dist: psutil>=7.2.2
+Requires-Dist: sentence-transformers>=5.6.0
+Requires-Dist: torch>=2.13.0
+Requires-Dist: voyager>=2.1.0
+
+## Usage
+```bash
+accelerate launch --multi_gpu main.py --input_path /home/henry/hard_negatives/test_pairs.parquet
+```
+
+```bibtex
+@software{shippole2026hardnegatives,
+  title        = {Hard Negatives: A Python Library for Mining Hard Negative Examples},
+  author       = {Shippole, Enrico and Chaffin, Antoine and Aarsen, Tom},
+  year         = {2026},
+  url          = {https://github.com/teraflop-ai/hard-negatives},
+  version      = {1.0.0},
+  note         = {Python package available on PyPI at \url{https://pypi.org/project/hard-negatives/}}
+}
+```

hard_negatives-0.1.0/src/hard_negatives.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+src/hard_negatives/data_loader.py
+src/hard_negatives/mine_negatives.py
+src/hard_negatives/prepare_model.py
+src/hard_negatives.egg-info/PKG-INFO
+src/hard_negatives.egg-info/SOURCES.txt
+src/hard_negatives.egg-info/dependency_links.txt
+src/hard_negatives.egg-info/requires.txt
+src/hard_negatives.egg-info/top_level.txt
+src/hard_negatives/args/voyager_args.py
+tests/test_mining.py

hard_negatives-0.1.0/src/hard_negatives.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

hard_negatives-0.1.0/src/hard_negatives.egg-info/requires.txt

@@ -0,0 +1,7 @@
+accelerate>=1.14.0
+datasets>=5.0.0
+flash-attn>=2.8.3.post1
+psutil>=7.2.2
+sentence-transformers>=5.6.0
+torch>=2.13.0
+voyager>=2.1.0

hard_negatives-0.1.0/src/hard_negatives.egg-info/top_level.txt

@@ -0,0 +1 @@
+hard_negatives