get_align_disc_sample.py

#!/usr/bin/env python3
import os
import numpy as np
import pickle as pkl
from tqdm import tqdm
from collections import defaultdict
import random
import argparse
from rapidfuzz.distance import Levenshtein

def float_to_str(f: float) -> str:
    return f'{int(f * 10):02d}'

def edit_distance(seq1, seq2) -> float:
    # normalized in [0,1]
    return Levenshtein.normalized_distance(seq1, seq2)

def safe_normalize(weights):
    s = float(sum(weights))
    if s <= 0:
        return [0.0 for _ in weights]
    return [w / s for w in weights]

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, required=True)
    parser.add_argument('--alpha', type=float, required=True)
    parser.add_argument('--beta', type=float, required=True)
    parser.add_argument('--gamma', type=float, required=True)
    parser.add_argument('--max_len', type=int, default=20)
    parser.add_argument('--k', type=int, default=20, help='samples per (target, position)')
    parser.add_argument('--neg_targets', type=int, default=20, help='# negative targets per test sequence')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--quiet', action='store_true')
    args = parser.parse_args()

    # Reproducibility
    random.seed(args.seed)
    np.random.seed(args.seed)

    with open(f'data_pkl/{args.dataset}_seq.pkl', 'rb') as f:
        train_sequences, test_sequences = pkl.load(f)

    train_sequences = [seq.tolist() for seq in train_sequences]
    test_sequences = [seq.tolist() for seq in test_sequences]

    # (Optional) validation behavior — keep as False to match your code
    validation = False
    if validation:
        test_sequences = [seq[:-1] for seq in test_sequences]

    # Sequence weights, normalized
    seq_weights = [(max(len(seq) - 1, 0)) ** args.alpha for seq in train_sequences]
    seq_weights = safe_normalize(seq_weights)

    # Precompute per-sequence position weights and accumulate target positions
    target2positions = defaultdict(list)  # target -> list of (seq_idx, target_index)
    target2weights   = defaultdict(list)  # target -> list of combined weights

    it = range(len(train_sequences))
    if not args.quiet:
        it = tqdm(it, desc='Index train')

    for seq_idx in it:
        seq = train_sequences[seq_idx]
        L = len(seq)
        if L <= 1:
            continue

        # Position weights
        if args.beta > 10:
            pos_w = [0.0] * (L - 1) + [1.0]
        else:
            pos_w = [0.0] + [i ** args.beta for i in range(1, L)]
            pos_w = safe_normalize(pos_w)

        sw = seq_weights[seq_idx]
        if sw == 0:
            continue

        # Only positions 1..L-1 are targets (consistent with your logic)
        for target_index in range(1, L):
            w = sw * pos_w[target_index]
            if w > 0.0:
                target = seq[target_index]
                target2positions[target].append((seq_idx, target_index))
                target2weights[target].append(w)

    all_targets = list(target2positions.keys())

    P, N = [], []
    it_test = test_sequences
    if not args.quiet:
        it_test = tqdm(test_sequences, desc='Evaluate')

    for test_seq in it_test:
        if len(test_seq) < 2:
            continue
        inp, tgt = test_seq[:-1], test_seq[-1]
        inp = inp[-args.max_len:]

        # Skip if no training positions for positive target
        if tgt not in target2positions or not target2positions[tgt]:
            continue

        # ----- POSITIVE -----
        pos_idx_pool = list(range(len(target2positions[tgt])))
        pos_samples = random.choices(
            pos_idx_pool,
            weights=target2weights[tgt],
            k=args.k
        )

        pos_num = 0.0
        pos_den = 0.0
        for i in pos_samples:
            seq_idx, target_index = target2positions[tgt][i]
            base_w = target2weights[tgt][i]

            # start window for this position
            lo = max(0, target_index - args.max_len)
            if lo >= target_index:
                continue  # empty window

            # start weights
            if args.gamma < -10:
                start_indices = [lo]
                start_weights = [1.0]
            else:
                k = target_index - lo
                start_indices = list(range(lo, target_index))
                start_weights = safe_normalize([(j + 1) ** args.gamma for j in range(k)])

            # if degenerate (all zeros), skip
            if not start_weights or sum(start_weights) == 0:
                continue

            # sample a start index
            start_index = random.choices(start_indices, weights=start_weights, k=1)[0]
            # convert to local index to fetch the weight again
            local_idx = start_index - lo
            st_w = start_weights[local_idx]

            train_input = train_sequences[seq_idx][start_index:target_index]
            if not train_input:
                continue

            dist = edit_distance(inp, train_input)
            sim = 1.0 - dist

            w = base_w * st_w
            pos_num += sim * w
            pos_den += w

        if pos_den == 0:
            # no valid positive sample from this test case
            continue
        pos_sim = pos_num / pos_den

        # ----- NEGATIVE -----
        neg_pool = [t for t in all_targets if t != tgt]
        if not neg_pool:
            # no negatives possible; skip this test case
            continue
        k_neg = min(args.neg_targets, len(neg_pool))
        neg_targets = random.sample(neg_pool, k_neg)

        neg_scores = []
        for nt in neg_targets:
            idx_pool = list(range(len(target2positions[nt])))
            if not idx_pool:
                continue
            neg_samples = random.choices(
                idx_pool,
                weights=target2weights[nt],
                k=args.k
            )

            neg_num = 0.0
            neg_den = 0.0
            for i in neg_samples:
                seq_idx, target_index = target2positions[nt][i]
                base_w = target2weights[nt][i]

                lo = max(0, target_index - args.max_len)
                if lo >= target_index:
                    continue

                if args.gamma < -10:
                    start_indices = [lo]
                    start_weights = [1.0]
                else:
                    k = target_index - lo
                    start_indices = list(range(lo, target_index))
                    start_weights = safe_normalize([(j + 1) ** args.gamma for j in range(k)])

                if not start_weights or sum(start_weights) == 0:
                    continue

                start_index = random.choices(start_indices, weights=start_weights, k=1)[0]
                local_idx = start_index - lo
                st_w = start_weights[local_idx]

                train_input = train_sequences[seq_idx][start_index:target_index]
                if not train_input:
                    continue

                dist = edit_distance(inp, train_input)
                sim = 1.0 - dist

                w = base_w * st_w
                neg_num += sim * w
                neg_den += w

            if neg_den > 0:
                neg_scores.append(neg_num / neg_den)

        if neg_scores:
            P.append(pos_sim)
            N.append(float(np.mean(neg_scores)))

    mean_P = float(np.nanmean(P)) if P else float('nan')
    mean_N = float(np.nanmean(N)) if N else float('nan')

    print("\n=== Similarity Report ===")
    print(f" dataset   : {args.dataset}")
    print(f" alpha     : {args.alpha:.4f}")
    print(f" beta      : {args.beta:.4f}")
    print(f" gamma     : {args.gamma:.4f}")
    print(f" max_len   : {args.max_len}")
    print(f" k         : {args.k}")
    print(f" neg_tgts  : {args.neg_targets}")
    print(f" P (mean)  : {mean_P:.6f}")
    print(f" N (mean)  : {mean_N:.6f}")
    if np.isfinite(mean_P) and np.isfinite(mean_N) and mean_N != 0:
        print(f" P/N ratio : {mean_P / mean_N:.6f}")
    print("=========================\n")

if __name__ == "__main__":
    main()