From 43a0be3d707b145bc5f0debd96424fa7a6318c6d Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 10 Jun 2026 22:48:47 -0700 Subject: [PATCH 1/2] added language model embed component --- cactus-engine/src/embed.cpp | 2 +- cactus-engine/src/engine.h | 6 + cactus-engine/src/model.cpp | 150 ++++++++++++++++++++-- python/cactus/transpile/component_plan.py | 2 +- python/cactus/transpile/hf_model.py | 1 + python/cactus/transpile/lower.py | 6 +- python/cactus/transpile/model_adapters.py | 130 ++++++++++++++++++- 7 files changed, 282 insertions(+), 15 deletions(-) diff --git a/cactus-engine/src/embed.cpp b/cactus-engine/src/embed.cpp index 3228fa6d7..6a380836a 100644 --- a/cactus-engine/src/embed.cpp +++ b/cactus-engine/src/embed.cpp @@ -56,7 +56,7 @@ int cactus_embed( return -1; } - std::vector embeddings = handle->model->get_embeddings(tokens, true, normalize); + std::vector embeddings = handle->model->get_text_embeddings(tokens, normalize); if (embeddings.empty()) { CACTUS_LOG_ERROR("embed", "Embedding returned empty result"); return -1; diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 471c3def9..beef0c4f1 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -652,6 +652,11 @@ class Model { std::vector get_embeddings(const std::vector& tokens, bool pooled = true, bool normalize = false, const std::string& profile_file = ""); + std::vector get_text_embeddings(const std::vector& tokens, bool normalize = false); + std::vector get_lm_embeddings(const std::vector& 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 get_image_embeddings(const std::string& image_path); std::vector get_audio_embeddings(const std::vector& 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; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 97b0cb2a6..0b5af7559 100644 --- a/cactus-engine/src/model.cpp +++ b/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,143 @@ void Model::maybe_roll_compact() { compress_kv_cache_keydiff(p); } +static void l2_normalize_inplace(std::vector& v) { + double norm = 0.0; + for (float x : v) norm += static_cast(x) * x; + float inv = static_cast(1.0 / std::max(std::sqrt(norm), 1e-12)); + for (float& x : v) x *= inv; +} + +static std::vector finalize_pooled_embedding(const std::vector& sum, size_t count, bool normalize) { + std::vector out(sum.size(), 0.0f); + if (count > 0) { + for (size_t h = 0; h < sum.size(); ++h) out[h] = static_cast(sum[h] / static_cast(count)); + } + if (normalize) l2_normalize_inplace(out); + return out; +} + +std::vector Model::get_text_embeddings(const std::vector& 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 Model::get_lm_embeddings(const std::vector& 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; + + // Recurrent-state caches (gated-deltanet) don't chain correctly across + // chunked-prefill boundaries, so cap them to a single chunk -- matches + // run_chunked_prefill. Such models embed up to effective_chunk tokens; longer + // inputs are truncated. KV/conv/sliding-window caches (qwen/gemma/lfm2) chain + // across chunks and span the whole sequence here. + 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(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(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(decoder_embed_->output_node_ids[out_idx]); + + std::vector 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(config_.pad_token_id); + write_int_input_at(*prefill_encoder_, "input_ids", i, static_cast(token)); + write_int_input_at(*prefill_encoder_, "position_ids", i, static_cast(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(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(reinterpret_cast(ptr)[i]) + : reinterpret_cast(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(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 Model::get_embeddings(const std::vector& tokens, bool pooled, bool normalize, const std::string& /*profile_file*/) { if (!components_.count("text_embedding")) { @@ -3250,12 +3389,7 @@ std::vector Model::get_embeddings(const std::vector& 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(v) * v; - float inv = static_cast(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(); diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index 837dda211..5c5df4114 100644 --- a/python/cactus/transpile/component_plan.py +++ b/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( diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index f0edbd570..43991440d 100644 --- a/python/cactus/transpile/hf_model.py +++ b/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", diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index 14e7eabea..e0e7a6604 100644 --- a/python/cactus/transpile/lower.py +++ b/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)) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 701fddd98..8299d9a09 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -1895,11 +1895,13 @@ def __init__( *, weights_dir: str | None = None, last_token_only: bool = True, + return_hidden: bool = False, ): super().__init__() self.model = model self.weights_dir = weights_dir self.last_token_only = bool(last_token_only) + self.return_hidden = bool(return_hidden) self.backbone = _lfm2_language_backbone(model) self.lm_head = getattr(model, "lm_head", None) if not isinstance(self.lm_head, torch.nn.Module): @@ -1939,6 +1941,8 @@ def forward( ) hidden_states = self.backbone.embedding_norm(hidden_states) + if self.return_hidden: + return hidden_states if self.last_token_only: hidden_states = hidden_states[:, -1:, :] return self.lm_head(hidden_states) @@ -3341,6 +3345,42 @@ def forward( return _gemma4_apply_final_logit_softcapping(self.model, logits) +class Gemma4DecoderEmbedChunkAdapter(Gemma4DecoderAdapter): + """Embedding-readout variant of the prefill chunk: runs the full backbone over + ALL tokens with final norm (no shared-KV-tail last-token shortcut) and returns + the full-sequence last hidden state (no lm_head).""" + + def forward( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor, + position_ids: torch.LongTensor, + ) -> torch.Tensor: + normalized_per_layer_inputs = per_layer_inputs + if normalized_per_layer_inputs.numel() == 0: + normalized_per_layer_inputs = None + attention_mask = torch.ones( + position_ids.shape, + dtype=torch.long, + device=position_ids.device, + ) + causal_mask_mapping = _gemma4_build_standard_causal_mask_mapping( + create_causal_mask=self._create_causal_mask, + create_sliding_window_causal_mask=self._create_sliding_window_causal_mask, + config=self.backbone.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + ) + return _gemma4_text_backbone_forward( + self.backbone, + inputs_embeds=inputs_embeds, + per_layer_inputs=normalized_per_layer_inputs, + causal_mask_mapping=causal_mask_mapping, + position_ids=position_ids, + ) + + def _build_gemma4_multimodal_component_specs( model: torch.nn.Module, *, @@ -3373,6 +3413,7 @@ def _require(component: str) -> None: _require("audio_encoder") _require("lm_encoder") _require("decoder_prefill_chunk") + _require("decoder_embed_chunk") _require("lm_encoder_text_chunk") _require("lm_encoder_media_chunk") _require("lm_encoder_step") @@ -3464,6 +3505,7 @@ def _require(component: str) -> None: lm_encoder_media_chunk = Gemma4LMEncoderMediaChunkAdapter(model, weights_dir=weights_dir).eval() decoder = Gemma4DecoderAdapter(model, weights_dir=weights_dir).eval() decoder_prefill = Gemma4DecoderPrefillChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_embed = Gemma4DecoderEmbedChunkAdapter(model, weights_dir=weights_dir).eval() decoder_step = Gemma4DecoderStepAdapter(model, weights_dir=weights_dir).eval() with torch.no_grad(): @@ -3610,6 +3652,34 @@ def _require(component: str) -> None: }, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, )) + if "decoder_embed_chunk" in expanded_components: + if decoder_inputs is not None: + embed_chunk_inputs = tuple( + value[:, :prefill_chunk_size, ...].contiguous() + if value.ndim >= 2 + else value + for value in decoder_inputs + ) + elif prefill_decoder_inputs is not None: + embed_chunk_inputs = tuple(prefill_decoder_inputs) + else: + raise RuntimeError("Gemma4 decoder_embed_chunk spec requires precomputed decoder inputs") + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed, + example_inputs=embed_chunk_inputs, + input_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, + output_keys=("last_hidden_state",), + graph_meta={ + **common_graph_meta, + "component": "decoder_embed_chunk", + "use_internal_kv_cache": True, + "max_cache_seq_len": cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, + )) if "lm_encoder_step" in expanded_components: if step_input_ids is None or step_position_ids is None: raise RuntimeError("Gemma4 lm_encoder_step spec requires step token inputs") @@ -4308,7 +4378,7 @@ def __init__(self, model: torch.nn.Module): self.model = model self.backbone = model.model - def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + def _encode_hidden_states(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: seq_len = int(inputs_embeds.shape[1]) allowed_positions = torch.tril( torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), @@ -4337,7 +4407,10 @@ def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> to past_key_values=None, use_cache=False, ) - hidden_states = self.backbone.norm(hidden_states) + return self.backbone.norm(hidden_states) + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = self._encode_hidden_states(inputs_embeds, position_ids) return self.model.lm_head(hidden_states[:, -1:, :]) def get_transpile_metadata(self): @@ -4362,6 +4435,14 @@ class Qwen3EmbedsCausalLMPrefillChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): pass +class Qwen3EmbedsCausalLMEmbedChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): + """Like the prefill chunk, but emits the full-sequence last hidden state + (no lm_head, no last-token slice) for embedding readout.""" + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + return self._encode_hidden_states(inputs_embeds, position_ids) + + def _build_qwen_causal_lm_component_specs( model: torch.nn.Module, *, @@ -4392,7 +4473,7 @@ def _build_qwen_causal_lm_component_specs( "adapter_family": family, } - chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk"} + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk", "decoder_embed_chunk"} wants_chunked = bool(chunk_components & requested_set) if wants_chunked and family != "qwen3": raise RuntimeError( @@ -4437,6 +4518,7 @@ def _build_qwen_causal_lm_component_specs( lm_encoder_text_chunk = Qwen3LMEncoderTextChunkAdapter(model).eval() decoder_media_step = Qwen3EmbedsCausalLMStepAdapter(model).eval() decoder_prefill_chunk = Qwen3EmbedsCausalLMPrefillChunkAdapter(model).eval() + decoder_embed_chunk = Qwen3EmbedsCausalLMEmbedChunkAdapter(model).eval() max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) @@ -4492,6 +4574,25 @@ def _build_qwen_causal_lm_component_specs( }, metadata={"family": family, "task": "causal_lm_logits"}, )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("last_hidden_state",), + graph_meta={ + **common_graph_meta, + "component": "decoder_embed_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) if "decoder_media_step" in requested_set: specs.append(ComponentModuleSpec( component="decoder_media_step", @@ -5533,6 +5634,7 @@ def _require(component: str) -> None: _require("lm_encoder_step") _require("lm_encoder_text_chunk") _require("decoder_prefill_chunk") + _require("decoder_embed_chunk") _require("decoder_step") if "lm_encoder" in requested_set: _require("vision_encoder") @@ -5567,6 +5669,7 @@ def _require(component: str) -> None: lm_encoder_text_chunk = Lfm2VlLMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() decoder = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False).eval() decoder_last_token = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=True).eval() + decoder_embed = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False, return_hidden=True).eval() image_features: torch.Tensor | None = None decoder_inputs: tuple[torch.Tensor, ...] | None = None @@ -5702,6 +5805,27 @@ def _require(component: str) -> None: }, metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, )) + if "decoder_embed_chunk" in expanded_components: + if decoder_inputs is None: + if image_features is None: + image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) + decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed, + example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("last_hidden_state",), + graph_meta={ + **common_graph_meta, + "component": "decoder_embed_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + }, + metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + )) if "decoder_step" in expanded_components: if decoder_inputs is None: if image_features is None: From 555a4996322e61c1fb8f12b0cbb265743114d226 Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 10 Jun 2026 22:54:58 -0700 Subject: [PATCH 2/2] removed server embedding gate for nomic, cleaned up code --- cactus-engine/src/model.cpp | 7 ------- python/cactus/server.py | 18 ++++++++++++++++-- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 0b5af7559..5e69eb74f 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -3230,11 +3230,6 @@ std::vector Model::get_lm_embeddings(const std::vector& tokens, } const size_t effective_chunk = component_tokens; - // Recurrent-state caches (gated-deltanet) don't chain correctly across - // chunked-prefill boundaries, so cap them to a single chunk -- matches - // run_chunked_prefill. Such models embed up to effective_chunk tokens; longer - // inputs are truncated. KV/conv/sliding-window caches (qwen/gemma/lfm2) chain - // across chunks and span the whole sequence here. bool recurrent_state = false; if (decoder_embed_->graph) { for (const auto& state : decoder_embed_->cache_states) { @@ -3330,8 +3325,6 @@ std::vector Model::get_embeddings(const std::vector& 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 wrapped; wrapped.reserve(tokens.size() + 2); if (tokenizer_) wrapped.push_back(tokenizer_->get_bos_token()); diff --git a/python/cactus/server.py b/python/cactus/server.py index b7d55c117..792fcc671 100644 --- a/python/cactus/server.py +++ b/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: