Kame Finetuning Workflow

KAME: TANDEM ARCHITECTURE FOR ENHANCING KNOWLEDGE IN REAL-TIME SPEECH-TO-SPEECH CONVERSATIONAL AI

KAME Inference code · Paper · Blog post · Hugging Face

Kame is an oracle-enabled extension of Moshi for full-duplex spoken dialogue. This repository provides the preprocessing, finetuning, checkpoint conversion, and inference workflow for Kame.

The public preprocessing entry point is a canonical dataset layout built from stereo audio and word-level transcripts:

• audio/<dialogue_id>.wav
• text/<dialogue_id>.json
• optional: oracle_raw/<dialogue_id>.json

This repository also includes a small sample under data/spokenwoz_sample/{audio,text,oracle_raw} so you can run the preprocessing steps on a bundled example before using your own data.

Installation

Python 3.12+ is required.

Install the project dependencies:

uv sync --python 3.12

This repository is intended to be used as a uv-managed finetuning workflow repo rather than as a standalone installable Python package.

If you use Weights & Biases for training logs:

wandb login

Canonical Dataset Layout

Each stereo wav must contain speaker A in the left channel and speaker B in the right channel. Each transcript JSON file must contain a word-level transcript for both speakers with timestamps aligned to the corresponding wav file.

The canonical text/*.json format is:

[
  {"speaker": "A", "word": "hello", "start": 0.46, "end": 1.52},
  {"speaker": "B", "word": "hi", "start": 1.82, "end": 2.04},
  {"speaker": "B", "word": "customer", "start": 2.04, "end": 2.703},
  {"speaker": "B", "word": "service", "start": 2.703, "end": 3.145},
  {"speaker": "B", "word": "how", "start": 3.145, "end": 3.366}
]

Each entry contains:

• speaker: "A" or "B"
• word: a word or tokenizer unit
• start: start time in seconds
• end: end time in seconds

If your dataset does not include word-level timestamps, create them first with your preferred alignment pipeline before continuing.

If oracle predictions are already present, oracle_raw/*.json should follow this format:

[
  {
    "timestamp_ms": 1500,
    "conversation_context": "A: hello there",
    "prediction": "hi how can I help you today",
    "total_word_count": 2,
    "trigger_word": "there",
    "recent_words": "hello there",
    "current_spoken_ratio": 0.75,
    "channel": 1,
    "hint": "hi how can I help you today"
  }
]

The hint field is intentionally empty when current_spoken_ratio <= 0.5.

Standard Preprocessing

1. Optional: Generate oracle_raw From Canonical Text Transcripts

If your dataset does not already include oracle predictions, generate them directly from text/*.json:

export OPENAI_API_KEY=...

uv run --extra oracle -m tools.generate_oracle_from_text \
  --text_dir data/my_dataset/text \
  --output_dir data/my_dataset/oracle_raw

By default this command assumes the canonical mapping A_channel=0 and B_channel=1.

2. Audio Tokenization

Convert all wav files into discrete Mimi tokens:

uv run -m tools.tokenize_audio \
  --audio_dir data/my_dataset/audio \
  --output_dir data/my_dataset/tokenized_audio

This creates data/my_dataset/tokenized_audio/*.npz. Each npz file contains audio tokens for A and B.

3. Text Tokenization

Convert the canonical transcripts into frame-level text streams:

uv run -m tools.tokenize_text \
  --word_transcript_dir data/my_dataset/text \
  --output_dir data/my_dataset/tokenized_text

This creates data/my_dataset/tokenized_text/*.npz. Each npz file contains text tokens for A and B.

If you use a different tokenizer, also pass --text_tokenizer_repo, --text_tokenizer_name, and the matching padding IDs.

4. Optional: Oracle Tokenization

If oracle_raw/*.json is present, convert it into event-level oracle records:

uv run -m tools.tokenize_oracle \
  --oracle_dir data/my_dataset/oracle_raw \
  --oracle_suffix ".json" \
  --tokenized_audio_dir data/my_dataset/tokenized_audio \
  --output_dir data/my_dataset/tokenized_oracle_a0b1_events \
  --A_channel 0 \
  --B_channel 1

This creates data/my_dataset/tokenized_oracle_a0b1_events/*.npz.

5. Build Parquet Files

Concatenate audio, text, and optional oracle streams into a parquet dataset ready for finetuning:

uv run -m tools.prepare_dataset \
  --tokenized_text_dir data/my_dataset/tokenized_text \
  --tokenized_audio_dir data/my_dataset/tokenized_audio \
  --tokenized_oracle_dir data/my_dataset/tokenized_oracle_a0b1_events \
  --output_prefix processed_data/my_dataset/train_text_oracle_a0b1_events

If you do not use oracle predictions, omit --tokenized_oracle_dir.

Model Initialization

The example scripts in this README assume a KAME finetuning model initialized from the original Kyutai weights. The current English example scripts use:

• init_models/moshiko-one_streams-bfloat16

To initialize from the original Kyutai weights:

uv run -m tools.init_moshi_for_ft \
  --moshi_lm_repo kyutai/moshiko-pytorch-bf16 \
  --save_dir init_models/moshiko-one_streams-bfloat16 \
  --model_dtype bfloat16

If you change the text tokenizer, also use --init_text_embeddings and keep the vocabulary size compatible.

Training

For a low-memory smoke test, run:

MAX_TRAIN_STEPS=3 bash examples/finetune_accelerate_cpu_offload.sh

This smoke example keeps the default finetuning target but uses a more conservative DeepSpeed configuration with CPU offload. It is intentionally slower, but is a better fit for validating that the public workflow runs end to end on a single GPU.

For a fuller training run, use:

bash examples/finetune_accelerate.sh

This reference script uses the default KAME finetuning target and a faster DeepSpeed configuration, but it may require substantial GPU memory. In practice, full finetuning may need multi-GPU execution depending on your hardware.

The current training implementation requires DeepSpeed, so both examples use Accelerate with a DeepSpeed config. On managed clusters you may wrap these commands in your own scheduler submission flow such as sbatch, but scheduler-specific scripts are intentionally omitted from this public repository.

Convert and Clean Checkpoints for Inference

After training, convert checkpoints in two stages:

1. Convert DeepSpeed checkpoints to fp32 safetensors

uv run -m tools.zero_to_fp32 \
  output/moshiko-finetuned/step_10000 \
  output/moshiko-finetuned/step_10000_fp32 \
  --moshi_lm_kwargs_path init_models/moshiko-one_streams-bfloat16/moshi_lm_kwargs.json

2. Clean the finetuning model for inference

uv run -m tools.clean_moshi \
  --moshi_ft_dir output/moshiko-finetuned/step_10000_fp32 \
  --save_dir output/moshiko-finetuned/step_10000_fp32_cleaned \
  --model_dtype float32

By default, clean_moshi reads oracle_embedding_mode from the training run's config.json, located in the parent directory of --moshi_ft_dir. This matches the standard layout above, where the config and checkpoint directories share a training root. The recorded metadata is the source of truth.

For a moved or nonstandard checkpoint layout, pass the config explicitly with --training_config_path /path/to/config.json. You may also pass --oracle_embedding_mode separate or tie as an explicit cross-check; export fails if it disagrees with the training metadata. For a legacy checkpoint whose config is unavailable or does not record this setting, the explicit mode is required as a fallback.

For tie, the post-ZeRO export copies the learned text_emb state into oracle_emb and verifies exact equality. For separate, it preserves the independently trained oracle_emb.

Oracle-Enabled Inference

Start the oracle-enabled server:

CLEAN_DIR=output/moshiko-finetuned/step_10000_fp32_cleaned
TOKENIZER_PATH=/path/to/tokenizer_spm_32k_3.model

uv run -m kame.server_oracle \
  --moshi-weight "$CLEAN_DIR/model.safetensors" \
  --config-path "$CLEAN_DIR/moshi_lm_kwargs.json" \
  --tokenizer "$TOKENIZER_PATH" \
  --host 0.0.0.0 \
  --port 8998

If the server runs on a remote node, use SSH port forwarding from your local machine:

ssh -L 8998:localhost:8998 user@remote-host

License

This repository is provided under the Apache 2.0 License, following the license of the upstream moshi-finetune repository. The SpokenWOZ sample data included in data/spokenwoz_sample is provided under CC BY-NC 4.0.

Attribution

kame_finetune is derived from the moshi-finetune codebase and adapted for the KAME training workflow.

Citation

If you use KAME or this finetuning workflow in your research, please cite:

@article{kuroki2025kame,
  title={KAME: Tandem Architecture for Enhancing Knowledge in Real-Time Speech-to-Speech Conversational AI},
  author={Kuroki, So and Kubo, Yotaro and Akiba, Takuya and Tang, Yujin},
  journal={arXiv preprint arXiv:2510.02327},
  year={2025}
}