Skip to content

Karen/contrastive train #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kar-m wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Karen/contrastive train #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4f30cec
added proper shuffling
kar-m Mar 10, 2026
70f1886
added contrastive training stuf
kar-m Mar 10, 2026
007178a
clean
kar-m Mar 10, 2026
0334bf2
added changes
kar-m Mar 10, 2026
788f9ba
aaaaaaa
kar-m Mar 11, 2026
f35f1ea
cleaned up diff
kar-m Mar 11, 2026
c42aec3
got alright results
kar-m Mar 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions setup
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,11 @@ fi

source "$VENV_DIR/bin/activate"

export NEEDLE_CACHE_DIR="${NEEDLE_CACHE_DIR:-/dev/shm/needle_cache}"
export NEEDLE_TOKENIZER_DIR="${NEEDLE_TOKENIZER_DIR:-/dev/shm/needle_tokenizer}"
export NEEDLE_LOCAL_UNIFIED_DIR="${NEEDLE_LOCAL_UNIFIED_DIR:-/dev/shm/needle_tool_calls_unified}"
mkdir -p "$NEEDLE_CACHE_DIR" "$NEEDLE_TOKENIZER_DIR" "$NEEDLE_LOCAL_UNIFIED_DIR"

echo "Installing dependencies..."
pip install --upgrade pip -q
pip install -e . -q
Expand Down
73 changes: 73 additions & 0 deletions src/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,55 @@ def main():
help="Number of mel frequency bins (default: 80)")
p.add_argument("--max-speech-samples", type=int, default=None,
help="Max voice-tool-call training samples (default: all)")
p.add_argument("--speech-replay-every", type=int, default=0,
help="Do 1 speech transcription replay step every N text steps (0=disabled)")
p.add_argument("--speech-gcs-prefix", type=str,
default="gs://cactus-dataset/speech-datav1/emilia-large")
p.add_argument("--reinit-decoder", action="store_true",
help="Reinitialise decoder weights from scratch when loading a checkpoint (keeps encoder)")

p = sub.add_parser("pretrain", add_help=False)
p.add_argument("--full", action="store_true")
p.add_argument("--checkpoint", type=str, default=None)
p.add_argument("--epochs", type=int, default=3)
p.add_argument("--batch-size", type=int, default=32)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--muon-lr", type=float, default=0.02)
p.add_argument("--d-model", type=int, default=512)
p.add_argument("--num-heads", type=int, default=16)
p.add_argument("--num-kv-heads", type=int, default=8)
p.add_argument("--num-layers", type=int, default=4)
p.add_argument("--num-dec-layers", type=int, default=4)
p.add_argument("--max-enc-len", type=int, default=256)
p.add_argument("--max-dec-len", type=int, default=1024)
p.add_argument("--max-samples", type=int, default=None)
p.add_argument("--max-speech-samples", type=int, default=None)
p.add_argument("--speech-replay-every", type=int, default=0,
help="Do 1 speech transcription replay step every N text steps (0=disabled)")
p.add_argument("--warmup-ratio", type=float, default=0.05)
p.add_argument("--wandb", action="store_true")
p.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "bfloat16"])
p.add_argument("--checkpoint-dir", type=str, default="checkpoints")
p.add_argument("--seed", type=int, default=42)
p.add_argument("--eval-every", type=int, default=1000)
p.add_argument("--max-eval-samples", type=int, default=None)
p.add_argument("--group-size", type=int, default=32)
p.add_argument("--activation", type=str, default="drelu", choices=["drelu", "swiglu", "geglu"])
p.add_argument("--num-memory-slots", type=int, default=64)
p.add_argument("--dropout", type=float, default=0.0,
help="Dropout rate for residual connections (default: 0.1)")
p.add_argument("--max-mel-len", type=int, default=1024,
help="Max mel spectrogram frames (default: 1024)")
p.add_argument("--n-mels", type=int, default=80,
help="Number of mel frequency bins (default: 80)")
p.add_argument("--speech-gcs-prefix", type=str, default="gs://cactus-dataset/speech-datav1/emilia-large")
p.add_argument("--speech-val-ratio", type=float, default=0.01)
p.add_argument("--toucan-config", type=str, default="Kimi-K2")
p.add_argument("--toucan-max-samples", type=int, default=None)
p.add_argument("--tool-contrastive-weight", type=float, default=1.0,
help="Toucan query-description contrastive loss weight")
p.add_argument("--audio-text-contrastive-weight", type=float, default=1.0,
help="Paired audio-text SigLIP loss weight")

p = sub.add_parser("tokenize", add_help=False)
p.add_argument("--max-samples", type=int, default=None,
Expand All @@ -72,6 +121,20 @@ def main():
help="Max decoder sequence length (default: 1024)")
p.add_argument("--batch-size", type=int, default=None,
help="Process in batches of this size, uploading shards to GCS incrementally")
p.add_argument("--shuffle-before-split", action="store_true",
help="Shuffle the unified dataset before splitting into train/val")
p.add_argument("--split-seed", type=int, default=42,
help="Seed for --shuffle-before-split (default: 42)")
p.add_argument("--speech-gcs-prefix", type=str, default="gs://cactus-dataset/speech-datav1/emilia-large")
p.add_argument("--speech-val-ratio", type=float, default=0.01)
p.add_argument("--max-speech-samples", type=int, default=None,
help="Optional cap for Emilia Stage 1 preprocessing")
p.add_argument("--toucan-config", type=str, default="Kimi-K2",
help="Toucan subset to parse and cache during tokenization")
p.add_argument("--toucan-max-samples", type=int, default=None,
help="Optional max samples for Toucan Stage 1 preprocessing")
p.add_argument("--overwrite-gcs-tokenizer", action="store_true",
help="Overwrite the shared GCS tokenizer after retraining locally")

p = sub.add_parser("run", add_help=False)
p.add_argument("--checkpoint", type=str, required=True)
Expand Down Expand Up @@ -140,6 +203,16 @@ def main():
if args.command == "tokenize":
from .tokenize_data import tokenize
tokenize(args)
elif args.command == "pretrain":
if getattr(args, "full", False):
args.d_model = 1536
args.num_heads = 24
args.num_kv_heads = 8
args.num_layers = 12
args.num_dec_layers = 4
args.num_memory_slots = 128
from .pretrain import pretrain
pretrain(args)
elif args.command == "train":
if getattr(args, "full", False):
args.d_model = 1536
Expand Down
Loading