import polars as pl
import numpy as np
from dupehell import generate
from pathlib import Path
import itertools
import random
import sys, io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
DOMAINS = ["kyc", "healthcare", "fintech", "publishing", "ecommerce"]
DIFFICULTIES = ["light", "medium", "hell"]
HN_RATIOS = [0.05, 0.1, 0.2]
SEED = 42
SIZE = 2000
MAX_SAMPLE_PAIRS = 50
OUTPUT_DIR = Path("hn_quality_output")
OUTPUT_DIR.mkdir(exist_ok=True)
def load_polars(result):
df = pl.read_ipc(result.dataset) if str(result.dataset).endswith(".ipc") else pl.read_parquet(result.dataset)
gt = pl.read_ipc(result.ground_truth) if str(result.ground_truth).endswith(".ipc") else pl.read_parquet(result.ground_truth)
return df, gt
def get_content_cols(df):
exclude = {"record_id", "master_id", "entity_type", "domain"}
return [c for c in df.columns if c.lower() not in exclude]
def similarity_score(va, vb):
if va is None and vb is None:
return 1.0
if va is None or vb is None:
return 0.0
try:
return 1.0 if va == vb else 0.0
except Exception:
return 1.0 if str(va) == str(vb) else 0.0
def pair_similarity(row1, row2, content_cols):
scores = [similarity_score(row1[c], row2[c]) for c in content_cols]
return float(np.mean(scores)) if scores else 0.0
def build_record_map(df):
m = {}
for row in df.iter_rows(named=True):
m[row["record_id"]] = row
return m
def form_exact_dup_pairs(df, gt, max_pairs=50):
exact_records = gt.filter(pl.col("match_type") == "exact_dup")["record_id"].to_list()
exact_df = df.filter(pl.col("record_id").is_in(exact_records))
grouped = {}
for row in exact_df.iter_rows(named=True):
mid = row["master_id"]
grouped.setdefault(mid, []).append(row["record_id"])
pairs = []
for mid, rids in grouped.items():
if len(rids) < 2:
continue
for a, b in itertools.combinations(rids, 2):
pairs.append((a, b))
if len(pairs) >= max_pairs:
break
random.shuffle(pairs)
return pairs[:max_pairs]
def form_hard_neg_pairs(df, gt, max_pairs=50):
hard_neg_records = gt.filter(pl.col("match_type") == "hard_neg")["record_id"].to_list()
mid_map = {}
for row in df.filter(pl.col("record_id").is_in(hard_neg_records)).iter_rows(named=True):
mid_map[row["record_id"]] = row["master_id"]
by_mid = {}
for rid, mid in mid_map.items():
by_mid.setdefault(mid, []).append(rid)
mids = list(by_mid.keys())
pairs = []
for i in range(len(mids)):
for j in range(i + 1, len(mids)):
for a in by_mid[mids[i]]:
for b in by_mid[mids[j]]:
pairs.append((a, b))
if len(pairs) >= max_pairs:
break
if len(pairs) >= max_pairs:
break
if len(pairs) >= max_pairs:
break
if len(pairs) >= max_pairs:
break
random.shuffle(pairs)
return pairs[:max_pairs]
def run_domain_test(domain, difficulty, size=SIZE, seed=SEED, hard_neg_ratio=0.3):
print(f" Generating {domain} (diff={difficulty}, size={size}, seed={seed}, hn_ratio={hard_neg_ratio})...")
try:
result = generate(
domain, size=size, seed=seed, difficulty=difficulty,
output_dir=str(OUTPUT_DIR), output_format="ipc",
hard_neg_ratio=hard_neg_ratio,
)
except Exception as e:
print(f" FAILED: {e}")
return None
df, gt = load_polars(result)
content_cols = get_content_cols(df)
record_map = build_record_map(df)
total_records = df.height
match_counts = gt.group_by("match_type").agg(pl.len()).to_pandas().set_index("match_type")["len"].to_dict()
hn_total = match_counts.get("hard_neg", 0)
print(f" Records={total_records}, GT={gt.height}, HN={hn_total}, match_types={match_counts}")
exact_pairs = form_exact_dup_pairs(df, gt, MAX_SAMPLE_PAIRS)
hard_pairs = form_hard_neg_pairs(df, gt, MAX_SAMPLE_PAIRS)
exact_sims = [pair_similarity(record_map[a], record_map[b], content_cols) for a, b in exact_pairs]
hard_sims = [pair_similarity(record_map[a], record_map[b], content_cols) for a, b in hard_pairs]
exact_sim = float(np.mean(exact_sims)) if exact_sims else 0.0
hard_sim = float(np.mean(hard_sims)) if hard_sims else 0.0
gap = exact_sim - hard_sim
status = "✅" if gap > 0.50 else "❌"
return {
"domain": domain,
"difficulty": difficulty,
"hn_ratio": hard_neg_ratio,
"total_records": total_records,
"gt_pairs": gt.height,
"hn_total": hn_total,
"exact_dup_pairs": len(exact_pairs),
"hard_neg_pairs": len(hard_pairs),
"exact_dup_sim": round(exact_sim, 4),
"hard_neg_sim": round(hard_sim, 4),
"gap": round(gap, 4),
"status": status,
}
def main():
results = []
print("=" * 70)
print("PART B: Cross-domain HN quality (medium difficulty)")
print("=" * 70)
for domain in DOMAINS:
r = run_domain_test(domain, "medium")
if r:
results.append(r)
print("\n" + "=" * 70)
print("PART C: Difficulty impact on kyc")
print("=" * 70)
for diff in DIFFICULTIES:
r = run_domain_test("kyc", diff)
if r:
results.append(r)
print("\n" + "=" * 70)
print("PART D: Hard neg ratio impact on kyc (medium)")
print("=" * 70)
for hn_ratio in HN_RATIOS:
r = run_domain_test("kyc", "medium", hard_neg_ratio=hn_ratio)
if r:
results.append(r)
pass_count = sum(1 for r in results if r["status"] == "✅")
fail_count = sum(1 for r in results if r["status"] == "❌")
total = len(results)
print("\n" + "=" * 110)
print(f"## Agent 15.1 (HN Quality) — {pass_count}/{total} ✅ ({fail_count} ❌)")
print("=" * 110)
header = f"| {'Domain':<14} | {'Difficulty':<10} | {'HN ratio':<9} | {'HN count':<9} | {'exact_dup_sim':<13} | {'hard_neg_sim':<13} | {'Gap':<7} | {'Status':<6} |"
sep = f"|{'':-<16}|{'':-<12}|{'':-<11}|{'':-<11}|{'':-<15}|{'':-<15}|{'':-<9}|{'':-<8}|"
print(header)
print(sep)
for r in results:
print(f"| {r['domain']:<14} | {r['difficulty']:<10} | {r['hn_ratio']:<9} | {r['hn_total']:<9} | {r['exact_dup_sim']:<13} | {r['hard_neg_sim']:<13} | {r['gap']:<7} | {r['status']:<6} |")
print("\n### Similarity Gap Summary\n")
for r in results:
sym = "❌" if r["status"] == "❌" else "✅"
print(f" {sym} {r['domain']}/{r['difficulty']} (hn_ratio={r['hn_ratio']}): gap={r['gap']:.4f} (exact_dup_sim={r['exact_dup_sim']:.4f}, hard_neg_sim={r['hard_neg_sim']:.4f})")
fails = [r for r in results if r["status"] == "❌"]
if fails:
print("\n### ❌ FAILURES (gap < 0.50)\n")
for f in fails:
print(f"| **BUG** | {f['domain']}/{f['difficulty']} | gap={f['gap']:.4f} | exact_dup_sim={f['exact_dup_sim']:.4f} | hard_neg_sim={f['hard_neg_sim']:.4f} |")
else:
print("\n✅ No failures — all gaps > 0.50")
print("\n### Difficulty Trend (kyc)\n")
for diff in DIFFICULTIES:
rs = [r for r in results if r["domain"] == "kyc" and r["difficulty"] == diff and r["hn_ratio"] == 0.3]
if rs:
r = rs[0]
print(f" {diff:>8}: gap={r['gap']:.4f} (exact={r['exact_dup_sim']:.4f}, hard_neg={r['hard_neg_sim']:.4f})")
print("\n### HN Ratio Impact (kyc, medium)\n")
for hn_ratio in HN_RATIOS + [0.3]:
rs = [r for r in results if r["domain"] == "kyc" and r["difficulty"] == "medium" and r["hn_ratio"] == hn_ratio]
if rs:
r = rs[0]
expected = int(SIZE * hn_ratio * 0.05)
print(f" hn_ratio={hn_ratio:>4}: expected HN ~{expected:>4}, actual HN={r['hn_total']:>4}, exact_pairs={r['exact_dup_pairs']}, hard_pairs={r['hard_neg_pairs']}")
print("\nDone.")
if __name__ == "__main__":
main()