wehub-resource-sync
358 lines
13 KiB
Python
358 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Mine hard negatives using the local search API.
|
|
|
|
For each (query, positive_chunk) training pair, queries the search API to find
|
|
top-K retrieval results. Non-positive results become hard negatives — these are
|
|
the actual mistakes the retrieval system makes in production.
|
|
|
|
Usage:
|
|
python mine_hard_negatives.py \
|
|
--input training/data/train.jsonl \
|
|
--output training/data/train_hn.jsonl \
|
|
--num-negatives 7
|
|
|
|
# With more candidates (slower but better negatives)
|
|
python mine_hard_negatives.py \
|
|
--input training/data/train.jsonl \
|
|
--output training/data/train_hn.jsonl \
|
|
--num-negatives 7 --n-docs 50
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
import time
|
|
|
|
import requests
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
stream=sys.stdout,
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SEARCH_URL = "http://localhost:30888/search"
|
|
|
|
|
|
def search_batch(
|
|
queries: list[str], n_docs: int = 20, nprobe: int = 128
|
|
) -> list[list[dict]]:
|
|
"""Query the search API with a batch of text queries."""
|
|
payload = {
|
|
"queries": [{"text": q} for q in queries],
|
|
"n_docs": n_docs,
|
|
"nprobe": nprobe,
|
|
}
|
|
resp = requests.post(SEARCH_URL, json=payload, timeout=120)
|
|
resp.raise_for_status()
|
|
results = resp.json()["results"]
|
|
return [r["hits"] for r in results]
|
|
|
|
|
|
def mine_from_search(
|
|
pairs: list[dict],
|
|
num_negatives: int = 7,
|
|
n_docs: int = 20,
|
|
batch_size: int = 64,
|
|
nprobe: int = 128,
|
|
filter_mode: str = "none",
|
|
margin: float | None = None,
|
|
) -> list[dict]:
|
|
"""Mine hard negatives by querying the search API.
|
|
|
|
For each pair, the search API returns top-K results. Any result whose path
|
|
differs from the positive chunk_path is a hard negative.
|
|
|
|
Args:
|
|
filter_mode: How to filter likely false negatives:
|
|
- "article": Skip chunks from the same article (by article_id).
|
|
Keeps hard negatives that score higher than positive.
|
|
- "margin": Skip negatives where neg_score > pos_score * margin.
|
|
Conservative, but may discard useful hard cases.
|
|
- "skip_top1": Skip the #1 non-positive result (likely false negative),
|
|
keep #2-#K. Compromise between article and margin.
|
|
- "none": No filtering, keep all non-positive hits.
|
|
margin: Only used when filter_mode="margin". Typical: 0.95-0.98.
|
|
"""
|
|
# Deduplicate queries for efficiency
|
|
unique_queries = list(dict.fromkeys(p["query"] for p in pairs))
|
|
query_to_idx = {q: i for i, q in enumerate(unique_queries)}
|
|
logger.info(f"{len(unique_queries)} unique queries (from {len(pairs)} pairs)")
|
|
|
|
# Collect all positive paths per unique query
|
|
query_positives = {}
|
|
for pair in pairs:
|
|
q = pair["query"]
|
|
if q not in query_positives:
|
|
query_positives[q] = set()
|
|
query_positives[q].add(pair["chunk_path"])
|
|
|
|
# Query search API in batches
|
|
all_hits = [None] * len(unique_queries)
|
|
n_batches = (len(unique_queries) + batch_size - 1) // batch_size
|
|
|
|
for i in range(0, len(unique_queries), batch_size):
|
|
batch_queries = unique_queries[i : i + batch_size]
|
|
batch_idx = i // batch_size + 1
|
|
|
|
try:
|
|
batch_hits = search_batch(batch_queries, n_docs=n_docs, nprobe=nprobe)
|
|
for j, hits in enumerate(batch_hits):
|
|
all_hits[i + j] = hits
|
|
except Exception as e:
|
|
logger.warning(f"Batch {batch_idx}/{n_batches} failed: {e}")
|
|
for j in range(len(batch_queries)):
|
|
all_hits[i + j] = []
|
|
|
|
done = min(i + batch_size, len(unique_queries))
|
|
if batch_idx % 10 == 0 or done == len(unique_queries):
|
|
logger.info(f" Searched: {done}/{len(unique_queries)}")
|
|
|
|
# Build set of positive article_ids per query for false-negative filtering
|
|
query_pos_articles = {}
|
|
for q in unique_queries:
|
|
pos_paths = query_positives[q]
|
|
# Extract article_id from hits that match positive paths
|
|
hits = all_hits[query_to_idx[q]] or []
|
|
article_ids = set()
|
|
for hit in hits:
|
|
if hit.get("path", "") in pos_paths:
|
|
article_ids.add(hit.get("article_id"))
|
|
query_pos_articles[q] = article_ids
|
|
|
|
# Extract hard negatives per unique query
|
|
query_negatives = {}
|
|
query_metadata = {}
|
|
stats = {
|
|
"total": 0,
|
|
"with_negs": 0,
|
|
"avg_negs": 0,
|
|
"avg_pos_rank": 0,
|
|
"same_article_filtered": 0,
|
|
"margin_filtered": 0,
|
|
"skip_top1_filtered": 0,
|
|
"harder_than_pos": 0,
|
|
}
|
|
pos_ranks = []
|
|
pos_rank_distribution = {str(i): 0 for i in range(1, n_docs + 1)}
|
|
pos_rank_distribution[f">{n_docs}"] = 0
|
|
|
|
for q in unique_queries:
|
|
positives = query_positives[q]
|
|
pos_article_ids = query_pos_articles[q]
|
|
hits = all_hits[query_to_idx[q]] or []
|
|
|
|
# Find positive score for stats
|
|
pos_score = None
|
|
for hit in hits:
|
|
if hit.get("path", "") in positives:
|
|
pos_score = hit["score"]
|
|
break
|
|
|
|
# Find negatives: exclude positives and same-article chunks (likely false negatives)
|
|
neg_paths = []
|
|
pos_rank = None
|
|
for rank, hit in enumerate(hits):
|
|
hit_path = hit.get("path", "")
|
|
if hit_path in positives:
|
|
if pos_rank is None:
|
|
pos_rank = rank
|
|
else:
|
|
# Filter likely false negatives
|
|
if filter_mode == "article":
|
|
if hit.get("article_id") in pos_article_ids:
|
|
stats["same_article_filtered"] += 1
|
|
continue
|
|
elif (
|
|
filter_mode == "margin"
|
|
and margin is not None
|
|
and pos_score is not None
|
|
):
|
|
if hit["score"] > pos_score * margin:
|
|
stats["margin_filtered"] += 1
|
|
continue
|
|
elif filter_mode == "skip_top1":
|
|
if len(neg_paths) == 0 and rank < 5:
|
|
# Skip the first non-positive hit (likely false negative)
|
|
stats["skip_top1_filtered"] += 1
|
|
continue
|
|
# filter_mode == "none": no filtering
|
|
|
|
if len(neg_paths) < num_negatives:
|
|
neg_paths.append(hit_path)
|
|
if pos_score is not None and hit["score"] >= pos_score:
|
|
stats["harder_than_pos"] += 1
|
|
|
|
query_negatives[q] = neg_paths
|
|
query_metadata[q] = {
|
|
"retrieve_top20": [
|
|
{
|
|
"rank": rank + 1,
|
|
"path": hit.get("path", ""),
|
|
"score": hit.get("score", 0.0),
|
|
"article_id": hit.get("article_id"),
|
|
}
|
|
for rank, hit in enumerate(hits)
|
|
],
|
|
"positive_score": pos_score if pos_score is not None else 0.0,
|
|
"positive_rank": pos_rank + 1 if pos_rank is not None else 0,
|
|
}
|
|
stats["total"] += 1
|
|
if neg_paths:
|
|
stats["with_negs"] += 1
|
|
if pos_rank is not None:
|
|
pos_ranks.append(pos_rank)
|
|
pos_rank_distribution[str(pos_rank + 1)] += 1
|
|
else:
|
|
pos_rank_distribution[f">{n_docs}"] += 1
|
|
|
|
# Write output
|
|
output_pairs = []
|
|
for pair in pairs:
|
|
neg_paths = query_negatives.get(pair["query"], [])
|
|
meta = query_metadata.get(pair["query"], {})
|
|
output_pairs.append(
|
|
{
|
|
**pair,
|
|
"neg_chunk_paths": neg_paths,
|
|
"retrieve_top20": meta.get("retrieve_top20", []),
|
|
"positive_score": meta.get("positive_score", 0.0),
|
|
"positive_rank": meta.get("positive_rank", 0),
|
|
}
|
|
)
|
|
|
|
# Stats
|
|
total_queries = len(unique_queries)
|
|
stats["avg_negs"] = (
|
|
sum(len(query_negatives[q]) for q in unique_queries) / total_queries
|
|
)
|
|
if pos_ranks:
|
|
stats["avg_pos_rank"] = sum(pos_ranks) / len(pos_ranks)
|
|
stats["pos_found_rate"] = len(pos_ranks) / total_queries
|
|
stats["pos_recall@1"] = sum(1 for r in pos_ranks if r == 0) / total_queries
|
|
stats["pos_recall@10"] = sum(1 for r in pos_ranks if r < 10) / total_queries
|
|
stats["pos_recall@20"] = sum(1 for r in pos_ranks if r < 20) / total_queries
|
|
stats["pos_rank_distribution"] = pos_rank_distribution
|
|
|
|
return output_pairs, stats
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--input", required=True, help="Input JSONL with {query, chunk_path}"
|
|
)
|
|
parser.add_argument(
|
|
"--output", required=True, help="Output JSONL with added neg_chunk_paths"
|
|
)
|
|
parser.add_argument("--num-negatives", type=int, default=7)
|
|
parser.add_argument(
|
|
"--n-docs",
|
|
type=int,
|
|
default=20,
|
|
help="Number of docs to retrieve per query from search API",
|
|
)
|
|
parser.add_argument(
|
|
"--filter-mode",
|
|
choices=["article", "margin", "skip_top1", "none"],
|
|
default="none",
|
|
help="False-negative filter: 'article' (skip same article_id), "
|
|
"'margin' (skip by score margin), 'none' (no filter)",
|
|
)
|
|
parser.add_argument(
|
|
"--margin",
|
|
type=float,
|
|
default=0.95,
|
|
help="Margin threshold (only for --filter-mode margin). Typical: 0.95-0.98",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size", type=int, default=64, help="Batch size for search API queries"
|
|
)
|
|
parser.add_argument(
|
|
"--chunk-path-prefix",
|
|
type=str,
|
|
default="/opt/dlami/nvme/kiwix_tiles/",
|
|
help="Prefix to prepend to relative chunk_path for matching search API results",
|
|
)
|
|
parser.add_argument(
|
|
"--nprobe", type=int, default=128, help="FAISS nprobe for search API"
|
|
)
|
|
parser.add_argument(
|
|
"--stats-output",
|
|
type=str,
|
|
default=None,
|
|
help="Optional JSON path to write mining stats",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Check search API
|
|
try:
|
|
resp = requests.get("http://localhost:30888/health", timeout=5)
|
|
resp.raise_for_status()
|
|
logger.info("Search API is healthy")
|
|
except Exception as e:
|
|
logger.error(f"Search API not available: {e}")
|
|
sys.exit(1)
|
|
|
|
# Load data
|
|
pairs = []
|
|
with open(args.input) as f:
|
|
for line in f:
|
|
pair = json.loads(line)
|
|
# Normalize chunk_path to absolute path for matching search API results
|
|
if args.chunk_path_prefix and not pair["chunk_path"].startswith("/"):
|
|
pair["chunk_path"] = args.chunk_path_prefix + pair["chunk_path"]
|
|
pairs.append(pair)
|
|
logger.info(f"Loaded {len(pairs)} pairs")
|
|
|
|
# Mine
|
|
t0 = time.time()
|
|
output_pairs, stats = mine_from_search(
|
|
pairs,
|
|
num_negatives=args.num_negatives,
|
|
n_docs=args.n_docs,
|
|
batch_size=args.batch_size,
|
|
nprobe=args.nprobe,
|
|
filter_mode=args.filter_mode,
|
|
margin=args.margin,
|
|
)
|
|
elapsed = time.time() - t0
|
|
|
|
# Write
|
|
with open(args.output, "w") as f:
|
|
for pair in output_pairs:
|
|
f.write(json.dumps(pair) + "\n")
|
|
|
|
n_with_negs = sum(1 for p in output_pairs if p["neg_chunk_paths"])
|
|
logger.info(f"Wrote {len(output_pairs)} pairs to {args.output}")
|
|
logger.info(
|
|
f" {n_with_negs} with negatives ({n_with_negs / len(output_pairs):.1%})"
|
|
)
|
|
logger.info(f" Avg negatives per query: {stats['avg_negs']:.1f}")
|
|
logger.info(f" Avg positive rank: {stats.get('avg_pos_rank', 'N/A')}")
|
|
if "pos_recall@1" in stats:
|
|
logger.info(f" Search API recall@1: {stats['pos_recall@1']:.3f}")
|
|
logger.info(f" Search API recall@10: {stats['pos_recall@10']:.3f}")
|
|
logger.info(f" Search API recall@20: {stats['pos_recall@20']:.3f}")
|
|
if stats.get("same_article_filtered"):
|
|
logger.info(f" Same-article filtered: {stats['same_article_filtered']}")
|
|
if stats.get("margin_filtered"):
|
|
logger.info(f" Margin-filtered: {stats['margin_filtered']}")
|
|
if stats.get("skip_top1_filtered"):
|
|
logger.info(f" Skip-top1 filtered: {stats['skip_top1_filtered']}")
|
|
if stats.get("harder_than_pos"):
|
|
logger.info(f" Negatives harder than positive: {stats['harder_than_pos']}")
|
|
logger.info(f" Positive rank distribution: {stats['pos_rank_distribution']}")
|
|
logger.info(f" Time: {elapsed:.0f}s")
|
|
|
|
if args.stats_output:
|
|
with open(args.stats_output, "w") as f:
|
|
json.dump(stats, f, indent=2, sort_keys=True)
|
|
logger.info(f"Wrote stats to {args.stats_output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|