Upload trainer.py

add trainer script

trainer.py

@@ -0,0 +1,432 @@
+# 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()
+