Move MOSS token row builder into codec runtime

Author: 0xShug0

CMakeLists.txt

@@ -801,7 +801,6 @@ audiocpp_add_model(moss
         src/models/moss/moss_tts_nano/prompt_builder.cpp
         src/models/moss/moss_tts_nano/session.cpp
         src/models/moss/moss_tts_nano/tokenizer_text.cpp
-        src/models/moss/shared/token_rows.cpp
         src/models/moss/moss_tts_local/depth_transformer.cpp
         src/models/moss/moss_tts_local/generator.cpp
         src/models/moss/moss_tts_local/loader.cpp

include/engine/community_models/moss_voicegen/session.h

@@ -9,7 +9,6 @@
 #include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 #include "engine/framework/modules/multi_codebook_embedding.h"
 #include "engine/framework/runtime/session_base.h"
-#include "engine/models/moss/shared/token_rows.h"
 
 #include <cstddef>
 #include <memory>

include/engine/community_models/moss_voicegen/tokenizer_text.h

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "engine/community_models/moss_voicegen/assets.h"
-#include "engine/models/moss/shared/token_rows.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 
 #include <memory>
 #include <optional>
@@ -23,7 +23,7 @@ class MossVoiceGenTextProcessor {
 
     // `instruction` describes the speaker to design. `language` must be the full language
     // name the model was trained on ("English", not "en"); an empty value renders "None".
-    moss::TokenRows build_generation_prefix(
+    engine::codecs::MossTokenRows build_generation_prefix(
         const std::string & text,
         const std::optional<std::string> & instruction,
         const std::optional<std::string> & language) const;

include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h

@@ -58,6 +58,27 @@ struct MossAudioTokenizerCodes {
     std::vector<std::vector<int32_t>> codebooks;
 };
 
+struct MossTokenRows {
+    std::vector<int32_t> text_tokens;
+    std::vector<int32_t> audio_codes;
+};
+
+class MossTokenRowBuilder {
+public:
+    MossTokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id);
+
+    void push_text_token(int32_t token_id);
+    void push_text_tokens(const std::vector<int32_t> & token_ids);
+    void push_audio_row(int32_t text_slot_token_id, const int32_t * codes, int64_t num_codebooks);
+    void push_audio_row(int32_t text_slot_token_id, const std::vector<std::vector<int32_t>> & codes, int64_t frame);
+    MossTokenRows finish();
+
+private:
+    int64_t num_codebooks_ = 0;
+    int32_t audio_pad_token_id_ = 0;
+    MossTokenRows rows_;
+};
+
 struct MossAudioTokenizerCodecRuntimeOptions {
     size_t weight_context_bytes = 256ull * 1024ull * 1024ull;
     size_t encoder_graph_arena_bytes = 2048ull * 1024ull * 1024ull;

include/engine/models/moss/moss_tts_local/generator.h

@@ -4,7 +4,6 @@
 #include "engine/models/moss/moss_tts_local/backbone.h"
 #include "engine/models/moss/moss_tts_local/depth_transformer.h"
 #include "engine/framework/modules/multi_codebook_embedding.h"
-#include "engine/models/moss/shared/token_rows.h"
 #include "engine/framework/sampling/torch_random.h"
 
 #include <cstddef>

include/engine/models/moss/moss_tts_local/tokenizer_text.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "engine/models/moss/shared/token_rows.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 #include "engine/models/moss/moss_tts_local/assets.h"
 
 #include <cstdint>
@@ -14,7 +14,7 @@ namespace engine::models::moss_tts_local {
 // Decoder input for a generation request: the text channel (input_ids[..., 0]) plus the
 // n_vq audio channels flattened row-major as [seq, n_vq] (input_ids[..., 1:]). Every audio
 // slot of the prompt carries audio_pad_token_id, matching MossTTSLocalProcessor._build_text_rows.
-using MossGenerationPrefix = moss::TokenRows;
+using MossGenerationPrefix = engine::codecs::MossTokenRows;
 
 // Reproduces the direct-generation branch of MossTTSLocalProcessor: it renders the
 // <user_inst> template, byte-level BPE encodes each piece with the Qwen tokenizer, and

include/engine/models/moss/shared/token_rows.h

@@ -85,12 +85,12 @@ MossVoiceGenTextProcessor::MossVoiceGenTextProcessor(std::shared_ptr<const MossV
 
 MossVoiceGenTextProcessor::~MossVoiceGenTextProcessor() = default;
 
-moss::TokenRows MossVoiceGenTextProcessor::build_generation_prefix(
+engine::codecs::MossTokenRows MossVoiceGenTextProcessor::build_generation_prefix(
     const std::string & text,
     const std::optional<std::string> & instruction,
     const std::optional<std::string> & language) const {
     const auto & config = impl_->assets->config;
-    moss::TokenRowBuilder builder(config.num_codebooks, static_cast<int32_t>(config.audio_pad_code));
+    engine::codecs::MossTokenRowBuilder builder(config.num_codebooks, static_cast<int32_t>(config.audio_pad_code));
 
     // Render the whole turn as one string and encode it in a single pass, the way the
     // reference processor does. Encoding the fragments separately would split merges

src/community_models/moss_voicegen/tokenizer_text.cpp

@@ -353,6 +353,64 @@ inline std::vector<float> causal_context_mask_window(
 
 namespace engine::codecs {
 
+MossTokenRowBuilder::MossTokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id)
+    : num_codebooks_(num_codebooks),
+      audio_pad_token_id_(audio_pad_token_id) {
+    if (num_codebooks_ <= 0) {
+        throw std::runtime_error("MOSS token row builder requires a positive codebook count");
+    }
+}
+
+void MossTokenRowBuilder::push_text_token(int32_t token_id) {
+    rows_.text_tokens.push_back(token_id);
+    rows_.audio_codes.insert(rows_.audio_codes.end(), static_cast<size_t>(num_codebooks_), audio_pad_token_id_);
+}
+
+void MossTokenRowBuilder::push_text_tokens(const std::vector<int32_t> & token_ids) {
+    for (const int32_t token_id : token_ids) {
+        push_text_token(token_id);
+    }
+}
+
+void MossTokenRowBuilder::push_audio_row(int32_t text_slot_token_id, const int32_t * codes, int64_t num_codebooks) {
+    if (num_codebooks != num_codebooks_) {
+        throw std::runtime_error("MOSS audio row codebook count mismatch");
+    }
+    if (codes == nullptr) {
+        throw std::runtime_error("MOSS audio row codes are missing");
+    }
+    rows_.text_tokens.push_back(text_slot_token_id);
+    rows_.audio_codes.insert(rows_.audio_codes.end(), codes, codes + num_codebooks);
+}
+
+void MossTokenRowBuilder::push_audio_row(
+    int32_t text_slot_token_id,
+    const std::vector<std::vector<int32_t>> & codes,
+    int64_t frame) {
+    if (static_cast<int64_t>(codes.size()) != num_codebooks_) {
+        throw std::runtime_error("MOSS audio row codebook count mismatch");
+    }
+    rows_.text_tokens.push_back(text_slot_token_id);
+    for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) {
+        const auto & channel = codes[static_cast<size_t>(codebook)];
+        if (frame < 0 || static_cast<size_t>(frame) >= channel.size()) {
+            throw std::runtime_error("MOSS audio row frame index is out of range");
+        }
+        rows_.audio_codes.push_back(channel[static_cast<size_t>(frame)]);
+    }
+}
+
+MossTokenRows MossTokenRowBuilder::finish() {
+    if (rows_.text_tokens.empty()) {
+        throw std::runtime_error("MOSS token rows must not be empty");
+    }
+    if (static_cast<int64_t>(rows_.audio_codes.size()) !=
+        static_cast<int64_t>(rows_.text_tokens.size()) * num_codebooks_) {
+        throw std::runtime_error("MOSS token rows audio code shape mismatch");
+    }
+    return std::move(rows_);
+}
+
 // Dequantizes MOSS-Audio-Tokenizer-v2 codes (RLFQ) into the codec's continuous
 // latent, i.e. the input to the codec decoder stack. Codes are the
 // [num_quantizers, steps] matrix produced by generation; the returned latent is

src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp

@@ -39,7 +39,7 @@ uint64_t mix_reference_audio_key(uint64_t key, uint64_t value) {
     return key;
 }
 
-uint64_t prefix_hash(const moss::TokenRows & prefix, int64_t num_codebooks) {
+uint64_t prefix_hash(const engine::codecs::MossTokenRows & prefix, int64_t num_codebooks) {
     uint64_t key = 1469598103934665603ull;
     for (size_t row = 0; row < prefix.text_tokens.size(); ++row) {
         key = mix_reference_audio_key(key, static_cast<uint32_t>(prefix.text_tokens[row]));
@@ -53,7 +53,7 @@ uint64_t prefix_hash(const moss::TokenRows & prefix, int64_t num_codebooks) {
     return key;
 }
 
-int64_t prefix_audio_nonpad_count(const moss::TokenRows & prefix, int32_t audio_pad_token_id) {
+int64_t prefix_audio_nonpad_count(const engine::codecs::MossTokenRows & prefix, int32_t audio_pad_token_id) {
     int64_t count = 0;
     for (const int32_t code : prefix.audio_codes) {
         if (code != audio_pad_token_id) {