training script simcse

train_simcse.py

@@ -0,0 +1,132 @@
+"""Unsupervised SimCSE on the FW make/model strings.
+
+This is the simcse_script.ipynb pipeline pointed at our own data instead of
+wiki1m/stsb. SimCSE's positive pair is the *same* sentence encoded twice -- the
+two views differ only through dropout -- and MultipleNegativesRankingLoss
+supplies the negatives from the rest of the batch.
+
+Run prepare_data.py first.
+
+  python simcse/train_simcse.py                      # defaults
+  python simcse/train_simcse.py --epochs 10 --batch-size 30
+"""
+
+import argparse
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from xml.parsers.expat import model
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
+from sentence_transformers.losses import MultipleNegativesRankingLoss
+from sentence_transformers.models import Pooling, Transformer
+from sentence_transformers.similarity_functions import SimilarityFunction
+from sentence_transformers.trainer import SentenceTransformerTrainer
+from sentence_transformers.training_args import SentenceTransformerTrainingArguments
+
+HERE = Path(__file__).resolve().parent
+DATA = HERE / "data"
+
+logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    with path.open() as f:
+        return [json.loads(line) for line in f]
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-name", default="distilbert/distilbert-base-uncased")
+    p.add_argument("--batch-size", type=int, default=30)
+    p.add_argument("--epochs", type=int, default=10)
+    # make/model strings are a handful of tokens; 32 is already generous
+    p.add_argument("--max-seq-length", type=int, default=32)
+    p.add_argument("--learning-rate", type=float, default=3e-5)
+    p.add_argument("--output-dir", default=None)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    out = Path(args.output_dir) if args.output_dir else HERE / "output" / (
+        "simcse-make-model-{}-bs{}-{}".format(
+            args.model_name.split("/")[-1],
+            args.batch_size,
+            datetime.now().strftime("%Y-%m-%d-%H-%M-%S"),
+        )
+    )
+
+    # ---- model ----------------------------------------------------------
+    word_embedding_model = Transformer(args.model_name, max_seq_length=args.max_seq_length)
+    pooling_model = Pooling(word_embedding_model.get_sentence_embedding_dimension(), "mean")
+    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
+
+    # ---- data -----------------------------------------------------------
+    train_rows = read_jsonl(DATA / "train.jsonl")
+    train_dataset = Dataset.from_list(train_rows)
+
+    eval_rows = read_jsonl(DATA / "eval.jsonl")
+
+    evaluator = EmbeddingSimilarityEvaluator(
+        sentences1=[r["sentence1"] for r in eval_rows],
+        sentences2=[r["sentence2"] for r in eval_rows],
+        scores=[r["score"] for r in eval_rows],
+        main_similarity=SimilarityFunction.COSINE,
+        name="make-model-dev",
+    )
+
+    logging.info("train sentences=%d  eval pairs=%d", len(train_dataset), len(eval_rows))
+
+    loss = MultipleNegativesRankingLoss(model)
+
+    logging.info("performance before training")
+    before = evaluator(model)
+    logging.info("%s", before)
+
+    # ---- train ----------------------------------------------------------
+    training_args = SentenceTransformerTrainingArguments(
+        output_dir=str(out),
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.batch_size,
+        warmup_ratio=0.1,
+        learning_rate=args.learning_rate,
+        eval_strategy="steps",
+        eval_steps=0.1,
+        logging_steps=0.05,
+        save_strategy="no",
+        fp16=True,
+        # identical strings in one batch become false negatives under MNRL
+        batch_sampler="no_duplicates",
+        report_to=[],
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        evaluator=evaluator,
+        loss=loss,
+    )
+    trainer.train()
+
+    after = evaluator(model)
+    logging.info("performance after training")
+    logging.info("%s", after)
+
+    final = out.parent / (out.name + "-final")
+    model.save_pretrained(str(final))
+    logging.info("saved -> %s", final)
+
+    key = "make-model-dev_spearman_cosine"
+    print("\n=== dev spearman (cosine) ===")
+    print(f"  before : {before.get(key):.4f}")
+    print(f"  after  : {after.get(key):.4f}")
+    print(f"  model  : {final}")
+   
+
+
+if __name__ == "__main__":
+    main()