diff options
Diffstat (limited to 'src/learner.rs')
| -rw-r--r-- | src/learner.rs | 183 |
1 files changed, 108 insertions, 75 deletions
diff --git a/src/learner.rs b/src/learner.rs index 91b24fe..b921c3a 100644 --- a/src/learner.rs +++ b/src/learner.rs @@ -5,21 +5,33 @@ use serde::{Deserialize, Serialize}; use crate::{config::Config, trajectory::Action}; -fn key(a: &str, b: &str) -> (String, String) { - let a = a.to_string(); - let b = b.to_string(); - if a < b { (a, b) } else { (b, a) } +const DIM: usize = 16; + +// ----------------------------------------------------------------------------- +// helpers + +fn dot(a: &[f32; DIM], b: &[f32; DIM]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() } -fn compute_sum( - cand: &str, - known_keys: &HashSet<&String>, - map: &HashMap<(String, String), f32>, -) -> f32 { - known_keys - .iter() - .filter_map(|k| map.get(&key(k, cand))) - .sum() +fn norm(a: &[f32; DIM]) -> f32 { + dot(a, a).sqrt() +} + +fn renorm(a: &mut [f32; DIM]) { + let n = norm(a).max(1e-6); + for x in a.iter_mut() { + *x /= n; + } +} + +fn init_vec() -> [f32; DIM] { + let mut v = [0.0; DIM]; + for x in v.iter_mut() { + *x = rand::random_range(-0.1..0.1); + } + renorm(&mut v); + v } fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, f32)> { @@ -50,10 +62,29 @@ fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, #[derive(Serialize, Deserialize, Default)] pub struct Learner { - similar: HashMap<(String, String), f32>, - different: HashMap<(String, String), f32>, + embedding: HashMap<String, [f32; DIM]>, } +impl Learner { + pub fn prune(&mut self, valid: &HashSet<String>) { + self.embedding.retain(|k, _v| valid.contains(k)); + } + + fn nudge(&mut self, target: &str, direction: &str, rate: f32) { + let t = *self + .embedding + .entry(direction.into()) + .or_insert_with(init_vec); + let v = self.embedding.entry(target.into()).or_insert_with(init_vec); + for i in 0..DIM { + v[i] += rate * (t[i] - v[i]); + } + renorm(v); + } +} +// ----------------------------------------------------------------------------- +// core + #[derive(Debug)] pub enum Learning { SkipExtend, @@ -63,65 +94,64 @@ pub enum Learning { } impl Learner { - pub fn prune(&mut self, valid: &HashSet<String>) { - self.similar - .retain(|(a, b), _v| valid.contains(a) && valid.contains(b)); - self.different - .retain(|(a, b), _v| valid.contains(a) && valid.contains(b)); - } - pub fn learn(&mut self, config: &Config, learning: &Learning) { - // global decay - for (_k, v) in self.similar.iter_mut().chain(self.different.iter_mut()) { - *v *= 1.0 - config.decay_rate; - } - // update match learning { Learning::SkipExtend => { - // This is the least information carrying case, it does not - // follow that the trajectory contains similar or different - // items, we may simply be seeking something in particular. + /* This is the least informative case, it does not + follow that the trajectory contains similar or different + items, we may simply be seeking something in particular. */ } + Learning::MoreExtend(trajectory, new) => { - // We're continuing a good run so `new` is similar to everything - // in `trajectory`, but to account for the possibilty that our - // mood has changed over the course of this streak we damp - // sub-linearly that update by a proxy of temporal distance. - for (distance, e) in trajectory.iter().rev().enumerate() { + /* We're continuing a good run so `new` is similar to everything + in `trajectory`, but to account for the possibilty that our + mood has changed over the course of this streak we dampen + this feedback. */ + for (distance, similar) in trajectory.iter().rev().enumerate() { let damp = ((distance + 1) as f32).powf(-0.5); - let v = self.similar.entry(key(e, new)).or_default(); - *v = (1.0 - config.learning_rate) * *v + config.learning_rate * damp; + self.nudge(new, similar, config.learning_rate * damp); + self.nudge(similar, new, config.learning_rate * damp * 0.5); // optional: symmetric, weaker } } + Learning::SkipToMore(trajectory, new) => { - // We have learnt that `new` is different to everything in - // `trajectory`, the strongest signal we have. - for e in trajectory { - let v = self.different.entry(key(e, new)).or_default(); - *v = (1.0 - config.learning_rate) * *v + config.learning_rate; + /* We have learnt that `new` is different to everything in + `trajectory`, the strongest signal we have. */ + for different in trajectory { + self.nudge(new, different, -config.learning_rate); } } + Learning::MoreToSkip(trajectory, new) => { - // `new` could be different to everything in `trajectory`, or we - // simply changed our minds, so we have only weak evidence of - // difference. The positive coherence of trajectory was taken - // care of during MoreExtend above. + /* `new` could be different to everything in `trajectory`, or we + simply changed our minds, so we have only weak evidence of + difference. The positive coherence of trajectory was taken + care of during MoreExtend above. */ let damp = (trajectory.len() + 1) as f32; - for e in trajectory { - let v = self.different.entry(key(e, new)).or_default(); - *v = (1.0 - config.learning_rate) * *v + config.learning_rate / damp; + for weakly_different in trajectory { + self.nudge(new, weakly_different, -config.learning_rate / damp); } } - }; - // global threshold drop - self.similar = self - .similar - .extract_if(|_k, v| (*v).abs() > config.min_score_thresh) - .collect(); - self.different = self - .different - .extract_if(|_k, v| (*v).abs() > config.min_score_thresh) - .collect(); + } + } + + fn mood(&self, trajectory: &[String]) -> Option<[f32; DIM]> { + let mut m = [0.0; DIM]; + let mut any = false; + for (distance, historical) in trajectory.iter().rev().enumerate() { + if let Some(v) = self.embedding.get(historical) { + let damp = ((distance + 1) as f32).powf(-0.5); + for i in 0..DIM { + m[i] += damp * v[i]; + } + any = true; + } + } + if !any || norm(&m) < 1e-6 { + return None; + } + renorm(&mut m); + Some(m) } pub fn sample( @@ -131,30 +161,33 @@ impl Learner { candidates: &HashSet<String>, temperature: f32, ) -> Option<(String, f32)> { - let trajectory: HashSet<_> = trajectory.iter().collect(); - let normalisation: f32 = f32::max(trajectory.len() as f32, 1.0); - - let candidates: Vec<_> = candidates - .iter() - .filter(|c| !trajectory.contains(c)) - .collect(); - + let seen: HashSet<_> = trajectory.iter().collect(); + let candidates: Vec<_> = candidates.iter().filter(|c| !seen.contains(c)).collect(); if candidates.is_empty() { return None; } - let items = candidates + let Some(mood) = self.mood(trajectory) else { + let items: Vec<_> = candidates.iter().map(|&c| (c, 0.0)).collect(); + return softmax_sample(&items, temperature); + }; + + let items: Vec<_> = candidates .iter() .map(|&c| { - let sim = compute_sum(c, &trajectory, &self.similar); - let dif = compute_sum(c, &trajectory, &self.different); - let score = match action { - Action::Skip => dif - sim, - Action::More => sim - dif, + let score = match self.embedding.get(c.as_str()) { + None => 0.0, + Some(v) => { + let cos = dot(v, &mood); + match action { + Action::More => cos, + Action::Skip => -cos.abs(), + } + } }; - (c, score / normalisation) + (c, score) }) - .collect::<Vec<_>>(); + .collect(); softmax_sample(&items, temperature) } |