added language model embed component

cactus-engine/src/embed.cpp

@@ -56,7 +56,7 @@ int cactus_embed(
             return -1;
         }
 
-        std::vector<float> embeddings = handle->model->get_embeddings(tokens, true, normalize);
+        std::vector<float> embeddings = handle->model->get_text_embeddings(tokens, normalize);
         if (embeddings.empty()) {
             CACTUS_LOG_ERROR("embed", "Embedding returned empty result");
             return -1;

cactus-engine/src/engine.h

@@ -652,6 +652,11 @@ class Model {
     std::vector<float> get_embeddings(const std::vector<uint32_t>& tokens, bool pooled = true,
                                        bool normalize = false, const std::string& profile_file = "");
 
+    std::vector<float> get_text_embeddings(const std::vector<uint32_t>& tokens, bool normalize = false);
+    std::vector<float> get_lm_embeddings(const std::vector<uint32_t>& tokens, bool normalize = false);
+    bool has_lm_embedding() const { return decoder_embed_ != nullptr; }
+    bool has_text_embedding() const { return components_.count("text_embedding") > 0; }
+
     std::vector<float> get_image_embeddings(const std::string& image_path);
 
     std::vector<float> get_audio_embeddings(const std::vector<float>& audio_features);
@@ -803,6 +808,7 @@ class Model {
     Component* encoder_ = nullptr;
     Component* decoder_ = nullptr;
     Component* decoder_prefill_ = nullptr;
+    Component* decoder_embed_ = nullptr;
     Component* prefill_encoder_ = nullptr;
     enum class DecodeRoute { CACHED_STEP, DIRECT_DECODER_STEP, FULL_CONTEXT_TEXT, ENCODER_CROSS_KV_STEP };
     DecodeRoute decode_route_ = DecodeRoute::CACHED_STEP;

cactus-engine/src/model.cpp

@@ -454,8 +454,6 @@ bool Model::init(const std::string& bundle_dir, size_t context_size,
         && !components_.count("decoder_step")
         && !components_.count("lm_encoder_step");
     if (is_text_embedding) {
-        // Embedding-only bundle: no decode route; get_embeddings loads the
-        // text_embedding component on demand.
         cache_max_seq_len_ = context_size;
         initialized_ = true;
         return true;
@@ -543,6 +541,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size,
              "audio_encoder",
              "lm_encoder_media_step",
              "decoder_prefill_chunk",
+             "decoder_embed_chunk",
              "lm_encoder",
          }) {
         if (components_.count(optional)) {
@@ -558,6 +557,9 @@ bool Model::init(const std::string& bundle_dir, size_t context_size,
         decoder_prefill_ = &components_.at("decoder_prefill_chunk");
         decoder_prefill_chunk_ = decoder_prefill_;
     }
+    if (components_.count("decoder_embed_chunk") && components_.at("decoder_embed_chunk").graph) {
+        decoder_embed_ = &components_.at("decoder_embed_chunk");
+    }
     if (components_.count("lm_encoder_text_chunk") && components_.at("lm_encoder_text_chunk").graph) {
         prefill_encoder_ = &components_.at("lm_encoder_text_chunk");
     }
@@ -3178,6 +3180,138 @@ void Model::maybe_roll_compact() {
     compress_kv_cache_keydiff(p);
 }
 
+static void l2_normalize_inplace(std::vector<float>& v) {
+    double norm = 0.0;
+    for (float x : v) norm += static_cast<double>(x) * x;
+    float inv = static_cast<float>(1.0 / std::max(std::sqrt(norm), 1e-12));
+    for (float& x : v) x *= inv;
+}
+
+static std::vector<float> finalize_pooled_embedding(const std::vector<double>& sum, size_t count, bool normalize) {
+    std::vector<float> out(sum.size(), 0.0f);
+    if (count > 0) {
+        for (size_t h = 0; h < sum.size(); ++h) out[h] = static_cast<float>(sum[h] / static_cast<double>(count));
+    }
+    if (normalize) l2_normalize_inplace(out);
+    return out;
+}
+
+std::vector<float> Model::get_text_embeddings(const std::vector<uint32_t>& tokens, bool normalize) {
+    if (has_text_embedding()) {
+        return get_embeddings(tokens, /*pooled=*/true, normalize);
+    }
+    if (has_lm_embedding()) {
+        return get_lm_embeddings(tokens, normalize);
+    }
+    throw std::runtime_error("get_text_embeddings: bundle has neither a text_embedding nor a decoder_embed_chunk component");
+}
+
+std::vector<float> Model::get_lm_embeddings(const std::vector<uint32_t>& tokens, bool normalize) {
+    if (!decoder_embed_) {
+        throw std::runtime_error("get_lm_embeddings: bundle has no decoder_embed_chunk component");
+    }
+    if (tokens.empty()) {
+        throw std::runtime_error("get_lm_embeddings: empty token sequence");
+    }
+    if (decode_route_ != DecodeRoute::CACHED_STEP || !encoder_) {
+        throw std::runtime_error("get_lm_embeddings: model does not support chunked LM embeddings");
+    }
+    if (!load_component_graph(*decoder_embed_)) {
+        throw std::runtime_error("get_lm_embeddings: failed to load decoder_embed_chunk graph");
+    }
+    if (prefill_encoder_ && !load_component_graph(*prefill_encoder_)) {
+        throw std::runtime_error("get_lm_embeddings: failed to load prefill encoder graph");
+    }
+    reset_component_cache_states(*decoder_embed_);
+
+    const size_t component_tokens = component_chunk_tokens(*decoder_embed_, "inputs_embeds");
+    if (component_tokens <= 1) {
+        throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk is not chunk-shaped");
+    }
+    const size_t effective_chunk = component_tokens;
+
+    bool recurrent_state = false;
+    if (decoder_embed_->graph) {
+        for (const auto& state : decoder_embed_->cache_states) {
+            for (int node_id : {state.key_node_id, state.value_node_id}) {
+                if (node_id < 0) continue;
+                if (decoder_embed_->graph->get_node_op_type(static_cast<size_t>(node_id)) == OpType::RECURRENT_CACHE_STATE) {
+                    recurrent_state = true;
+                }
+            }
+        }
+    }
+    const size_t embed_limit = recurrent_state ? std::min(tokens.size(), effective_chunk) : tokens.size();
+
+    size_t encoder_chunk = 0;
+    if (prefill_encoder_ && input_index(*prefill_encoder_, "input_ids") >= 0 && input_index(*prefill_encoder_, "position_ids") >= 0) {
+        encoder_chunk = component_chunk_tokens(*prefill_encoder_, "input_ids");
+        if (encoder_chunk == 0 || effective_chunk % encoder_chunk != 0) encoder_chunk = 0;
+    }
+
+    const int out_idx = output_index(*decoder_embed_, "last_hidden_state");
+    if (out_idx < 0 || static_cast<size_t>(out_idx) >= decoder_embed_->output_node_ids.size()) {
+        throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk missing last_hidden_state output");
+    }
+    const size_t out_node = static_cast<size_t>(decoder_embed_->output_node_ids[out_idx]);
+
+    std::vector<double> sum;
+    size_t count = 0;
+
+    size_t processed = 0;
+    while (processed < embed_limit) {
+        for (auto& buf : decoder_embed_->input_buffers) std::fill(buf.begin(), buf.end(), 0);
+        if (encoder_chunk > 0) {
+            for (size_t chunk_offset = 0; chunk_offset < effective_chunk; chunk_offset += encoder_chunk) {
+                for (auto& buf : prefill_encoder_->input_buffers) std::fill(buf.begin(), buf.end(), 0);
+                for (size_t i = 0; i < encoder_chunk; ++i) {
+                    const size_t index = processed + chunk_offset + i;
+                    const uint32_t token = index < tokens.size() ? tokens[index] : static_cast<uint32_t>(config_.pad_token_id);
+                    write_int_input_at(*prefill_encoder_, "input_ids", i, static_cast<int64_t>(token));
+                    write_int_input_at(*prefill_encoder_, "position_ids", i, static_cast<int64_t>(processed + chunk_offset + i));
+                }
+                prefill_encoder_->graph->execute();
+                copy_component_outputs_to_chunk_inputs_range(*prefill_encoder_, *decoder_embed_, chunk_offset);
+            }
+        } else {
+            for (size_t i = 0; i < effective_chunk; ++i) {
+                const size_t index = processed + i;
+                const uint32_t token = index < tokens.size() ? tokens[index] : static_cast<uint32_t>(config_.pad_token_id);
+                run_encoder_step(token, processed + i);
+                copy_component_outputs_to_chunk_inputs(*encoder_, *decoder_embed_, i);
+            }
+        }
+        decoder_embed_->graph->execute();
+
+        const auto& desc = decoder_embed_->graph->get_output_buffer(out_node);
+        void* ptr = decoder_embed_->graph->get_output(out_node);
+        const size_t hidden = desc.shape.empty() ? 0 : desc.shape.back();
+        const size_t seq = (desc.shape.size() >= 2) ? desc.shape[desc.shape.size() - 2] : 1;
+        if (hidden == 0 || !ptr) {
+            throw std::runtime_error("get_lm_embeddings: decoder_embed_chunk produced no last_hidden_state");
+        }
+        if (sum.empty()) sum.assign(hidden, 0.0);
+        const bool is_fp16 = desc.precision == Precision::FP16;
+        auto read_at = [&](size_t i) -> float {
+            return is_fp16 ? static_cast<float>(reinterpret_cast<const __fp16*>(ptr)[i])
+                           : reinterpret_cast<const float*>(ptr)[i];
+        };
+        size_t real_rows = std::min(effective_chunk, tokens.size() - processed);
+        real_rows = std::min(real_rows, seq);
+        for (size_t t = 0; t < real_rows; ++t) {
+            for (size_t h = 0; h < hidden; ++h) sum[h] += static_cast<double>(read_at(t * hidden + h));
+        }
+        count += real_rows;
+        processed += effective_chunk;
+    }
+
+    reset_component_cache_states(*decoder_embed_);
+    decoder_embed_->graph->release_runtime_buffers();
+    decoder_embed_->graph->release_all_weight_pages();
+    unload_component_graph(*decoder_embed_);
+    return finalize_pooled_embedding(sum, count, normalize);
+}
+
 std::vector<float> Model::get_embeddings(const std::vector<uint32_t>& tokens, bool pooled,
                                           bool normalize, const std::string& /*profile_file*/) {
     if (!components_.count("text_embedding")) {
@@ -3191,8 +3325,6 @@ std::vector<float> Model::get_embeddings(const std::vector<uint32_t>& tokens, bo
         throw std::runtime_error("get_embeddings: failed to load embedding component graph");
     }
 
-    // Embedding encoders (nomic / XLM-R) wrap the sequence with BOS/EOS, matching
-    // the reference tokenizer's add_special_tokens behavior.
     std::vector<uint32_t> wrapped;
     wrapped.reserve(tokens.size() + 2);
     if (tokenizer_) wrapped.push_back(tokenizer_->get_bos_token());
@@ -3250,12 +3382,7 @@ std::vector<float> Model::get_embeddings(const std::vector<uint32_t>& tokens, bo
         for (size_t h = 0; h < hidden; ++h) result[h] = read_at(h);
     }
 
-    if (normalize) {
-        double norm = 0.0;
-        for (float v : result) norm += static_cast<double>(v) * v;
-        float inv = static_cast<float>(1.0 / std::max(std::sqrt(norm), 1e-12));
-        for (float& v : result) v *= inv;
-    }
+    if (normalize) l2_normalize_inplace(result);
 
     comp->graph->release_runtime_buffers();
     comp->graph->release_all_weight_pages();

python/cactus/server.py

@@ -31,7 +31,6 @@
 
 LLM_MODEL_TYPES = {"gemma", "gemma3n", "gemma4", "lfm2", "qwen", "qwen3p5", "needle", "youtu"}
 STT_MODEL_TYPES = {"whisper", "parakeet_tdt", "parakeet-tdt"}
-EMBED_MODEL_TYPES = {"bert", "nomic"}
 
 
 @dataclass(frozen=True)
@@ -41,6 +40,7 @@ class ModelInfo:
     model_type: str
     context_length: int
     created: int
+    supports_embedding: bool = False
 
 
 class ModelRegistry:
@@ -65,6 +65,19 @@ def _read_config_field(model_dir: Path, field: str) -> str:
                 return line.split("=", 1)[1].strip()
         return ""
 
+    @staticmethod
+    def _supports_embedding(model_dir: Path) -> bool:
+        """A bundle can produce text embeddings if it carries a text_embedding
+        component (Nomic/BERT encoder) or a decoder_embed_chunk component
+        (causal-LM hidden-state embeddings)."""
+        manifest = model_dir / "components" / "manifest.json"
+        try:
+            data = json.loads(manifest.read_text(encoding="utf-8"))
+        except Exception:
+            return False
+        names = {str(c.get("component", "")) for c in data.get("components", []) if isinstance(c, dict)}
+        return "text_embedding" in names or "decoder_embed_chunk" in names
+
     @classmethod
     def _info_for_dir(cls, path: Path) -> ModelInfo | None:
         display_id = path.expanduser().name
@@ -85,6 +98,7 @@ def _info_for_dir(cls, path: Path) -> ModelInfo | None:
             model_type=cls._read_config_field(resolved, "model_type"),
             context_length=context_length,
             created=int(stat.st_mtime),
+            supports_embedding=cls._supports_embedding(resolved),
         )
 
     def _discover(self, root: Path) -> None:
@@ -616,7 +630,7 @@ async def create_transcription(
     async def create_embeddings(request: Request, req: EmbeddingRequest):
         reg: ModelRegistry = request.app.state.registry
         info = reg.require(req.model)
-        if info.model_type not in EMBED_MODEL_TYPES:
+        if not info.supports_embedding:
             raise HTTPException(status_code=400, detail=f"Model '{req.model}' is not an embedding model")
         inputs = [req.input] if isinstance(req.input, str) else list(req.input)
         if not inputs:

python/cactus/transpile/component_plan.py

@@ -34,7 +34,7 @@ def _plan_from_profile(
                 return ComponentPlan(
                     task="causal_lm_logits",
                     components=("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk",
-                                "decoder_prefill_chunk", "decoder_media_step"),
+                                "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step"),
                     force_component_pipeline=True,
                 )
             return ComponentPlan(

python/cactus/transpile/hf_model.py

@@ -317,6 +317,7 @@ def _write_component_bundle(
             "lm_encoder",
             "lm_encoder_text_chunk",
             "decoder_prefill_chunk",
+            "decoder_embed_chunk",
             "decoder",
             "decoder_cross_kv",
             "lm_encoder_step",

python/cactus/transpile/lower.py

@@ -219,7 +219,8 @@ def _should_lower_attention_with_internal_kv_cache(ir: IRGraph, node: IRNode) ->
     if node.op not in {"attention", "scaled_dot_product_attention"}:
         return False
     component = str(ir.meta.get("component", "") or "").strip().lower()
-    if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step"}:
+    if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step",
+                         "decoder_embed_chunk"}:
         return False
     if len(node.inputs) < 3:
         return False
@@ -390,7 +391,8 @@ def _should_lower_conv1d_with_internal_conv_cache(ir: IRGraph, node: IRNode, x:
     if node.op != "conv1d":
         return False
     component = str(ir.meta.get("component", "") or "").strip().lower()
-    if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step"}:
+    if component not in {"decoder_step", "decoder_prefill_chunk", "decoder_media_step",
+                         "decoder_embed_chunk"}:
         return False
     stride = int(node.attrs.get("stride", 1))
     padding = int(node.attrs.get("padding", 0))