hypernet_sp/evals/specificity_train.py

"""Train the INDEPENDENT specificity head: a single linear probe on frozen BGE-small embeddings that
fires on 'future-relevant proper/specific information' (proper nouns, codes, prices, dates, numeric
values) and stays quiet on generic words/phrases. This is the elegant winner from the bake-off
(evals/specificity_probe.py): BGE linear probe AUC ~0.96, beating tokenizer fragmentation (0.87) and
killing the form-frequency proxy (0.30). Reuses the SAME BGE we already run for routing + retrieval —
the probe is just one more linear head on the existing 384-d vector.

Output: evals/specificity_clf.joblib  {"clf": LogisticRegression, "labels": ["generic","specific"]}.
Optionally merges evals/specificity_gen.jsonl (open-model-generated spans), same as the intent pipeline.
Run: python3.12 evals/specificity_train.py
"""
import sys, os, json
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "runtime"))
import numpy as np

# pin-worthy spans: specific VALUES / proper entities the user might introduce and need recalled later
SPECIFIC = [
 "$120", "$500", "$2.50", "$1,200", "14C", "B3", "level B3", "bay 12", "spot 47", "gate 22",
 "QX7-2291", "POL-55821", "EMP-90832", "5588", "555-0199", "hunter2", "abc123", "O negative",
 "Naruhito", "Sapporo", "Mochi", "Helsinki", "Apollo", "Mei", "Hilton", "Mia", "Aki", "Toyota",
 "42 Oak Street", "3pm", "4pm", "Friday", "Monday", "next Monday", "80 cm", "7 years old",
 "15% off", "20 percent", "room 305", "flight DL472", "seat 14C", "PIN 4417", "model X100",
 "the Helsinki office", "building 7", "aisle 9", "track 3", "$45 each", "version 2.1",
 "Dr. Tanaka", "the Apollo project", "March 3rd", "9:30am", "2024-01-15", "carriage 6",
 "ticket #88231", "wifi password hunter2", "policy POL-55821", "code QX7-2291",
]
# generic words / phrases / common nouns / fillers — nothing worth pinning
GENERIC = [
 "the", "planning", "explain", "weekend", "trip", "theme", "party", "idea", "help", "something",
 "really", "maybe", "about", "really nice", "a little", "how are you", "good morning", "let me",
 "i think", "the meeting", "my budget", "the total", "some ideas", "the cake", "the venue",
 "the project", "the room", "the office", "a few", "later today", "this morning", "the number",
 "the price", "the cost", "the color", "the time", "the date", "the place", "the food", "the plan",
 "do it again", "thanks a lot", "sounds good", "no problem", "of course", "for sure", "right now",
 "the budget", "the schedule", "the guest list", "the order", "the booking", "the reservation",
 "make it bigger", "a good option", "something cheaper", "the next step", "the first one",
 "tell me more", "what do you think", "any suggestions", "the usual", "as before", "the same",
]


def main():
    from rag import BGERetriever
    bge = BGERetriever()
    spec = list(dict.fromkeys(SPECIFIC))
    gen = list(dict.fromkeys(GENERIC))
    gen_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "specificity_gen.jsonl")
    ngen = 0
    if os.path.exists(gen_path) and "--hand-only" not in sys.argv:
        have = set(spec) | set(gen)
        for l in open(gen_path):
            o = json.loads(l); t = o["text"]
            if t in have:
                continue
            (spec if o["label"] == "specific" else gen).append(t); have.add(t); ngen += 1
    print(f"merged {ngen} generated spans")

    texts = spec + gen
    y = np.array([1] * len(spec) + [0] * len(gen))
    X = bge._encode(texts, is_query=False)
    print(f"dataset: {len(y)} spans | specific={int(y.sum())} generic={int((1-y).sum())}")

    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import cross_val_predict, StratifiedKFold
    from sklearn.metrics import classification_report, roc_auc_score
    clf = LogisticRegression(max_iter=2000, C=2.0, class_weight="balanced")
    proba = cross_val_predict(clf, X, y, cv=StratifiedKFold(5, shuffle=True, random_state=0),
                              method="predict_proba")[:, 1]
    pred = (proba >= 0.5).astype(int)
    print(f"\n5-fold CV  AUC={roc_auc_score(y, proba):.3f}")
    print(classification_report(y, pred, target_names=["generic", "specific"], digits=3))
    print("CV mistakes:")
    for t, yt, pp in zip(texts, y, proba):
        if (pp >= 0.5) != bool(yt):
            print(f"  {'SPEC' if yt else 'gen '}->{pp:.2f}: {t}")

    clf.fit(X, y)
    import joblib
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "specificity_clf.joblib")
    joblib.dump({"clf": clf, "labels": ["generic", "specific"]}, out)
    print("saved", out)
    print("SPECIFICITY_TRAIN_DONE")


if __name__ == "__main__":
    main()