Csm 1b Hf

🚀 CSM-1B-HF

このモデルは、Hugging Face実装用のSesame CSM 1Bモデルの重みです。これは、音声合成や会話音声生成などのタスクに利用でき、Hugging Faceのtransformersライブラリと完全に互換性があります。

🚀 クイックスタート

CSM-HFは、Sesameの会話音声モデル（CSM）のHugging Face実装です。このコードベースは、Sesameが提供するpytorchコードを完全に書き換えたもので、推論からトレーニングまでHugging Faceのtransformersと完全に互換性があるように設計されています。

✨ 主な機能

• CSMModelクラスを作成しました。
• バックボーンとデコーダのtorchtuneモデルをHF transformersのLllamaModelに置き換えました。
• モデルの入力を準備するためのプロセッサクラスを追加しました。
• ラベルサポートとデコーダトレーニングのアモーティゼーションを追加しました。
• 音声生成のために、モデルクラスにgenerate_frameとgenerateメソッドを追加しました。
• HuggingFaceのTrainerを完全にサポートしています。

📦 インストール

このモデルはHugging Face Hubからダウンロードして使用できます。依存関係のインストールやモデルのロードに関する詳細は、GitHubリポジトリを参照してください。

💻 使用例

基本的な使用法

import torch
from modeling_csm import CSMModel
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer
from tokenizers.processors import TemplateProcessing
from moshi.models import loaders
from processor import CSMProcessor
import torchaudio

device = 'cuda'

def load_llama3_tokenizer():
    """
    https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
    """
    tokenizer_name = "meta-llama/Llama-3.2-1B"
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    bos = tokenizer.bos_token
    eos = tokenizer.eos_token
    tokenizer._tokenizer.post_processor = TemplateProcessing(
        single=f"{bos}:0 $A:0 {eos}:0",
        pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1",
        special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)],
    )

    return tokenizer

text_tokenizer = load_llama3_tokenizer()

mimi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MIMI_NAME)
audio_tokenizer = loaders.get_mimi(mimi_weight, device=device)
audio_tokenizer.set_num_codebooks(32)

processor = CSMProcessor(text_tokenizer, audio_tokenizer)


def load_audio(path, target_sr):
    audio, sr = torchaudio.load(path)
    audio = audio.squeeze(0)
    if sr != target_sr:
        audio = torchaudio.functional.resample(audio, orig_freq=sr, new_freq=target_sr)
    return audio


model = CSMModel.from_pretrained("thomasgauthier/csm-1b-hf", torch_dtype=torch.bfloat16)
model.to('cuda')


inputs = processor(
    messages=[
        {
        "role": "speaker_0",
        "content": [
            {"type": "text", "text": "<AUDIO_CLIP_TRANSCRIPT>"},
            {"type": "audio"} # This placeholder is required for audio tokenization (it maps to the first element in the `audios` list passed to the processor)
        ]
    },
            {
        "role": "speaker_0",
        "content": [
            {"type": "text", "text": "Hello, this is voice cloning speaking"},
            # does not include audio as the model will generate it
        ]
    }
        ], 
    audios=[load_audio('AUDIO_CLIP_FOR_VOICE_CLONING.wav', audio_tokenizer.sample_rate)],
    return_tensors="pt"
)

import torch

with torch.inference_mode():
    # Generate up to 50 new frames
    gen_frames = model.generate(
        input_ids=inputs['input_ids'].cuda(),
        attention_mask=inputs['attention_mask'].cuda(),
        max_new_frames=50,
        topk=50,
        temperature=1.0,
        use_cache=True,
        stop_on_all_zeros=True,

    )

decoded_audio = audio_tokenizer.decode(gen_frames.permute(0, 2, 1)).squeeze(0).squeeze(0)

audio_array = (decoded_audio * 32768).to(torch.int16).cpu().numpy()

# Audio can be played with the following code:
# from IPython.display import Audio
# Audio(audio_array, rate=audio_tokenizer.sample_rate)

📚 ドキュメント

アーキテクチャ

モデルのアーキテクチャについては、ARCHITECTURE.md（O1によって書かれたもの）を参照してください。

トレーニング

データ形式

CSM-HFは、トレーニングデータをJSONL形式で期待しています。各行は会話を含むJSONオブジェクトです。各会話は以下の要素で構成されています。

• messages：メッセージオブジェクトの配列で、各メッセージには以下の要素があります。
  • role：話者識別子（例："speaker_0", "speaker_1"）
  • content：コンテンツオブジェクトの配列で、以下のタイプがあります。
    • テキスト：{"type": "text", "text": "The message text"}
    • 音声：{"type": "audio", "url": "path/to/audio/file.wav"}
• training_mask：トレーニングに使用するメッセージ（true）またはコンテキスト（false）を示すブール型の配列

データ形式の例：

{
  "messages": [
    {
      "role": "speaker_0",
      "content": [
        {"type": "text", "text": "We have a chance for a new life here."},
        {"type": "audio", "url": "clips/example_audio.wav"}
      ]
    },
    {
      "role": "speaker_1",
      "content": [
        {"type": "text", "text": "Uncle?"},
        {"type": "audio", "url": "clips/response_audio.wav"}
      ]
    }
  ],
  "training_mask": [false, true]
}

トレーニングプロセス

モデルは二段階の自己回帰アーキテクチャを使用しています。

1. バックボーン（フレーム間処理）：
   • フレームの全シーケンスを処理します。
   • 各フレームはすべてのコードブックの結合埋め込みを表します。
   • 発話間の長距離依存関係を処理します。
2. デコーダ（フレーム内処理）：
   • 一度に1つのフレームを処理します。
   • 32のコードブックを順次生成します（1つのセマンティック + 31の音響）。
   • 各コードブックはシーケンス内のトークンとして扱われます。

トレーニングでは、コンピュートアモーティゼーション技術を活用しています。

• ゼロ番目（セマンティック）のコードブックはすべてのフレームでトレーニングされます。
• 残りのコードブック（1 - 31）は、amortization_ratioのフレームでのみトレーニングされます。
• これにより、品質を維持しながらメモリ使用量を大幅に削減できます。

モデルをトレーニングするには、以下のコマンドを実行します。

python train.py \
  --train_file path/to/training_data.jsonl \
  --output_dir ./output \
  --num_train_epochs 3 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --learning_rate 5e-6

今後の予定

• [x] 二段階自己回帰アーキテクチャの実装
• [x] マルチコードブック音声トークン化
• [x] 効率的なトレーニングのためのコンピュートアモーティゼーション
• [x] テキスト/音声が交互に配置されたデータセットの準備
• [x] 個別のバックボーン/デコーダ損失を持つカスタムトレーニングループ
• [x] デコーダアモーティゼーションのためのエポック繰り返しの適切な処理
• [x] メモリ最適化技術（混合精度、勾配累積）
• [ ] 効率的な微調整のためのLoRAサポート
• [ ] torch.compileによる高速推論
• [ ] プロンプトチューニング/プレフィックス最適化による音声クローニング
• [ ] DPOのサポート
• [ ] RL（GRPO、RLOOなど）のサポート

📄 ライセンス

このプロジェクトは、Apache-2.0ライセンスの下でライセンスされています。

謝辞

特別な感謝を以下の組織とツールに送ります。

• Sesame Labs：元のアーキテクチャの設計と実装
• Hugging Face：Transformersライブラリとトレーニングインフラ
• ClaudeとChatGPT：ドキュメントとコード開発の支援

このプロジェクトは、オープンソースコミュニティの研究とツールに基づいています。このようなプロジェクトを可能にする協調精神に感謝しています。