# static-embedding-japanese trainer.py # base: https://huggingface.co./blog/static-embeddings # MIT License import logging import os import random from pathlib import Path from sentence_transformers import ( SentenceTransformer, SentenceTransformerModelCardData, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, ) from sentence_transformers.evaluation import NanoBEIREvaluator from sentence_transformers.losses import MatryoshkaLoss, MultipleNegativesRankingLoss from sentence_transformers.models.StaticEmbedding import StaticEmbedding from sentence_transformers.training_args import BatchSamplers, MultiDatasetBatchSamplers from transformers import AutoTokenizer from datasets import Dataset, DatasetDict, load_dataset EXP = "030" print("EXP:", EXP) PROJECT_ROOT = Path(__file__).resolve().parents[1] print(PROJECT_ROOT) EN_TARGET_DATASETS = [ # "gooaq", # non-commarical "msmarco", "squad", # "s2orc", # large "allnli", # "paq", # large "trivia_qa", # "msmarco_10m", "swim_ir", # "pubmedqa", "miracl", # "mldr", # non-commarical "mr_tydi", ] JA_TARGET_DATASETS = [ "hpprc_emb__auto-wiki-nli-triplet", "hpprc_emb__auto-wiki-qa", "hpprc_emb__auto-wiki-qa-nemotron", "hpprc_emb__auto-wiki-qa-pair", "hpprc_emb__baobab-wiki-retrieval", # "hpprc_emb__jagovfaqs", JMTEB task のtestに正解が含まれている "hpprc_emb__janli-triplet", "hpprc_emb__jaquad", "hpprc_emb__jqara", # JMTEB task のドメイン "hpprc_emb__jsnli-triplet", "hpprc_emb__jsquad", "hpprc_emb__miracl", # JMTEB task のドメイン "hpprc_emb__mkqa", "hpprc_emb__mkqa-triplet", # "hpprc_emb__mmarco", 文字化け等が含みノイジー "hpprc_emb__mr-tydi", # JMTEB task のドメイン "hpprc_emb__nu-mnli-triplet", "hpprc_emb__nu-snli-triplet", # "hpprc_emb__paws-x-triplet", JMTEB task のtestに含まれている？ "hpprc_emb__quiz-no-mori", "hpprc_emb__quiz-works", "hpprc_emb__snow-triplet", "hpprc_llmjp-kaken", "hpprc_llmjp_warp_html", "hpprc_mqa_ja", "hpprc_msmarco_ja", ] AUG_FACTOR_DATASETS = { "hpprc_emb__miracl": 20, "hpprc_emb__mr-tydi": 20, "hpprc_emb__jqara": 10, "hpprc_emb__baobab-wiki-retrieval": 5, "hpprc_emb__mkqa": 5, "hpprc_emb__auto-wiki-qa-nemotron": 2, "hpprc_msmarco_ja": 2, } os.environ["TOKENIZERS_PARALLELISM"] = "false" logging.basicConfig( format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO ) random.seed(12) def _load_train_eval_datasets_en(): """ Either load the train and eval datasets from disk or load them from the datasets library & save them to disk. Upon saving to disk, we quit() to ensure that the datasets are not loaded into memory before training. """ en_train_dataset_dir = PROJECT_ROOT / "datasets" / "en_train_dataset" en_eval_dataset_dir = PROJECT_ROOT / "datasets" / "en_eval_dataset" try: train_dataset = DatasetDict.load_from_disk(en_train_dataset_dir) eval_dataset = DatasetDict.load_from_disk(en_eval_dataset_dir) return train_dataset, eval_dataset except FileNotFoundError: print("Loading gooaq dataset...") gooaq_dataset = load_dataset("sentence-transformers/gooaq", split="train") gooaq_dataset_dict = gooaq_dataset.train_test_split(test_size=10_000, seed=12) gooaq_train_dataset: Dataset = gooaq_dataset_dict["train"] gooaq_eval_dataset: Dataset = gooaq_dataset_dict["test"] print("Loaded gooaq dataset.") print("Loading msmarco dataset...") msmarco_dataset = load_dataset( "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1", "triplet", split="train", ) msmarco_dataset_dict = msmarco_dataset.train_test_split( test_size=10_000, seed=12 ) msmarco_train_dataset: Dataset = msmarco_dataset_dict["train"] msmarco_eval_dataset: Dataset = msmarco_dataset_dict["test"] print("Loaded msmarco dataset.") print("Loading squad dataset...") squad_dataset = load_dataset("sentence-transformers/squad", split="train") squad_dataset_dict = squad_dataset.train_test_split(test_size=10_000, seed=12) squad_train_dataset: Dataset = squad_dataset_dict["train"] squad_eval_dataset: Dataset = squad_dataset_dict["test"] print("Loaded squad dataset.") print("Loading s2orc dataset...") s2orc_dataset = load_dataset( "sentence-transformers/s2orc", "title-abstract-pair", split="train[:100000]" ) s2orc_dataset_dict = s2orc_dataset.train_test_split(test_size=10_000, seed=12) s2orc_train_dataset: Dataset = s2orc_dataset_dict["train"] s2orc_eval_dataset: Dataset = s2orc_dataset_dict["test"] print("Loaded s2orc dataset.") print("Loading allnli dataset...") allnli_train_dataset = load_dataset( "sentence-transformers/all-nli", "triplet", split="train" ) allnli_eval_dataset = load_dataset( "sentence-transformers/all-nli", "triplet", split="dev" ) print("Loaded allnli dataset.") print("Loading paq dataset...") paq_dataset = load_dataset("sentence-transformers/paq", split="train") paq_dataset_dict = paq_dataset.train_test_split(test_size=10_000, seed=12) paq_train_dataset: Dataset = paq_dataset_dict["train"] paq_eval_dataset: Dataset = paq_dataset_dict["test"] print("Loaded paq dataset.") print("Loading trivia_qa dataset...") trivia_qa = load_dataset("sentence-transformers/trivia-qa", split="train") trivia_qa_dataset_dict = trivia_qa.train_test_split(test_size=5_000, seed=12) trivia_qa_train_dataset: Dataset = trivia_qa_dataset_dict["train"] trivia_qa_eval_dataset: Dataset = trivia_qa_dataset_dict["test"] print("Loaded trivia_qa dataset.") print("Loading msmarco_10m dataset...") msmarco_10m_dataset = load_dataset( "bclavie/msmarco-10m-triplets", split="train" ) msmarco_10m_dataset_dict = msmarco_10m_dataset.train_test_split( test_size=10_000, seed=12 ) msmarco_10m_train_dataset: Dataset = msmarco_10m_dataset_dict["train"] msmarco_10m_eval_dataset: Dataset = msmarco_10m_dataset_dict["test"] print("Loaded msmarco_10m dataset.") print("Loading swim_ir dataset...") swim_ir_dataset = load_dataset( "nthakur/swim-ir-monolingual", "en", split="train" ).select_columns(["query", "text"]) swim_ir_dataset_dict = swim_ir_dataset.train_test_split( test_size=10_000, seed=12 ) swim_ir_train_dataset: Dataset = swim_ir_dataset_dict["train"] swim_ir_eval_dataset: Dataset = swim_ir_dataset_dict["test"] print("Loaded swim_ir dataset.") # NOTE: 20 negatives print("Loading pubmedqa dataset...") pubmedqa_dataset = load_dataset( "sentence-transformers/pubmedqa", "triplet-20", split="train" ) pubmedqa_dataset_dict = pubmedqa_dataset.train_test_split( test_size=100, seed=12 ) pubmedqa_train_dataset: Dataset = pubmedqa_dataset_dict["train"] pubmedqa_eval_dataset: Dataset = pubmedqa_dataset_dict["test"] print("Loaded pubmedqa dataset.") # NOTE: A lot of overlap with anchor/positives print("Loading miracl dataset...") miracl_dataset = load_dataset( "sentence-transformers/miracl", "en-triplet-all", split="train" ) miracl_dataset_dict = miracl_dataset.train_test_split(test_size=10_000, seed=12) miracl_train_dataset: Dataset = miracl_dataset_dict["train"] miracl_eval_dataset: Dataset = miracl_dataset_dict["test"] print("Loaded miracl dataset.") # NOTE: A lot of overlap with anchor/positives print("Loading mldr dataset...") mldr_dataset = load_dataset( "sentence-transformers/mldr", "en-triplet-all", split="train" ) mldr_dataset_dict = mldr_dataset.train_test_split(test_size=10_000, seed=12) mldr_train_dataset: Dataset = mldr_dataset_dict["train"] mldr_eval_dataset: Dataset = mldr_dataset_dict["test"] print("Loaded mldr dataset.") # NOTE: A lot of overlap with anchor/positives print("Loading mr_tydi dataset...") mr_tydi_dataset = load_dataset( "sentence-transformers/mr-tydi", "en-triplet-all", split="train" ) mr_tydi_dataset_dict = mr_tydi_dataset.train_test_split( test_size=10_000, seed=12 ) mr_tydi_train_dataset: Dataset = mr_tydi_dataset_dict["train"] mr_tydi_eval_dataset: Dataset = mr_tydi_dataset_dict["test"] print("Loaded mr_tydi dataset.") train_dataset = DatasetDict( { "gooaq": gooaq_train_dataset, "msmarco": msmarco_train_dataset, "squad": squad_train_dataset, "s2orc": s2orc_train_dataset, "allnli": allnli_train_dataset, "paq": paq_train_dataset, "trivia_qa": trivia_qa_train_dataset, "msmarco_10m": msmarco_10m_train_dataset, "swim_ir": swim_ir_train_dataset, "pubmedqa": pubmedqa_train_dataset, "miracl": miracl_train_dataset, "mldr": mldr_train_dataset, "mr_tydi": mr_tydi_train_dataset, } ) eval_dataset = DatasetDict( { "gooaq": gooaq_eval_dataset, "msmarco": msmarco_eval_dataset, "squad": squad_eval_dataset, "s2orc": s2orc_eval_dataset, "allnli": allnli_eval_dataset, "paq": paq_eval_dataset, "trivia_qa": trivia_qa_eval_dataset, "msmarco_10m": msmarco_10m_eval_dataset, "swim_ir": swim_ir_eval_dataset, "pubmedqa": pubmedqa_eval_dataset, "miracl": miracl_eval_dataset, "mldr": mldr_eval_dataset, "mr_tydi": mr_tydi_eval_dataset, } ) train_dataset.save_to_disk(str(en_train_dataset_dir)) eval_dataset.save_to_disk(str(en_eval_dataset_dir)) return train_dataset, eval_dataset def load_train_eval_datasets_en(target_dataset_names: list[str] = []): print("Loading train and eval datasets...") if len(target_dataset_names) == 0: return DatasetDict(), DatasetDict() train_dataset, eval_dataset = _load_train_eval_datasets_en() ds_names = list(train_dataset.keys()) for ds_name in ds_names: if ds_name not in target_dataset_names: del train_dataset[ds_name] del eval_dataset[ds_name] else: print( "target en ds", ds_name, len(train_dataset[ds_name]), len(eval_dataset[ds_name]), ) return train_dataset, eval_dataset def load_train_eval_datasets_jp(target_dataset_names: list[str] = []): print("Loading train and eval datasets...") jp_train_dataset_dir = PROJECT_ROOT / "datasets" / "jp_train_dataset" jp_eval_dataset_dir = PROJECT_ROOT / "datasets" / "jp_eval_dataset" train_dataset_dict = {} eval_dataset_dict = {} for ds_name in target_dataset_names: print("loading jp ds", ds_name) try: train_ds = Dataset.load_from_disk(f"{jp_train_dataset_dir}/{ds_name}") eval_ds = Dataset.load_from_disk(f"{jp_eval_dataset_dir}/{ds_name}") except FileNotFoundError: print(f"{ds_name} not found, loading from datasets library...") ds = load_dataset( "hotchpotch/sentence_transformer_japanese", ds_name, split="train" ) ds_size = len(ds) test_size = min(3000, ds_size // 100) splitted = ds.train_test_split(test_size=test_size, seed=12) train_ds = splitted["train"] eval_ds = splitted["test"] # save train_ds.save_to_disk(f"{jp_train_dataset_dir}/{ds_name}") eval_ds.save_to_disk(f"{jp_eval_dataset_dir}/{ds_name}") train_dataset_dict[ds_name] = train_ds eval_dataset_dict[ds_name] = eval_ds return DatasetDict(train_dataset_dict), DatasetDict(eval_dataset_dict) def main(): # 1. Load a model to finetune with 2. (Optional) model card data print("Loading model...") static_embedding = StaticEmbedding( AutoTokenizer.from_pretrained("hotchpotch/xlm-roberta-japanese-tokenizer"), embedding_dim=1024, ) model = SentenceTransformer( modules=[static_embedding], model_card_data=SentenceTransformerModelCardData( language="ja", license="mit", model_name="Static Embeddings with japanese tokenizer finetuned on various datasets", ), ) # 3. Set up training & evaluation datasets - each dataset is trained with MNRL (with MRL) print("Loading datasets...") train_dataset_en, eval_dataset_en = load_train_eval_datasets_en(EN_TARGET_DATASETS) train_dataset_jp, eval_dataset_jp = load_train_eval_datasets_jp(JA_TARGET_DATASETS) # merge print("Merging datasets...") train_dataset = DatasetDict({**train_dataset_en, **train_dataset_jp}) eval_dataset = DatasetDict({**eval_dataset_en, **eval_dataset_jp}) for ds_name, aug_factor in AUG_FACTOR_DATASETS.items(): columns = train_dataset[ds_name].column_names def data_aug(example): result = {} for col in columns: result[col] = example[col] * aug_factor return result before_len = len(train_dataset[ds_name]) train_dataset[ds_name] = train_dataset[ds_name].map( data_aug, batched=True, num_proc=11 ) print("data augmented", ds_name, before_len, len(train_dataset[ds_name])) for train_ds_name in train_dataset.keys(): print( "train ds", train_ds_name, len(train_dataset[train_ds_name]), len(eval_dataset[train_ds_name]), ) # 4. Define a loss function loss = MultipleNegativesRankingLoss(model) loss = MatryoshkaLoss(model, loss, matryoshka_dims=[32, 64, 128, 256, 512, 1024]) # 5. (Optional) Specify training arguments run_name = f"static-retrieval-mrl-jp-v1_{EXP}" args = SentenceTransformerTrainingArguments( # Required parameter: output_dir=f"models/{run_name}", # Optional training parameters: num_train_epochs=2, per_device_train_batch_size=2048 * 3, # gradient_accumulation_steps=4, per_device_eval_batch_size=2048, learning_rate=2e-1, lr_scheduler_type="cosine", # optim="adafactor", warmup_ratio=0.1, fp16=False, # Set to False if you get an error that your GPU can't run on FP16 bf16=True, # Set to True if you have a GPU that supports BF16 batch_sampler=BatchSamplers.NO_DUPLICATES, # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL, # Optional tracking/debugging parameters: eval_strategy="steps", eval_steps=200, save_strategy="steps", save_steps=200, save_total_limit=20, logging_steps=20, logging_first_step=True, dataloader_prefetch_factor=4, dataloader_num_workers=15, run_name=run_name, # Will be used in W&B if `wandb` is installed ) # 6. (Optional) Create an evaluator & evaluate the base model evaluator = NanoBEIREvaluator() evaluator(model) # 7. Create a trainer & train trainer = SentenceTransformerTrainer( model=model, args=args, train_dataset=train_dataset, eval_dataset=eval_dataset, loss=loss, evaluator=evaluator, ) trainer.train() # (Optional) Evaluate the trained model on the evaluator after training evaluator(model) # 8. Save the trained model model.save_pretrained(f"{PROJECT_ROOT}/models/{run_name}/final") # 9. (Optional) Push it to the Hugging Face Hub # model.push_to_hub(run_name, private=True) if __name__ == "__main__": main()