File size: 4,627 Bytes
ead1d01 | 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 123 124 125 126 127 128 129 130 131 132 133 | """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()
|