Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317//! One mutable ONNX session per worker. No concurrent calls into a session.use anyhow::{Context, Result, ensure};use ort::{ memory::{AllocationDevice, AllocatorType, MemoryInfo, MemoryType}, session::Session, value::{DynValue, Tensor},};use serde::{Deserialize, Serialize};use std::{path::Path, time::Instant};use tokenizers::Tokenizer;/// Requests, the stop rule and context fitting are shared with the browser's worker.pub use tokenling_core::turn::{Generation, TOP_K, TOP_P};use tokenling_core::world;
#[derive(Clone, Debug, Deserialize, Serialize)]#[serde(rename_all = "camelCase")]pub struct Generated { pub text: String, pub prompt: String, pub reason: String, pub prompt_tokens: usize, pub prefill_ms: f64, pub tps: f64, pub generated_tokens: usize, pub total_ms: f64,}pub struct Brain { session: Session, tokenizer: Tokenizer, past_names: Vec<String>, inputs: Vec<String>, heads: usize, width: usize, gpu: Option<i32>,}impl Brain { pub fn load(model_dir: &Path, threads: usize, gpu: Option<i32>) -> Result<Self> { let tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json")) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let cfg: serde_json::Value = serde_json::from_slice(&std::fs::read(model_dir.join("config.json"))?)?; let heads = cfg["n_head"].as_u64().context("config requires n_head")? as usize; let width = cfg["n_embd"].as_u64().context("config requires n_embd")? as usize / heads; let mut builder = Session::builder()?; builder = builder .with_intra_threads(threads) .map_err(|e| anyhow::anyhow!(e.to_string()))?; builder = builder .with_inter_threads(1) .map_err(|e| anyhow::anyhow!(e.to_string()))?; for (key, value) in [ ("session.disable_quant_qdq", "1"), ("session.disable_prepacking", "1"), ("session.intra_op.allow_spinning", "0"), ] { builder = builder .with_config_entry(key, value) .map_err(|e| anyhow::anyhow!(e.to_string()))?; } if let Some(device) = gpu { builder = builder .with_execution_providers([ort::ep::CUDA::default() .with_device_id(device) .with_tf32(false) .build() .error_on_failure()]) .map_err(|e| anyhow::anyhow!(e.to_string()))?; } let session = builder .commit_from_file(model_dir.join("onnx/model_q4.onnx")) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let inputs: Vec<_> = session .inputs() .iter() .map(|i| i.name().to_owned()) .collect(); let past_names = inputs .iter() .filter(|n| n.starts_with("past_key_values.")) .cloned() .collect(); let mut brain = Self { session, tokenizer, past_names, inputs, heads, width, gpu, }; brain.generate(&Generation { header: String::new(), lines: vec![], prefix: "hi".into(), mode: "line".into(), max_new: 1, temperature: 0.0, top_k: TOP_K, top_p: TOP_P, ctx: 1024, })?; Ok(brain) } pub fn count(&self, text: &str) -> Result<usize> { Ok(self .tokenizer .encode(text, false) .map_err(|e| anyhow::anyhow!(e.to_string()))? .len()) } pub fn fit(&self, g: &Generation) -> Result<String> { world::fit( &g.header, &g.lines, &g.prefix, g.max_new, g.ctx, &mut |text| self.count(text).map_err(|e| e.to_string()), ) .map_err(anyhow::Error::msg) } pub fn generate(&mut self, g: &Generation) -> Result<Generated> { self.generate_cancellable(g, || false) } pub fn generate_cancellable( &mut self, g: &Generation, cancelled: impl Fn() -> bool, ) -> Result<Generated> { ensure!((1..=220).contains(&g.max_new), "invalid output budget"); ensure!( matches!(g.mode.as_str(), "line" | "diary"), "invalid generation mode" ); ensure!( g.temperature.is_finite() && (0.0..=2.0).contains(&g.temperature), "invalid temperature" ); ensure!( (1..=50257).contains(&g.top_k) && g.top_p > 0.0 && g.top_p <= 1.0, "invalid sampling settings" ); let prompt = self.fit(g)?; let encoded = self .tokenizer .encode(prompt.clone(), false) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let ids: Vec<i64> = encoded.get_ids().iter().map(|i| *i as i64).collect(); let prompt_tokens = ids.len(); let mut current = ids; let mut generated = Vec::new(); let mut past: Vec<(String, DynValue)> = self .past_names .iter() .map(|name| { Ok(( name.clone(), Tensor::<f32>::from_array(([1, self.heads, 0, self.width], Vec::<f32>::new()))? .into_dyn(), )) }) .collect::<Result<_>>()?; let started = Instant::now(); let mut prefill_ms = 0.0; let mut reason = "length"; let mut text = String::new(); let mut binding = self .gpu .map(|_| self.session.create_binding()) .transpose()?; let gpu_memory = self .gpu .map(|device| { MemoryInfo::new( AllocationDevice::CUDA, device, AllocatorType::Device, MemoryType::Default, ) }) .transpose()?; let cpu_memory = MemoryInfo::new( AllocationDevice::CPU, 0, AllocatorType::Device, MemoryType::CPUOutput, )?; for step in 0..g.max_new { ensure!(!cancelled(), "generation cancelled"); let total = prompt_tokens + step; let positions: Vec<i64> = if step == 0 { (0..prompt_tokens as i64).collect() } else { vec![(total - 1) as i64] }; let mut input = vec![ ( "input_ids".into(), Tensor::<i64>::from_array(([1, current.len()], current))?.into_dyn(), ), ( "attention_mask".into(), Tensor::<i64>::from_array(([1, total], vec![1; total]))?.into_dyn(), ), ( "position_ids".into(), Tensor::<i64>::from_array(([1, positions.len()], positions))?.into_dyn(), ), ]; input.retain(|(name, _)| self.inputs.contains(name)); input.append(&mut past); let mut output = if let Some(binding) = &mut binding { binding.clear(); for (name, value) in &input { binding.bind_input(name, value)?; } // Sampling needs only logits on the CPU. Keep all KV outputs // on CUDA and feed their Values directly into the next step. // Fresh output bindings avoid aliasing a previous cache buffer. binding.bind_output_to_device("logits", &cpu_memory)?; for name in &self.past_names { binding.bind_output_to_device( name.replace("past_key_values.", "present."), gpu_memory.as_ref().unwrap(), )?; } let output = self.session.run_binding(binding)?; binding.synchronize_outputs()?; output } else { self.session.run(input)? }; let (_, logits) = output["logits"].try_extract_tensor::<f32>()?; let token = sample(&logits[logits.len() - 50257..], g)?; if step == 0 { prefill_ms = started.elapsed().as_secs_f64() * 1000.0; } if token == 50256 { reason = "eos"; break; } generated.push(token); text = self .tokenizer .decode(&generated, true) .map_err(|e| anyhow::anyhow!(e.to_string()))?; if let Some((why, kept)) = world::stop(&g.mode, &text) { text = kept; reason = why; break; } for name in &self.past_names { let value = output .remove(name.replace("past_key_values.", "present.")) .context("missing KV output")?; past.push((name.clone(), value)); } current = vec![token as i64]; } let decode_s = (started.elapsed().as_secs_f64() - prefill_ms / 1000.0).max(1e-6); Ok(Generated { text, prompt, reason: reason.into(), prompt_tokens, prefill_ms, tps: generated.len().saturating_sub(1) as f64 / decode_s, generated_tokens: generated.len(), total_ms: started.elapsed().as_secs_f64() * 1000.0, }) }}
fn sample(logits: &[f32], g: &Generation) -> Result<u32> { let temperature = g.temperature; ensure!( logits.iter().all(|v| v.is_finite()), "non-finite model logits" ); if temperature == 0.0 { return Ok(logits .iter() .enumerate() .max_by(|a, b| a.1.total_cmp(b.1)) .unwrap() .0 as u32); } let mut ranked: Vec<_> = logits.iter().copied().enumerate().collect(); let k = g.top_k.clamp(1, ranked.len()); ranked.select_nth_unstable_by(k - 1, |a, b| b.1.total_cmp(&a.1)); ranked.truncate(k); ranked.sort_unstable_by(|a, b| b.1.total_cmp(&a.1)); let high = ranked[0].1; let probabilities: Vec<_> = ranked .iter() .map(|(_, v)| ((*v - high) / temperature).exp()) .collect(); let total: f32 = probabilities.iter().sum(); let mut kept = 0.0; let mut n = 0; for p in &probabilities { kept += p; n += 1; if kept / total >= g.top_p { break; } } let mut draw = rand::random::<f32>() * kept; for (i, p) in probabilities.iter().take(n).enumerate() { draw -= p; if draw <= 0.0 { return Ok(ranked[i].0 as u32); } } Ok(ranked[n - 1].0 as u32)}