Bump version to 1.6.3
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ryan Ward
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
4ba9706e35
commit
59cfd0e2f5
@@ -0,0 +1,115 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::Tensor;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct GenerationConfig {
|
||||
pub max_tokens: usize,
|
||||
pub temperature: f64,
|
||||
pub top_p: f64,
|
||||
}
|
||||
|
||||
impl Default for GenerationConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_tokens: 32,
|
||||
temperature: 0.7,
|
||||
top_p: 0.9,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GenerationConfig {
|
||||
pub fn deterministic() -> Self {
|
||||
Self {
|
||||
max_tokens: 32,
|
||||
temperature: 0.0,
|
||||
top_p: 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn creative() -> Self {
|
||||
Self {
|
||||
max_tokens: 64,
|
||||
temperature: 0.9,
|
||||
top_p: 0.95,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sample(logits: &Tensor, config: &GenerationConfig) -> Result<u32> {
|
||||
if config.temperature == 0.0 {
|
||||
return greedy(logits);
|
||||
}
|
||||
|
||||
let logits = logits.to_dtype(candle_core::DType::F32)?;
|
||||
let logits_vec = logits.to_vec1::<f32>()?;
|
||||
|
||||
// Apply temperature
|
||||
let scaled: Vec<f64> = logits_vec
|
||||
.iter()
|
||||
.map(|&x| (x as f64) / config.temperature)
|
||||
.collect();
|
||||
|
||||
// Softmax
|
||||
let max_val = scaled.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
|
||||
let exp: Vec<f64> = scaled.iter().map(|&x| (x - max_val).exp()).collect();
|
||||
let sum: f64 = exp.iter().sum();
|
||||
let probs: Vec<f64> = exp.iter().map(|&x| x / sum).collect();
|
||||
|
||||
// Top-p (nucleus) sampling
|
||||
let mut indexed_probs: Vec<(usize, f64)> = probs.iter().copied().enumerate().collect();
|
||||
indexed_probs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
let mut cumulative = 0.0;
|
||||
let mut candidates = Vec::new();
|
||||
for (idx, prob) in &indexed_probs {
|
||||
cumulative += prob;
|
||||
candidates.push((*idx, *prob));
|
||||
if cumulative >= config.top_p {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Renormalize candidates
|
||||
let candidate_sum: f64 = candidates.iter().map(|(_, p)| p).sum();
|
||||
let threshold = simple_rng() * candidate_sum;
|
||||
|
||||
let mut acc = 0.0;
|
||||
for (idx, prob) in &candidates {
|
||||
acc += prob;
|
||||
if acc >= threshold {
|
||||
return Ok(*idx as u32);
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to top candidate
|
||||
Ok(candidates[0].0 as u32)
|
||||
}
|
||||
|
||||
fn greedy(logits: &Tensor) -> Result<u32> {
|
||||
let logits = logits.to_dtype(candle_core::DType::F32)?;
|
||||
let logits_vec = logits.to_vec1::<f32>()?;
|
||||
let max_idx = logits_vec
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.map(|(idx, _)| idx)
|
||||
.unwrap_or(0);
|
||||
Ok(max_idx as u32)
|
||||
}
|
||||
|
||||
fn simple_rng() -> f64 {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::time::SystemTime;
|
||||
|
||||
let mut hasher = DefaultHasher::new();
|
||||
SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_nanos()
|
||||
.hash(&mut hasher);
|
||||
std::thread::current().id().hash(&mut hasher);
|
||||
let hash = hasher.finish();
|
||||
(hash as f64) / (u64::MAX as f64)
|
||||
}
|
||||
@@ -0,0 +1,309 @@
|
||||
mod generation;
|
||||
mod model_loader;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context as _, Result};
|
||||
use candle_core::{DType, Tensor};
|
||||
use candle_transformers::models::llama::{Cache, Config, Llama, LlamaConfig};
|
||||
use tokenizers::Tokenizer;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
pub use generation::GenerationConfig;
|
||||
|
||||
/// A token that can be used to cancel an in-progress generation.
|
||||
/// Clone it and pass to `generate_cancellable`, then call `cancel()` to
|
||||
/// interrupt the generation loop between tokens.
|
||||
#[derive(Clone)]
|
||||
pub struct CancellationToken {
|
||||
cancelled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl CancellationToken {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cancelled: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancel(&self) {
|
||||
self.cancelled.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.cancelled.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
const HF_REPO: &str = "HuggingFaceTB/SmolLM2-135M-Instruct";
|
||||
const MODEL_FILENAME: &str = "model.safetensors";
|
||||
const TOKENIZER_FILENAME: &str = "tokenizer.json";
|
||||
const CONFIG_FILENAME: &str = "config.json";
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum Device {
|
||||
Cpu,
|
||||
#[cfg(feature = "metal")]
|
||||
Metal,
|
||||
#[cfg(feature = "cuda")]
|
||||
Cuda(usize),
|
||||
}
|
||||
|
||||
impl Device {
|
||||
pub fn best_available() -> Self {
|
||||
#[cfg(feature = "metal")]
|
||||
{
|
||||
return Device::Metal;
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
return Device::Cuda(0);
|
||||
}
|
||||
#[cfg(not(any(feature = "metal", feature = "cuda")))]
|
||||
{
|
||||
Device::Cpu
|
||||
}
|
||||
}
|
||||
|
||||
fn to_candle_device(&self) -> Result<candle_core::Device> {
|
||||
match self {
|
||||
Device::Cpu => Ok(candle_core::Device::Cpu),
|
||||
#[cfg(feature = "metal")]
|
||||
Device::Metal => Ok(candle_core::Device::new_metal(0)?),
|
||||
#[cfg(feature = "cuda")]
|
||||
Device::Cuda(ordinal) => Ok(candle_core::Device::new_cuda(*ordinal)?),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct InferenceEngine {
|
||||
model: Arc<Mutex<Llama>>,
|
||||
cache: Arc<Mutex<Cache>>,
|
||||
config: Config,
|
||||
tokenizer: Tokenizer,
|
||||
device: candle_core::Device,
|
||||
}
|
||||
|
||||
impl InferenceEngine {
|
||||
pub async fn new(device: Device) -> Result<Self> {
|
||||
let candle_device = device.to_candle_device()?;
|
||||
let model_dir = model_loader::ensure_model_available().await?;
|
||||
|
||||
let config_path = model_dir.join(CONFIG_FILENAME);
|
||||
let config_str = tokio::fs::read_to_string(&config_path)
|
||||
.await
|
||||
.with_context(|| format!("reading config from {}", config_path.display()))?;
|
||||
let llama_config: LlamaConfig =
|
||||
serde_json::from_str(&config_str).context("parsing model config")?;
|
||||
let config = llama_config.into_config(false);
|
||||
|
||||
let model_path = model_dir.join(MODEL_FILENAME);
|
||||
let vb = unsafe {
|
||||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||||
&[model_path],
|
||||
DType::F32,
|
||||
&candle_device,
|
||||
)?
|
||||
};
|
||||
|
||||
let model = Llama::load(vb, &config).context("loading SmolLM2 model")?;
|
||||
let cache = Cache::new(true, DType::F32, &config, &candle_device)?;
|
||||
|
||||
let tokenizer_path = model_dir.join(TOKENIZER_FILENAME);
|
||||
let tokenizer_bytes = tokio::fs::read(&tokenizer_path)
|
||||
.await
|
||||
.with_context(|| format!("reading tokenizer from {}", tokenizer_path.display()))?;
|
||||
let tokenizer =
|
||||
Tokenizer::from_bytes(&tokenizer_bytes).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
log::info!(
|
||||
"Local inference engine initialized (device: {:?}, model: {})",
|
||||
device,
|
||||
HF_REPO
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
model: Arc::new(Mutex::new(model)),
|
||||
cache: Arc::new(Mutex::new(cache)),
|
||||
config,
|
||||
tokenizer,
|
||||
device: candle_device,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn tokenizer(&self) -> &Tokenizer {
|
||||
&self.tokenizer
|
||||
}
|
||||
|
||||
pub fn device(&self) -> &candle_core::Device {
|
||||
&self.device
|
||||
}
|
||||
|
||||
pub async fn generate_cancellable(
|
||||
&self,
|
||||
prompt: &str,
|
||||
config: &GenerationConfig,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<String> {
|
||||
let encoding = self
|
||||
.tokenizer
|
||||
.encode(prompt, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
let input_ids = encoding.get_ids().to_vec();
|
||||
let mut tokens = input_ids.clone();
|
||||
|
||||
let eos_token_id = self
|
||||
.tokenizer
|
||||
.token_to_id("</s>")
|
||||
.or_else(|| self.tokenizer.token_to_id("<|endoftext|>"))
|
||||
.or_else(|| self.tokenizer.token_to_id("<|im_end|>"));
|
||||
|
||||
let model = self.model.lock().await;
|
||||
let mut cache = self.cache.lock().await;
|
||||
|
||||
if cancel.is_cancelled() {
|
||||
anyhow::bail!("cancelled before generation started");
|
||||
}
|
||||
|
||||
// Reset cache for new generation
|
||||
*cache = Cache::new(true, DType::F32, &self.config, &self.device)?;
|
||||
|
||||
// Prefill: process all input tokens at once
|
||||
let input_tensor = Tensor::new(input_ids.as_slice(), &self.device)?.unsqueeze(0)?;
|
||||
let logits = model.forward(&input_tensor, 0, &mut cache)?;
|
||||
|
||||
if cancel.is_cancelled() {
|
||||
anyhow::bail!("cancelled during prefill");
|
||||
}
|
||||
|
||||
// Sample first generated token from the last logits position
|
||||
let first_logits = logits.squeeze(0)?;
|
||||
let next_token = generation::sample(&first_logits, config)?;
|
||||
|
||||
if let Some(eos) = eos_token_id {
|
||||
if next_token == eos {
|
||||
let generated_tokens = &tokens[input_ids.len()..];
|
||||
let output = self
|
||||
.tokenizer
|
||||
.decode(generated_tokens, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
return Ok(output.trim().to_string());
|
||||
}
|
||||
}
|
||||
tokens.push(next_token);
|
||||
|
||||
// Decode: generate one token at a time, checking cancellation between tokens
|
||||
for _i in 1..config.max_tokens {
|
||||
if cancel.is_cancelled() {
|
||||
anyhow::bail!("cancelled during generation");
|
||||
}
|
||||
|
||||
let pos = tokens.len() - 1;
|
||||
let next_input = Tensor::new(&[*tokens.last().unwrap()], &self.device)?.unsqueeze(0)?;
|
||||
let logits = model.forward(&next_input, pos, &mut cache)?;
|
||||
|
||||
let next_logits = logits.squeeze(0)?;
|
||||
let next_token = generation::sample(&next_logits, config)?;
|
||||
|
||||
if let Some(eos) = eos_token_id {
|
||||
if next_token == eos {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tokens.push(next_token);
|
||||
}
|
||||
|
||||
let generated_tokens = &tokens[input_ids.len()..];
|
||||
let output = self
|
||||
.tokenizer
|
||||
.decode(generated_tokens, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
Ok(output.trim().to_string())
|
||||
}
|
||||
|
||||
pub async fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
|
||||
let encoding = self
|
||||
.tokenizer
|
||||
.encode(prompt, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
let input_ids = encoding.get_ids().to_vec();
|
||||
let mut tokens = input_ids.clone();
|
||||
|
||||
let eos_token_id = self
|
||||
.tokenizer
|
||||
.token_to_id("</s>")
|
||||
.or_else(|| self.tokenizer.token_to_id("<|endoftext|>"))
|
||||
.or_else(|| self.tokenizer.token_to_id("<|im_end|>"));
|
||||
|
||||
let model = self.model.lock().await;
|
||||
let mut cache = self.cache.lock().await;
|
||||
|
||||
// Reset cache for new generation
|
||||
*cache = Cache::new(true, DType::F32, &self.config, &self.device)?;
|
||||
|
||||
// Prefill: process all input tokens at once
|
||||
let input_tensor = Tensor::new(input_ids.as_slice(), &self.device)?.unsqueeze(0)?;
|
||||
let logits = model.forward(&input_tensor, 0, &mut cache)?;
|
||||
|
||||
// Sample first generated token from the last logits position
|
||||
let first_logits = logits.squeeze(0)?;
|
||||
let next_token = generation::sample(&first_logits, config)?;
|
||||
|
||||
if let Some(eos) = eos_token_id {
|
||||
if next_token == eos {
|
||||
let generated_tokens = &tokens[input_ids.len()..];
|
||||
let output = self
|
||||
.tokenizer
|
||||
.decode(generated_tokens, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
return Ok(output.trim().to_string());
|
||||
}
|
||||
}
|
||||
tokens.push(next_token);
|
||||
|
||||
// Decode: generate one token at a time
|
||||
for _i in 1..config.max_tokens {
|
||||
let pos = tokens.len() - 1;
|
||||
let next_input = Tensor::new(&[*tokens.last().unwrap()], &self.device)?.unsqueeze(0)?;
|
||||
let logits = model.forward(&next_input, pos, &mut cache)?;
|
||||
|
||||
let next_logits = logits.squeeze(0)?;
|
||||
let next_token = generation::sample(&next_logits, config)?;
|
||||
|
||||
if let Some(eos) = eos_token_id {
|
||||
if next_token == eos {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tokens.push(next_token);
|
||||
}
|
||||
|
||||
let generated_tokens = &tokens[input_ids.len()..];
|
||||
let output = self
|
||||
.tokenizer
|
||||
.decode(generated_tokens, true)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
Ok(output.trim().to_string())
|
||||
}
|
||||
|
||||
pub fn model_dir() -> PathBuf {
|
||||
model_loader::model_cache_dir()
|
||||
}
|
||||
|
||||
pub async fn is_model_downloaded() -> bool {
|
||||
let dir = Self::model_dir();
|
||||
tokio::fs::metadata(dir.join(MODEL_FILENAME)).await.is_ok()
|
||||
&& tokio::fs::metadata(dir.join(TOKENIZER_FILENAME))
|
||||
.await
|
||||
.is_ok()
|
||||
&& tokio::fs::metadata(dir.join(CONFIG_FILENAME)).await.is_ok()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
|
||||
use super::{CONFIG_FILENAME, HF_REPO, MODEL_FILENAME, TOKENIZER_FILENAME};
|
||||
|
||||
pub fn model_cache_dir() -> PathBuf {
|
||||
directories::ProjectDirs::from("", "", "galaxy")
|
||||
.map(|dirs| dirs.cache_dir().join("models").join("smollm2-135m"))
|
||||
.unwrap_or_else(|| PathBuf::from(".galaxy/models/smollm2-135m"))
|
||||
}
|
||||
|
||||
pub async fn ensure_model_available() -> Result<PathBuf> {
|
||||
let cache_dir = model_cache_dir();
|
||||
|
||||
let model_path = cache_dir.join(MODEL_FILENAME);
|
||||
let tokenizer_path = cache_dir.join(TOKENIZER_FILENAME);
|
||||
let config_path = cache_dir.join(CONFIG_FILENAME);
|
||||
|
||||
if model_path.exists() && tokenizer_path.exists() && config_path.exists() {
|
||||
log::debug!("Model already cached at {}", cache_dir.display());
|
||||
return Ok(cache_dir);
|
||||
}
|
||||
|
||||
log::info!("Downloading SmolLM2-135M model from HuggingFace...");
|
||||
tokio::fs::create_dir_all(&cache_dir)
|
||||
.await
|
||||
.with_context(|| format!("creating cache dir {}", cache_dir.display()))?;
|
||||
|
||||
let api = hf_hub::api::tokio::Api::new().context("initializing HuggingFace API")?;
|
||||
let repo = api.model(HF_REPO.to_string());
|
||||
|
||||
let downloaded_model = repo
|
||||
.get(MODEL_FILENAME)
|
||||
.await
|
||||
.context("downloading model weights")?;
|
||||
let downloaded_tokenizer = repo
|
||||
.get(TOKENIZER_FILENAME)
|
||||
.await
|
||||
.context("downloading tokenizer")?;
|
||||
let downloaded_config = repo
|
||||
.get(CONFIG_FILENAME)
|
||||
.await
|
||||
.context("downloading config")?;
|
||||
|
||||
// hf-hub caches files itself, but we symlink/copy to our canonical location
|
||||
// for predictable access
|
||||
link_or_copy(&downloaded_model, &model_path).await?;
|
||||
link_or_copy(&downloaded_tokenizer, &tokenizer_path).await?;
|
||||
link_or_copy(&downloaded_config, &config_path).await?;
|
||||
|
||||
log::info!("Model downloaded and cached at {}", cache_dir.display());
|
||||
Ok(cache_dir)
|
||||
}
|
||||
|
||||
async fn link_or_copy(src: &std::path::Path, dst: &std::path::Path) -> Result<()> {
|
||||
if dst.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Try symlink first (saves disk space)
|
||||
#[cfg(unix)]
|
||||
{
|
||||
if tokio::fs::symlink(src, dst).await.is_ok() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to copy
|
||||
tokio::fs::copy(src, dst)
|
||||
.await
|
||||
.with_context(|| format!("copying {} to {}", src.display(), dst.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user