holographic-memory 0.6.3

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::{ensure, Result};
use candle_core::{DType, Device, D};

use super::qwen3::Qwen3;
use super::{encode, encoder_batches, last_hidden, load, prompts, ModelSource};

/// Qwen3-Embedding: queries carry the task instruction, documents none; last-token pooling and
/// L2 normalization, inputs truncated to 2048 tokens (the reference `max_seq_length`).
pub struct Embedder {
    model: Qwen3,
    tokenizer: tokenizers::Tokenizer,
    device: Device,
    dim: usize,
    batch: usize,
    query_task: EmbeddingQueryTask,
}

/// Query instruction; document embeddings are identical for both tasks.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum EmbeddingQueryTask {
    #[default]
    UserMessages,
    Conversation,
}

const MAX_TOKENS: usize = 2048;
const BATCH_TOKENS: usize = 16384;
/// Default padded batch on Metal/CUDA: on an M4 (Metal), 96 user turns embedded in 3.5 s at
/// batch 4 against 4.3 s at 1, 3.7 s at 2, 4.0 s at 8 and 5.6 s at 16
/// (`benchmarks/results/local_models_batching.json`). CUDA is unmeasured.
const GPU_BATCH: usize = 4;

impl Embedder {
    pub fn load(source: &ModelSource, device: &Device) -> Result<Self> {
        let l = load(source, device, super::dtype_for(device, false))?;
        let mut tokenizer = l.tokenizer;
        tokenizer
            .with_truncation(Some(tokenizers::TruncationParams {
                max_length: MAX_TOKENS,
                ..Default::default()
            }))
            .map_err(|e| anyhow::anyhow!("tokenizer truncation: {e}"))?;
        tokenizer.with_padding(None);
        let dim = l.config.hidden_size;
        Ok(Self {
            model: Qwen3::new(&l.config, l.vb, false)?,
            tokenizer,
            device: device.clone(),
            dim,
            batch: super::default_encode_batch(device, GPU_BATCH),
            query_task: EmbeddingQueryTask::default(),
        })
    }

    pub fn set_query_task(&mut self, task: EmbeddingQueryTask) {
        self.query_task = task;
    }

    pub fn query_task(&self) -> EmbeddingQueryTask {
        self.query_task
    }

    pub fn dimensions(&self) -> usize {
        self.dim
    }

    /// Sets the most sequences per forward pass (at least 1).
    pub fn set_batch(&mut self, batch: usize) {
        self.batch = batch.max(1);
    }

    /// L2-normalized embeddings, one per text, in input order.
    pub fn embed<S: AsRef<str>>(&self, texts: &[S], is_query: bool) -> Result<Vec<Vec<f32>>> {
        let ids: Vec<Vec<u32>> = texts
            .iter()
            .map(|t| {
                let t = t.as_ref();
                if is_query {
                    let query = match self.query_task {
                        EmbeddingQueryTask::UserMessages => prompts::embed_query(t),
                        EmbeddingQueryTask::Conversation => prompts::conversation_embed_query(t),
                    };
                    encode(&self.tokenizer, &query, true)
                } else {
                    encode(&self.tokenizer, t, true)
                }
            })
            .collect::<Result<_>>()?;
        ensure!(ids.iter().all(|x| !x.is_empty()), "empty token sequence");
        let lens: Vec<usize> = ids.iter().map(Vec::len).collect();
        let mut out = vec![Vec::new(); texts.len()];
        for batch in encoder_batches(&self.device, &lens, BATCH_TOKENS, self.batch) {
            let last = last_hidden(&self.model, &ids, &batch)?.to_dtype(DType::F32)?;
            let norm = last
                .sqr()?
                .sum_keepdim(D::Minus1)?
                .sqrt()?
                .clamp(1e-12, f64::MAX)?;
            let rows = last.broadcast_div(&norm)?.to_vec2::<f32>()?;
            for (&i, row) in batch.iter().zip(rows) {
                out[i] = row;
            }
        }
        Ok(out)
    }
}