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