data preparation script simcse

aa559f7 verified 2 days ago

4.67 kB

"""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()

	"""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()