open-standards/joule-code/examples/l3-jouleclaw/src/hybrid_backend.rs

//! In-process L3 backend: the **hybrid safetensors families** behind
//! JouleClaw's `HybridLm` trait.
//!
//! Where [`crate::runtime_backend::JouleclawRuntimeBackend`] runs GGUF
//! transformer checkpoints on the kernel-graph runtime, this backend runs the
//! safetensors families — NemotronH (interleaved mixers), Granite-4.0-H
//! (dual-residual, dense or MoE), Falcon-H1 (parallel attention ∥ SSM),
//! Cohere2-MoE (North Mini Code: full/sliding attention + sigmoid MoE), and
//! Mellum2 (YaRN + QK-norm + softmax-renorm MoE) — straight from their BF16
//! safetensors checkpoints via
//! `jouleclaw_loader_gguf::hybrid_lm`. One enum, one `generate()` loop:
//! batched prefill (bit-identical to sequential decode by the loader's parity
//! guarantees), then greedy or seeded sampling.
//!
//! The family is auto-detected from `config.json`'s `model_type`, so a
//! consumer points at a snapshot directory and gets whichever hybrid lives
//! there. Same [`LlmBackend`] trait as every other backend — it drops into
//! [`crate::jouleclaw_l3_stage`] unchanged.

use jouleclaw_llm::{FinishReason, LlmBackend, LlmError, LlmRequest, LlmResponse};
use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::cohere2_moe::Cohere2MoeModel;
use jouleclaw_loader_gguf::falcon_h1::FalconH1Model;
use jouleclaw_loader_gguf::granite_h::GraniteHModel;
use jouleclaw_loader_gguf::hybrid_lm::{generate, GenerateParams};
use jouleclaw_loader_gguf::mellum::MellumModel;
use jouleclaw_loader_gguf::nemotron_h::NemotronHModel;
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use jouleclaw_loader_gguf::sample::SamplingConfig;
use std::path::{Path, PathBuf};

/// The loaded hybrid, by family.
enum Hybrid {
    NemotronH(NemotronHModel),
    GraniteH(GraniteHModel),
    FalconH1(FalconH1Model),
    Cohere2Moe(Cohere2MoeModel),
    Mellum(MellumModel),
}

/// An L3 backend over any of the hybrid safetensors families.
pub struct HybridBackend {
    model: Hybrid,
    tokenizer: BpeTokenizer,
    name: String,
    /// BOS prepended to every prompt (family convention).
    bos: Option<u32>,
    /// EOS ids that stop generation (from `config.json`).
    eos: Vec<u32>,
    typical_joules: f64,
}

impl HybridBackend {
    /// Load a hybrid checkpoint from a snapshot directory
    /// (`config.json` + `model*.safetensors` + `tokenizer.json`),
    /// auto-detecting the family from `model_type`.
    pub fn from_dir(dir: impl AsRef<Path>) -> Result<Self, LlmError> {
        let dir = dir.as_ref();
        let unavailable = |what: &str, e: String| LlmError::Unavailable(format!("{what}: {e}"));
        let cfg_bytes = std::fs::read(dir.join("config.json"))
            .map_err(|e| unavailable("read config.json", e.to_string()))?;
        let cfg_json: serde_json::Value = serde_json::from_slice(&cfg_bytes)
            .map_err(|e| unavailable("parse config.json", e.to_string()))?;
        let model_type = cfg_json["model_type"].as_str().unwrap_or_default().to_string();

        // Collect every model*.safetensors shard, sorted.
        let mut shard_paths: Vec<PathBuf> = std::fs::read_dir(dir)
            .map_err(|e| unavailable("read dir", e.to_string()))?
            .filter_map(|e| e.ok().map(|e| e.path()))
            .filter(|p| {
                p.extension().is_some_and(|e| e == "safetensors")
                    && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("model"))
            })
            .collect();
        shard_paths.sort();
        if shard_paths.is_empty() {
            return Err(unavailable("checkpoint", "no model*.safetensors found".into()));
        }
        let mut shards = shard_paths
            .iter()
            .map(|p| read_safetensors_file(p).map_err(|e| format!("{}: {e:?}", p.display())))
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| unavailable("map safetensors", e))?;

        let (model, bos) = match model_type.as_str() {
            "nemotron_h" => {
                let cfg = serde_json::from_slice(&cfg_bytes)
                    .map_err(|e| unavailable("nemotron_h config", e.to_string()))?;
                let st = shards.pop().unwrap();
                if !shards.is_empty() {
                    return Err(unavailable("nemotron_h", "expected single-file checkpoint".into()));
                }
                let m = NemotronHModel::bind(cfg, st)
                    .map_err(|e| unavailable("bind nemotron_h", format!("{e:?}")))?;
                (Hybrid::NemotronH(m), Some(1)) // BOS = 1, per the family bring-up
            }
            "granitemoehybrid" => {
                let cfg = serde_json::from_slice(&cfg_bytes)
                    .map_err(|e| unavailable("granite config", e.to_string()))?;
                let m = GraniteHModel::bind_sharded(cfg, shards)
                    .map_err(|e| unavailable("bind granite", format!("{e:?}")))?;
                (Hybrid::GraniteH(m), None)
            }
            "falcon_h1" => {
                let cfg = serde_json::from_slice(&cfg_bytes)
                    .map_err(|e| unavailable("falcon_h1 config", e.to_string()))?;
                let st = shards.pop().unwrap();
                if !shards.is_empty() {
                    return Err(unavailable("falcon_h1", "expected single-file checkpoint".into()));
                }
                let m = FalconH1Model::bind(cfg, st)
                    .map_err(|e| unavailable("bind falcon_h1", format!("{e:?}")))?;
                (Hybrid::FalconH1(m), None)
            }
            "cohere2_moe" => {
                let cfg = serde_json::from_slice(&cfg_bytes)
                    .map_err(|e| unavailable("cohere2_moe config", e.to_string()))?;
                let m = Cohere2MoeModel::bind_sharded(cfg, shards)
                    .map_err(|e| unavailable("bind cohere2_moe", format!("{e:?}")))?;
                // bos_token_id from the config (Cohere BOS).
                let bos = cfg_json["bos_token_id"].as_u64().map(|v| v as u32);
                (Hybrid::Cohere2Moe(m), bos)
            }
            "mellum" => {
                let cfg = serde_json::from_slice(&cfg_bytes)
                    .map_err(|e| unavailable("mellum config", e.to_string()))?;
                let m = MellumModel::bind_sharded(cfg, shards)
                    .map_err(|e| unavailable("bind mellum", format!("{e:?}")))?;
                let bos = cfg_json["bos_token_id"].as_u64().map(|v| v as u32);
                (Hybrid::Mellum(m), bos)
            }
            other => {
                return Err(unavailable(
                    "model_type",
                    format!("unsupported hybrid family {other:?}"),
                ));
            }
        };

        // EOS may be an int or a list in HF configs.
        let eos = match &cfg_json["eos_token_id"] {
            serde_json::Value::Number(n) => n.as_u64().map(|v| v as u32).into_iter().collect(),
            serde_json::Value::Array(a) => {
                a.iter().filter_map(|v| v.as_u64().map(|v| v as u32)).collect()
            }
            _ => Vec::new(),
        };
        let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json"))
            .map_err(|e| unavailable("tokenizer.json", format!("{e:?}")))?;
        let name = dir
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or("hybrid-model")
            .to_string();
        Ok(Self { model, tokenizer, name, bos, eos, typical_joules: 1.0 })
    }
}

impl LlmBackend for HybridBackend {
    fn model_name(&self) -> &str {
        &self.name
    }

    fn complete(&self, request: &LlmRequest) -> Result<LlmResponse, LlmError> {
        let mut prompt: Vec<u32> = Vec::new();
        if let Some(b) = self.bos {
            prompt.push(b);
        }
        prompt.extend(self.tokenizer.encode(&request.prompt));
        if prompt.is_empty() {
            return Err(LlmError::Other("empty prompt after tokenization".into()));
        }
        let sampling = if request.temperature <= 0.0 {
            SamplingConfig::greedy()
        } else {
            SamplingConfig::temperature(request.temperature, 0)
        };
        let params = GenerateParams::greedy(request.max_tokens as usize)
            .with_stop_tokens(&self.eos)
            .with_sampling(sampling);

        // Fabric attribution window: GFLOPs per member for
        // THIS call, each priced by the member's loaded matmul
        // figure — a projection, not a meter reading; reported
        // only when the fabric priced some of it. The L3 stage
        // grades it by its weakest input
        // (`crate::projection_grade`): ModelBased only when
        // every member was priced by a Measured or Benchmarked
        // figure.
        jouleclaw_loader_gguf::fabric_dispatch::ledger_reset();
        let out = match &self.model {
            Hybrid::NemotronH(m) => generate(m, &prompt, &params),
            Hybrid::GraniteH(m) => generate(m, &prompt, &params),
            Hybrid::FalconH1(m) => generate(m, &prompt, &params),
            Hybrid::Cohere2Moe(m) => generate(m, &prompt, &params),
            Hybrid::Mellum(m) => generate(m, &prompt, &params),
        };
        let ledger = jouleclaw_loader_gguf::fabric_dispatch::ledger_snapshot();
        let projected_j: f64 = ledger.iter().filter_map(|e| e.projected_uj).sum::<f64>() / 1e6;
        let energy_joules = (projected_j > 0.0).then_some(projected_j);
        let stopped = out.len() < request.max_tokens as usize;
        Ok(LlmResponse {
            text: self.tokenizer.decode(&out),
            finish_reason: if stopped { FinishReason::Stop } else { FinishReason::Length },
            input_tokens: prompt.len() as u32,
            output_tokens: out.len() as u32,
            // The fabric projection when the ledger priced the
            // call (graded by the stage from the ledger's
            // inputs); honest None otherwise (the run binary can
            // still meter host energy over the whole call).
            energy_joules,
        })
    }

    fn typical_joules_per_call(&self) -> f64 {
        self.typical_joules
    }
}