amberjs 1.17.0

//! Amber Native AI Inference Engine powered by HuggingFace Candle
//!
//! Provides local GGUF / SafeTensors loading, Apple Silicon Metal GPU acceleration,
//! autoregressive KV-cache decode, zero-copy Tensor math, and stream decoding.

use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use anyhow::{bail, Context, Result};
use candle_core::quantized::gguf_file;
use candle_core::{Device, Tensor};
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::quantized_llama::ModelWeights as QLlama;
use candle_transformers::models::quantized_qwen2::ModelWeights as QQwen2;
use once_cell::sync::Lazy;
use tokenizers::Tokenizer;

use crate::ai_engine::generator::{GenerateOptions, GenerateResult};

/// Global counter for active model instances
static MODEL_ID_COUNTER: AtomicU64 = AtomicU64::new(1);

/// Global model registry holding loaded Candle models
static MODEL_REGISTRY: Lazy<Mutex<HashMap<u64, Arc<Mutex<CandleModel>>>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

/// Returns the best available compute device (Metal on macOS, CUDA on Linux, CPU fallback)
pub fn get_device(pref: Option<&str>) -> Device {
    match pref.map(|s| s.to_lowercase()).as_deref() {
        Some("metal") | Some("mps") => {
            #[cfg(feature = "metal")]
            {
                Device::new_metal(0).unwrap_or(Device::Cpu)
            }
            #[cfg(not(feature = "metal"))]
            {
                Device::Cpu
            }
        }
        Some("cuda") | Some("gpu") => {
            #[cfg(feature = "cuda")]
            {
                Device::new_cuda(0).unwrap_or(Device::Cpu)
            }
            #[cfg(not(feature = "cuda"))]
            {
                Device::Cpu
            }
        }
        Some("cpu") => Device::Cpu,
        _ => {
            // Auto-detection
            #[cfg(feature = "metal")]
            {
                if let Ok(dev) = Device::new_metal(0) {
                    return dev;
                }
            }
            #[cfg(feature = "cuda")]
            {
                if let Ok(dev) = Device::new_cuda(0) {
                    return dev;
                }
            }
            Device::Cpu
        }
    }
}

/// Supported model architecture variants
pub enum ModelArchitecture {
    Llama(QLlama),
    Qwen2(QQwen2),
    FallbackEdge,
}

/// A loaded local Candle model
pub struct CandleModel {
    pub id: u64,
    pub path: String,
    pub device: Device,
    pub arch: ModelArchitecture,
    pub tokenizer: Option<Tokenizer>,
    pub metadata: HashMap<String, serde_json::Value>,
    pub context_length: usize,
    pub eos_token_id: u32,
}

impl CandleModel {
    /// Loads a local model from a file path or creates an edge fallback
    pub fn load(path_str: &str, options_json: &str) -> Result<Arc<Mutex<Self>>> {
        let opts: serde_json::Value =
            serde_json::from_str(options_json).unwrap_or(serde_json::json!({}));

        let device_str = opts.get("device").and_then(|v| v.as_str());
        let device = get_device(device_str);

        let path = Path::new(path_str);
        let id = MODEL_ID_COUNTER.fetch_add(1, Ordering::SeqCst);

        // If path doesn't exist on disk, fallback to EdgeGenerator
        if !path.exists() {
            let model = Self {
                id,
                path: path_str.to_string(),
                device,
                arch: ModelArchitecture::FallbackEdge,
                tokenizer: None,
                metadata: HashMap::new(),
                context_length: 4096,
                eos_token_id: 2,
            };
            let arc_model = Arc::new(Mutex::new(model));
            let mut reg = MODEL_REGISTRY.lock().unwrap();
            reg.insert(id, arc_model.clone());
            return Ok(arc_model);
        }

        let mut file =
            File::open(path).with_context(|| format!("Failed to open model at {:?}", path))?;

        // Read GGUF content
        let content = gguf_file::Content::read(&mut file)
            .with_context(|| format!("Failed to read GGUF metadata from {:?}", path))?;

        // Extract metadata
        let mut metadata = HashMap::new();
        for (k, v) in &content.metadata {
            let json_val = match v {
                gguf_file::Value::U8(n) => serde_json::json!(n),
                gguf_file::Value::I8(n) => serde_json::json!(n),
                gguf_file::Value::U16(n) => serde_json::json!(n),
                gguf_file::Value::I16(n) => serde_json::json!(n),
                gguf_file::Value::U32(n) => serde_json::json!(n),
                gguf_file::Value::I32(n) => serde_json::json!(n),
                gguf_file::Value::F32(n) => serde_json::json!(n),
                gguf_file::Value::U64(n) => serde_json::json!(n),
                gguf_file::Value::I64(n) => serde_json::json!(n),
                gguf_file::Value::F64(n) => serde_json::json!(n),
                gguf_file::Value::Bool(b) => serde_json::json!(b),
                gguf_file::Value::String(s) => serde_json::json!(s),
                gguf_file::Value::Array(arr) => serde_json::json!(arr.len()),
            };
            metadata.insert(k.clone(), json_val);
        }

        let arch_name = metadata
            .get("general.architecture")
            .and_then(|v| v.as_str())
            .unwrap_or("llama")
            .to_lowercase();

        let context_length = metadata
            .get(&format!("{}.context_length", arch_name))
            .and_then(|v| v.as_u64())
            .unwrap_or(4096) as usize;

        let eos_token_id = metadata
            .get("tokenizer.ggml.eos_token_id")
            .and_then(|v| v.as_u64())
            .unwrap_or(2) as u32;

        // Try to load tokenizer
        let tokenizer = Self::find_and_load_tokenizer(path);

        // Load quantized weights based on architecture
        let arch = if arch_name.contains("qwen") {
            let model = QQwen2::from_gguf(content, &mut file, &device)
                .with_context(|| "Failed to construct quantized Qwen2 model")?;
            ModelArchitecture::Qwen2(model)
        } else {
            let model = QLlama::from_gguf(content, &mut file, &device)
                .with_context(|| "Failed to construct quantized Llama model")?;
            ModelArchitecture::Llama(model)
        };

        let candle_model = Self {
            id,
            path: path_str.to_string(),
            device,
            arch,
            tokenizer,
            metadata,
            context_length,
            eos_token_id,
        };

        let arc_model = Arc::new(Mutex::new(candle_model));
        let mut reg = MODEL_REGISTRY.lock().unwrap();
        reg.insert(id, arc_model.clone());

        Ok(arc_model)
    }

    /// Searches for tokenizer.json adjacent to model file
    fn find_and_load_tokenizer(model_path: &Path) -> Option<Tokenizer> {
        let candidates = [
            model_path.with_file_name("tokenizer.json"),
            model_path.with_extension("tokenizer.json"),
            PathBuf::from("tokenizer.json"),
        ];

        for cand in &candidates {
            if cand.exists() {
                if let Ok(tok) = Tokenizer::from_file(cand) {
                    return Some(tok);
                }
            }
        }
        None
    }

    /// Autoregressive text generation
    pub fn generate(&mut self, prompt: &str, options: &GenerateOptions) -> Result<GenerateResult> {
        let mut full_text = String::new();
        let mut count = 0;

        let res = self.generate_stream(prompt, options, |token_str| {
            full_text.push_str(token_str);
            count += 1;
            true
        })?;

        Ok(GenerateResult {
            text: full_text,
            tokens_generated: count,
            finish_reason: res.finish_reason,
        })
    }

    /// Streaming text generation with callback per token
    pub fn generate_stream<F>(
        &mut self,
        prompt: &str,
        options: &GenerateOptions,
        callback: F,
    ) -> Result<GenerateResult>
    where
        F: FnMut(&str) -> bool,
    {
        match &mut self.arch {
            ModelArchitecture::FallbackEdge => {
                crate::ai_engine::EdgeGenerator::generate_stream(prompt, options, callback)
            }
            ModelArchitecture::Llama(model) => Self::run_llama_stream(
                model,
                self.tokenizer.as_ref(),
                &self.device,
                self.eos_token_id,
                prompt,
                options,
                callback,
            ),
            ModelArchitecture::Qwen2(model) => Self::run_qwen_stream(
                model,
                self.tokenizer.as_ref(),
                &self.device,
                self.eos_token_id,
                prompt,
                options,
                callback,
            ),
        }
    }

    fn run_llama_stream<F>(
        model: &mut QLlama,
        tokenizer: Option<&Tokenizer>,
        device: &Device,
        eos_id: u32,
        prompt: &str,
        options: &GenerateOptions,
        mut callback: F,
    ) -> Result<GenerateResult>
    where
        F: FnMut(&str) -> bool,
    {
        let seed = 299792458;
        let temp = if options.temperature <= 0.0 {
            None
        } else {
            Some(options.temperature as f64)
        };
        let top_p = if options.top_p <= 0.0 || options.top_p >= 1.0 {
            None
        } else {
            Some(options.top_p as f64)
        };
        let mut logits_processor = LogitsProcessor::new(seed, temp, top_p);

        let input_tokens = if let Some(tok) = tokenizer {
            let encoding = tok
                .encode(prompt, true)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            encoding.get_ids().to_vec()
        } else {
            prompt.bytes().map(|b| b as u32).collect()
        };

        if input_tokens.is_empty() {
            return Ok(GenerateResult {
                text: String::new(),
                tokens_generated: 0,
                finish_reason: "stop".to_string(),
            });
        }

        let mut index_pos = 0;
        let input = Tensor::new(&input_tokens[..], device)?.unsqueeze(0)?;
        let mut logits = model.forward(&input, index_pos)?;
        logits = logits.squeeze(0)?;
        let mut next_token = logits_processor.sample(&logits)?;
        index_pos += input_tokens.len();

        let mut generated_count = 0;
        let max_tokens = options.max_tokens.max(1);
        let mut finish_reason = "length".to_string();

        for _ in 0..max_tokens {
            generated_count += 1;

            if next_token == eos_id {
                finish_reason = "stop".to_string();
                break;
            }

            let token_str = if let Some(tok) = tokenizer {
                tok.decode(&[next_token], false).unwrap_or_default()
            } else {
                String::from_utf8_lossy(&[next_token as u8]).to_string()
            };

            let should_stop = options.stop_sequences.iter().any(|s| token_str.contains(s));
            if should_stop {
                finish_reason = "stop".to_string();
                break;
            }

            if !callback(&token_str) {
                finish_reason = "interrupted".to_string();
                break;
            }

            let next_input = Tensor::new(&[next_token], device)?.unsqueeze(0)?;
            logits = model.forward(&next_input, index_pos)?.squeeze(0)?;
            index_pos += 1;
            next_token = logits_processor.sample(&logits)?;
        }

        Ok(GenerateResult {
            text: String::new(),
            tokens_generated: generated_count,
            finish_reason,
        })
    }

    fn run_qwen_stream<F>(
        model: &mut QQwen2,
        tokenizer: Option<&Tokenizer>,
        device: &Device,
        eos_id: u32,
        prompt: &str,
        options: &GenerateOptions,
        mut callback: F,
    ) -> Result<GenerateResult>
    where
        F: FnMut(&str) -> bool,
    {
        let seed = 299792458;
        let temp = if options.temperature <= 0.0 {
            None
        } else {
            Some(options.temperature as f64)
        };
        let top_p = if options.top_p <= 0.0 || options.top_p >= 1.0 {
            None
        } else {
            Some(options.top_p as f64)
        };
        let mut logits_processor = LogitsProcessor::new(seed, temp, top_p);

        let input_tokens = if let Some(tok) = tokenizer {
            let encoding = tok
                .encode(prompt, true)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            encoding.get_ids().to_vec()
        } else {
            prompt.bytes().map(|b| b as u32).collect()
        };

        if input_tokens.is_empty() {
            return Ok(GenerateResult {
                text: String::new(),
                tokens_generated: 0,
                finish_reason: "stop".to_string(),
            });
        }

        let mut index_pos = 0;
        let input = Tensor::new(&input_tokens[..], device)?.unsqueeze(0)?;
        let mut logits = model.forward(&input, index_pos)?;
        logits = logits.squeeze(0)?;
        let mut next_token = logits_processor.sample(&logits)?;
        index_pos += input_tokens.len();

        let mut generated_count = 0;
        let max_tokens = options.max_tokens.max(1);
        let mut finish_reason = "length".to_string();

        for _ in 0..max_tokens {
            generated_count += 1;

            if next_token == eos_id {
                finish_reason = "stop".to_string();
                break;
            }

            let token_str = if let Some(tok) = tokenizer {
                tok.decode(&[next_token], false).unwrap_or_default()
            } else {
                String::from_utf8_lossy(&[next_token as u8]).to_string()
            };

            let should_stop = options.stop_sequences.iter().any(|s| token_str.contains(s));
            if should_stop {
                finish_reason = "stop".to_string();
                break;
            }

            if !callback(&token_str) {
                finish_reason = "interrupted".to_string();
                break;
            }

            let next_input = Tensor::new(&[next_token], device)?.unsqueeze(0)?;
            logits = model.forward(&next_input, index_pos)?.squeeze(0)?;
            index_pos += 1;
            next_token = logits_processor.sample(&logits)?;
        }

        Ok(GenerateResult {
            text: String::new(),
            tokens_generated: generated_count,
            finish_reason,
        })
    }
}

/// Retrieves a model by its unique handle ID
pub fn get_model(id: u64) -> Option<Arc<Mutex<CandleModel>>> {
    let reg = MODEL_REGISTRY.lock().unwrap();
    reg.get(&id).cloned()
}

// ---------------------------------------------------------------------------
// Native Hardware Accelerated Tensor Math (Candle Core)
// ---------------------------------------------------------------------------

/// High-performance 2D matrix multiplication using Candle Tensor
pub fn candle_matmul(
    a_data: &[f32],
    m: usize,
    k1: usize,
    b_data: &[f32],
    k2: usize,
    n: usize,
) -> Result<Vec<f32>> {
    if k1 != k2 {
        bail!(
            "Dimension mismatch in matmul: [{}x{}] and [{}x{}]",
            m,
            k1,
            k2,
            n
        );
    }
    // Fast path: for small matrices, CPU SIMD avoids GPU kernel dispatch and compiler latency
    if m * k1 < 256 && k2 * n < 256 {
        let mut out = vec![0.0f32; m * n];
        for i in 0..m {
            for k in 0..k1 {
                let a_ik = a_data[i * k1 + k];
                for j in 0..n {
                    out[i * n + j] += a_ik * b_data[k * n + j];
                }
            }
        }
        return Ok(out);
    }

    let dev = get_device(None);
    let tensor_a = Tensor::from_slice(a_data, (m, k1), &dev)?;
    let tensor_b = Tensor::from_slice(b_data, (k2, n), &dev)?;
    let result = tensor_a.matmul(&tensor_b)?;
    let flat_result = result.flatten_all()?.to_vec1::<f32>()?;
    Ok(flat_result)
}

/// Numerically stable softmax along inner axis using Candle Tensor
pub fn candle_softmax(data: &[f32], shape: &[usize]) -> Result<Vec<f32>> {
    // Fast path: small vectors
    if data.len() < 256 {
        let max_val = data.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exps: Vec<f32> = data.iter().map(|&x| (x - max_val).exp()).collect();
        let sum: f32 = exps.iter().sum();
        if sum > 0.0 {
            return Ok(exps.into_iter().map(|x| x / sum).collect());
        }
    }

    let dev = get_device(None);
    let tensor = Tensor::from_slice(data, shape, &dev)?;
    let sm = candle_nn::ops::softmax_last_dim(&tensor)?;
    let flat = sm.flatten_all()?.to_vec1::<f32>()?;
    Ok(flat)
}

/// Fast dot product using Candle
pub fn candle_dot(a: &[f32], b: &[f32]) -> Result<f32> {
    if a.len() != b.len() {
        bail!(
            "Length mismatch for dot product: {} vs {}",
            a.len(),
            b.len()
        );
    }
    // Fast path: small vectors
    if a.len() < 256 {
        let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
        return Ok(dot);
    }

    let dev = get_device(None);
    let tensor_a = Tensor::from_slice(a, a.len(), &dev)?;
    let tensor_b = Tensor::from_slice(b, b.len(), &dev)?;
    let dot = (tensor_a * tensor_b)?.sum_all()?.to_scalar::<f32>()?;
    Ok(dot)
}

/// Fast L2 norm using Candle
pub fn candle_norm(a: &[f32]) -> Result<f32> {
    // Fast path: small vectors
    if a.len() < 256 {
        let sq: f32 = a.iter().map(|x| x * x).sum();
        return Ok(sq.sqrt());
    }

    let dev = get_device(None);
    let tensor = Tensor::from_slice(a, a.len(), &dev)?;
    let sq = tensor.sqr()?.sum_all()?.to_scalar::<f32>()?;
    Ok(sq.sqrt())
}