139 lines
4.0 KiB
Rust
139 lines
4.0 KiB
Rust
//! Candle host pipeline for CALCULET-compiled language models.
|
|
//!
|
|
//! This replaces the tokenizer, Qwen chat-template, sampling, generation loop,
|
|
//! and CALRT tensor glue from the captured llama.cpp fork. The compiled
|
|
//! prefill/decode graphs remain Calbin models executed by `calculet-calrt`.
|
|
|
|
mod backend;
|
|
mod calrt;
|
|
mod engine;
|
|
mod error;
|
|
mod template;
|
|
mod tokenizer;
|
|
|
|
pub use backend::CandleChatBackend;
|
|
pub use calrt::{CalrtCompiledModel, CalrtModelPlan};
|
|
pub use engine::{
|
|
CandleEngine, CompiledModel, DEFAULT_SEQUENCE_ID, GenerationConfig, ModelInput, ModelPhase,
|
|
ModelStatus,
|
|
};
|
|
pub use error::{Error, Result};
|
|
pub use template::{
|
|
ParsedAssistantOutput, QWEN_END_OF_TEXT, QWEN_IM_END, parse_qwen3_output, render_qwen3,
|
|
};
|
|
pub use tokenizer::{HuggingFaceTokenizer, TokenCodec};
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use agentos_protocol::{
|
|
CHAT_PROTOCOL_VERSION, ChatMessage, ChatRequest, MessageRole, StopReason, ToolChoice,
|
|
};
|
|
use candle_core::{Device, Tensor};
|
|
use serde_json::Map;
|
|
use std::collections::BTreeMap;
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Debug)]
|
|
struct TestTokenizer {
|
|
pieces: BTreeMap<u32, &'static str>,
|
|
}
|
|
|
|
impl TestTokenizer {
|
|
fn new() -> Self {
|
|
Self {
|
|
pieces: BTreeMap::from([(0, "prompt"), (1, "done"), (2, QWEN_IM_END)]),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TokenCodec for TestTokenizer {
|
|
fn encode(&self, _text: &str) -> Result<Vec<u32>> {
|
|
Ok(vec![0])
|
|
}
|
|
|
|
fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
|
|
Ok(token_ids
|
|
.iter()
|
|
.filter_map(|token| self.pieces.get(token))
|
|
.copied()
|
|
.collect())
|
|
}
|
|
|
|
fn token_to_id(&self, token: &str) -> Option<u32> {
|
|
self.pieces
|
|
.iter()
|
|
.find_map(|(id, piece)| (*piece == token).then_some(*id))
|
|
}
|
|
|
|
fn vocabulary_size(&self) -> usize {
|
|
self.pieces.len()
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct TestModel {
|
|
step: usize,
|
|
}
|
|
|
|
impl CompiledModel for TestModel {
|
|
fn status(&self) -> ModelStatus {
|
|
ModelStatus {
|
|
ready: true,
|
|
reason: "test model is ready".into(),
|
|
details: serde_json::Value::Null,
|
|
}
|
|
}
|
|
|
|
fn vocabulary_size(&self) -> usize {
|
|
4
|
|
}
|
|
|
|
fn maximum_sequence_length(&self) -> usize {
|
|
32
|
|
}
|
|
|
|
fn reset_sequence(&mut self, _sequence_id: u32) -> Result<()> {
|
|
self.step = 0;
|
|
Ok(())
|
|
}
|
|
|
|
fn forward(&mut self, _input: &ModelInput) -> Result<Tensor> {
|
|
let token = if self.step == 0 { 1 } else { 2 };
|
|
self.step += 1;
|
|
let mut logits = vec![-100.0_f32; self.vocabulary_size()];
|
|
logits[token] = 100.0;
|
|
logits[3] = 200.0;
|
|
Ok(Tensor::from_vec(
|
|
logits,
|
|
self.vocabulary_size(),
|
|
&Device::Cpu,
|
|
)?)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn runs_a_deterministic_candle_generation_loop_without_hardware() {
|
|
let mut engine = CandleEngine::new(
|
|
Arc::new(TestTokenizer::new()),
|
|
TestModel::default(),
|
|
GenerationConfig::default(),
|
|
)
|
|
.unwrap();
|
|
let response = engine
|
|
.generate(&ChatRequest {
|
|
protocol_version: CHAT_PROTOCOL_VERSION.into(),
|
|
messages: vec![ChatMessage::text(MessageRole::User, "run")],
|
|
tools: Vec::new(),
|
|
max_output_tokens: 4,
|
|
tool_choice: ToolChoice::Auto,
|
|
parallel_tool_calls: false,
|
|
metadata: Map::new(),
|
|
})
|
|
.unwrap();
|
|
assert_eq!(response.content, "done");
|
|
assert_eq!(response.stop_reason, StopReason::EndTurn);
|
|
assert_eq!(response.usage.output_tokens, 2);
|
|
}
|
|
}
|