File size: 4,667 Bytes
aa559f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Build the SimCSE training corpus and a labelled eval set from FW_cleaned_data_new_base_15.csv.
Only Make + Model are used, as requested.
Outputs (written to simcse/data/):
train.jsonl - {"sentence1": s, "sentence2": s} unsupervised SimCSE pairs
eval.jsonl - {"sentence1", "sentence2", "score"} labelled pairs for EmbeddingSimilarityEvaluator
groups.json - canonical string -> list of raw surface forms (kept for inspection / later use)
The labels come from the canonical columns already present in the CSV: two raw
"<Make> <Model>" strings that normalize to the same "<MAKE> <NEW_MODELS>" are a
positive pair (score 1.0); strings from different canonical groups are negatives
(score 0.0). Nothing but Make/Model text ever reaches the model -- the canonical
columns are used only to build labels.
"""
import json
import random
import re
from collections import defaultdict
from pathlib import Path
import pandas as pd
CSV = Path(__file__).resolve().parents[1] / "data" / "FW_cleaned_data_new_base_15.csv"
OUT = Path(__file__).resolve().parent / "data"
SEED = 42
N_NEG = 2000 # negative eval pairs; positives come from the data (~1.3k)
def norm(s: str) -> str:
"""Lowercase, strip punctuation to spaces, collapse whitespace."""
s = str(s).lower().strip()
s = re.sub(r"[^a-z0-9]+", " ", s)
return re.sub(r"\s+", " ", s).strip()
def dedupe_make(make: str, model: str) -> str:
"""Join make+model, dropping a make token repeated at the head of the model.
The raw data contains rows like Make='audi', Model='audi a3'.
"""
make_t, model_t = norm(make).split(), norm(model).split()
i = 0
while i < len(make_t) and i < len(model_t) and make_t[i] == model_t[i]:
i += 1
# only strip if the whole make prefix was echoed and something survives
if i == len(make_t) and len(model_t) > i:
model_t = model_t[i:]
return " ".join(make_t + model_t)
def main():
rng = random.Random(SEED)
OUT.mkdir(parents=True, exist_ok=True)
df = pd.read_csv(CSV, low_memory=False)
df = df.dropna(subset=["Make", "Model", "MAKE", "NEW_MODELS"])
raw = (df["Make"].map(norm) + " " + df["Model"].map(norm)).str.strip()
canon = df.apply(lambda r: dedupe_make(r["MAKE"], r["NEW_MODELS"]), axis=1)
pairs = pd.DataFrame({"raw": raw, "canon": canon}).drop_duplicates()
pairs = pairs[(pairs["raw"] != "") & (pairs["canon"] != "")]
groups = defaultdict(set)
for r, c in pairs.itertuples(index=False):
groups[c].add(r)
groups = {c: sorted(v) for c, v in groups.items()}
# ---- training corpus -------------------------------------------------
# One sentence per canonical group plus every distinct surface form. Using
# the canonical string too lets the model see the clean form it should map
# everything toward. Duplicates are dropped because SimCSE's in-batch
# negatives would otherwise treat identical strings as negatives.
corpus = sorted({s for forms in groups.values() for s in forms} | set(groups))
rng.shuffle(corpus)
with (OUT / "train.jsonl").open("w") as f:
for s in corpus:
f.write(json.dumps({"sentence1": s, "sentence2": s}) + "\n")
# ---- labelled eval set ----------------------------------------------
positives = []
for c, forms in groups.items():
variants = sorted(set(forms) | {c})
for i in range(len(variants)):
for j in range(i + 1, len(variants)):
positives.append((variants[i], variants[j], 1.0))
canon_list = sorted(groups)
negatives, seen = [], set()
while len(negatives) < N_NEG:
a, b = rng.sample(canon_list, 2)
sa = rng.choice(sorted(set(groups[a]) | {a}))
sb = rng.choice(sorted(set(groups[b]) | {b}))
key = tuple(sorted((sa, sb)))
if sa == sb or key in seen:
continue
seen.add(key)
negatives.append((sa, sb, 0.0))
eval_rows = positives + negatives
rng.shuffle(eval_rows)
with (OUT / "eval.jsonl").open("w") as f:
for s1, s2, score in eval_rows:
f.write(json.dumps({"sentence1": s1, "sentence2": s2, "score": score}) + "\n")
with (OUT / "groups.json").open("w") as f:
json.dump(groups, f, indent=1, sort_keys=True)
print(f"canonical groups : {len(groups)}")
print(f"train sentences : {len(corpus)} -> {OUT / 'train.jsonl'}")
print(f"eval pairs : {len(eval_rows)} ({len(positives)} pos / {len(negatives)} neg)"
f" -> {OUT / 'eval.jsonl'}")
print("\nsample train:", corpus[:5])
print("sample eval :", eval_rows[:3])
if __name__ == "__main__":
main()
|