Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175//! One pet line or diary page: what to ask the brain, how lookups continue//! the same line, and how the result is filtered. The runtimes only generate.use crate::{ types::{DiaryPage, WorldState}, world,};use serde::{Deserialize, Serialize};
/// A request to the brain: continue `header`, the newest `lines` that fit and/// `prefix`. Prompts end at `Name:` with no trailing space. Sampling keeps the/// `top_k` likeliest tokens, then the smallest set holding `top_p` of their/// probability; temperature 0 is greedy (used only by checks).#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]#[serde(rename_all = "camelCase")]pub struct Generation { pub header: String, pub lines: Vec<String>, pub prefix: String, pub mode: String, pub max_new: usize, pub temperature: f32, pub top_k: usize, pub top_p: f32, pub ctx: usize,}pub const TEMPERATURE: f32 = 0.85;pub const TOP_K: usize = 40;pub const TOP_P: f32 = 0.92;/// Lookups per line before the line must end.const MAX_LOOKUPS: u32 = 3;/// Filtered lines are generated again this many times in all.const TRIES: u32 = 3;
/// The pet's line, with up to three lookup round-trips. Idle turns see a/// shorter history: prefill is most of their cost.#[derive(Clone, Debug, Deserialize, Serialize)]pub struct ChatTurn { pub kind: String, world: WorldState, header: String, lines: Vec<String>, text: String, hop: u32, tries: u32, body: Option<String>,}impl ChatTurn { pub fn new(world: WorldState, lines: Vec<String>, kind: &str) -> Self { Self { kind: if kind == "idle" { "idle" } else { "react" }.into(), header: world::header("chat", &world), world, lines, text: String::new(), hop: 0, tries: 0, body: None, } } pub fn request(&self) -> Generation { let idle = self.kind == "idle"; Generation { header: self.header.clone(), lines: self.lines.clone(), prefix: format!("{}:{}", self.world.pet, self.text), mode: "line".into(), max_new: if idle { 40 } else { 64 }, temperature: TEMPERATURE, top_k: TOP_K, top_p: TOP_P, ctx: if idle { 512 } else { 1024 }, } } /// Adds a generated piece that stopped for `reason`. Returns true when the /// line needs another generation: after a lookup (its result is appended /// here) or after a filtered line. pub fn feed(&mut self, text: &str, reason: &str) -> bool { if self.body.is_some() { return false; } self.text += text; if reason == "lookup" && self.hop < MAX_LOOKUPS { if let Some(result) = world::lookup_tail(&self.text) .and_then(|(name, arg)| world::tool_result(&self.world, &name, &arg)) { self.text += &format!("{{{result}}}"); self.hop += 1; return true; } } if world::not_cute(&self.text) { self.tries += 1; self.text.clear(); self.hop = 0; if self.tries < TRIES { return true; } self.body = Some(String::new()); return false; } self.body = Some(world::compact(&self.text)); false } /// The line so far, with lookup results, for live display. pub fn text(&self) -> &str { &self.text } /// The finished line, empty when every try was filtered. pub fn body(&self) -> Option<&str> { self.body.as_deref() }}
/// A diary page for one day, written from that day's log.#[derive(Clone, Debug, Deserialize, Serialize)]pub struct DiaryTurn { pub date: String, pub day: i64, pub weather: String, pub stage: String, /// Debug rewrites replace today's page. pub replace: bool, header: String,}pub const DIARY_START: &str = "dear diary,";impl DiaryTurn { pub fn new(world: &WorldState, date: String, replace: bool) -> Self { Self { date, day: world.day, weather: world.weather.clone(), stage: world.stage.clone(), replace, header: world::header("diary", world), } } pub fn request(&self) -> Generation { Generation { header: self.header.clone(), lines: vec![], prefix: DIARY_START.into(), mode: "diary".into(), max_new: 220, temperature: TEMPERATURE, top_k: TOP_K, top_p: TOP_P, ctx: 1024, } } /// The page from the brain's continuation: up to three known doodle /// stickers, text without tags. A page with unsafe text is dropped whole /// rather than publishing its fragments; the day stays pending. pub fn finish(&self, generated: &str) -> Option<DiaryPage> { let text = format!("{DIARY_START}{generated}"); if world::not_cute(&text) { return None; } let doodles = world::tags(&text) .into_iter() .filter(|t| t.name == "doodle") .filter_map(|t| t.arg) .filter(|d| world::DOODLES.contains(&d.as_str())) .take(3) .collect(); Some(DiaryPage { date: self.date.clone(), day: self.day, weather: self.weather.clone(), stage: self.stage.clone(), text: world::clean_page(&text), doodles, }) }}