train script

train.py

@@ -0,0 +1,105 @@
+from datasets import load_dataset, Dataset
+import pandas as pd
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
+
+
+def clean_text(x):
+    if x is None:
+        return ""
+    x = str(x).strip()
+    x = " ".join(x.split())
+    return x
+
+
+def build_doc_fast(context):
+    return f"text: {context}"
+
+
+def main():
+    dataset_name = "phamson02/large-vi-legal-queries"
+
+    # Load + clean
+    ds = load_dataset(dataset_name, split="train")
+    df = ds.to_pandas()
+
+    print("Raw shape:", df.shape)
+
+    for col in ["domain", "title", "header", "aspect", "context", "query"]:
+        if col not in df.columns:
+            df[col] = ""
+        df[col] = df[col].apply(clean_text)
+
+    df = df[(df["query"] != "") & (df["context"] != "")]
+    df = df.drop_duplicates(subset=["query", "context"]).reset_index(drop=True)
+
+    print("Cleaned rows:", len(df))
+
+    train_df = pd.DataFrame(
+        {
+            "anchor": df["query"].tolist(),
+            "positive": [build_doc_fast(context) for context in df["context"].tolist()],
+        }
+    )
+
+    print(train_df.head())
+
+    train_dataset = Dataset.from_pandas(train_df, preserve_index=False)
+    print(train_dataset[0])
+
+    # IMPORTANT: no .to("cuda") here under torchrun / DDP
+    model = SentenceTransformer(
+        "google/embeddinggemma-300m",
+        model_kwargs={
+            # "torch_dtype": "auto",
+            # "attn_implementation": "flash_attention_2",
+        },
+    )
+    model.max_seq_length = 512
+
+    loss = CachedMultipleNegativesRankingLoss(
+        model,
+        mini_batch_size=32,
+        gather_across_devices=False,
+    )
+
+    task_name = "Retrieval"
+    training_args = SentenceTransformerTrainingArguments(
+        prompts=model.prompts[task_name],
+        torch_compile=False,
+        output_dir="./embeddinggemma-300m-vilegal",
+        num_train_epochs=1,
+        per_device_train_batch_size=1024,
+        gradient_accumulation_steps=1,
+        learning_rate=2e-5,
+        warmup_ratio=0.1,
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        report_to="none",
+        remove_unused_columns=False,
+        dataloader_num_workers=8,
+        dataloader_persistent_workers=True,
+        # Often helpful for DDP stability/perf with Transformer training:
+        ddp_find_unused_parameters=False,
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        loss=loss,
+    )
+
+    trainer.train()
+
+    # Save only once from the main process
+    trainer.save_model("./embeddinggemma-300m-vilegal")
+
+
+if __name__ == "__main__":
+    main()