services:
  nomen-ai-train:
    build: .
    image: nomen-ai:latest
    environment:
      HF_TOKEN: ${HF_TOKEN}
      TRACKIO_PROJECT: nomen-ai
      TRACKIO_SPACE_ID: krystv/ml-intern-nomen001
      BASE_MODEL: Qwen/Qwen2.5-1.5B-Instruct
      SFT_DATASET: krystv/nomen-ai-sft
      DPO_DATASET: krystv/nomen-ai-dpo
      SFT_HUB_MODEL_ID: krystv/nomen-ai-sft-lora
      DPO_HUB_MODEL_ID: krystv/nomen-ai-dpo-lora
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    shm_size: 8gb
    command: bash scripts/train_all_colab.sh