This repository contains the training code for SEA-Embedding, an open and reproducible text embedding model for Southeast Asian languages. The training pipeline is built on top of Sentence Transformers and consists of two stages: pre-training (Step 1) and fine-tuning (Step 2).

Requires Python 3.9+ and PyTorch 1.11.0+.

conda create -n sea-embedding python=3.10
conda activate sea-embedding

# Install the package in editable mode
pip install -e .

# Install additional training dependencies
pip install -r requirements.txt

# Install uv if not already installed
pip install uv

# Sync all dependencies from the lockfile
uv sync

# Or install directly from requirements
uv pip install -r requirements.txt
uv pip install -e .

The training pipeline is split into two steps, each with a corresponding SLURM script for multi-node, multi-GPU execution.

Edit slurm_run_multinode_step1.sh to set your cluster parameters:

#SBATCH --nodes=NUMBER_OF_NODES      # number of nodes to use
#SBATCH --nodelist=YOUR_NODE         # specific node names
#SBATCH --gres=gpu:GPU               # number of GPUs per node

Then set the model and output name inside the script:

model="MODEL FROM HF"    # e.g. "aisingapore/sea-lion-7b"
model_name="model_name"  # used for output directory naming

Submit the job:

sbatch slurm_run_multinode_step1.sh

Logs are saved to log/step1/step1-<job_id>.out.

Edit slurm_run_multinode_step2.sh to set the model from Step 1:

model="MODEL FROM previous step"  # path or HF repo of Step 1 output
model_name="model_name"

Update the SLURM header as in Step 1, then submit:

sbatch slurm_run_multinode_step2.sh

Logs are saved to log/step2/step2-<job_id>.out.

Both scripts sweep over the following hyperparameters by default:

If you use SEA-Embedding in your work, please cite:

@misc{limkonchotiwat2026seaembeddingopenreproducibletext,
      title={SEA-Embedding: Open and Reproducible Text Embeddings for Southeast Asia}, 
      author={Peerat Limkonchotiwat and Raymond Ng and Sarana Nutanong and Jian Gang Ngui},
      year={2026},
      eprint={2606.03027},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2606.03027}, 
}