diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/cli.rs | 60 | ||||
| -rw-r--r-- | src/config.rs | 47 | ||||
| -rw-r--r-- | src/learner.rs | 201 | ||||
| -rw-r--r-- | src/main.rs | 146 | ||||
| -rw-r--r-- | src/state.rs | 81 | ||||
| -rw-r--r-- | src/trajectory.rs | 39 |
6 files changed, 387 insertions, 187 deletions
diff --git a/src/cli.rs b/src/cli.rs new file mode 100644 index 0000000..f1563b3 --- /dev/null +++ b/src/cli.rs @@ -0,0 +1,60 @@ +use std::env; + +#[derive(Debug, PartialEq)] +pub enum Action { + NewConfig, + NewState, + DumpTrajectory, + DumpLearner, + DumpAlbums, + Run, +} + +const NAME: &str = env!("CARGO_PKG_NAME"); +const VERSION: &str = env!("CARGO_PKG_VERSION"); +const DESCRIPTION: &str = env!("CARGO_PKG_DESCRIPTION"); + +fn print_help() { + println!( + "\ + {NAME}: {DESCRIPTION} +Usage: {NAME} [OPTION] + +Options: + --help Show this help message. + --version Show version information. + --new-config Write a default config file and exit. + --new-state Write a default state file and exit. + --dump-trajectory Dump the play trajectory as JSON to stdout. + --dump-learner Dump learner weights as JSON to stdout. + --dump-albums Dump album list from state as JSON to stdout. +<no option> Suggest and play the next album" + ); +} + +pub fn parse() -> Option<Action> { + let mut args = env::args(); + + let _binary_name = args.next(); + + match args.next().as_deref() { + None => Some(Action::Run), + Some("--help") | Some("-h") => { + print_help(); + None + } + Some("--version") | Some("-V") | Some("-v") => { + println!("{NAME} {VERSION}"); + None + } + Some("--new-config") => Some(Action::NewConfig), + Some("--new-state") => Some(Action::NewState), + Some("--dump-trajectory") => Some(Action::DumpTrajectory), + Some("--dump-learner") => Some(Action::DumpLearner), + Some("--dump-albums") => Some(Action::DumpAlbums), + Some(unknown) => { + eprintln!("Unknown option: {unknown}"); + std::process::exit(1); + } + } +} diff --git a/src/config.rs b/src/config.rs index 0be1d09..f38b571 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,26 +1,44 @@ use directories::{ProjectDirs, UserDirs}; use serde::{Deserialize, Serialize}; +use std::fs; + +use expanduser::expanduser; #[derive(Serialize, Deserialize)] pub struct Config { pub music_roots: Vec<String>, pub audio_exts: Vec<String>, pub skip_window_secs: u64, - pub learning_rate: f32, - pub decay_rate: f32, - pub min_score_thresh: f32, pub temperature: f32, + pub low_weight: f32, + pub mid_weight: f32, + pub high_weight: f32, + pub max_weight: f32, } -pub fn default_state_path() -> String { +fn default_config_path() -> String { + let name = "mood.json"; ProjectDirs::from("qualifier", "organisation", "mood") - .and_then(|pd| { - pd.config_dir() - .join("mood.ron") - .to_str() - .map(|x| x.to_string()) - }) - .unwrap_or("~/.config/mood/mood.ron".to_string()) + .and_then(|pd| pd.config_dir().join(name).to_str().map(|x| x.to_string())) + .unwrap_or(format!("~/.config/mood/{name}")) +} + +impl Config { + pub fn try_load() -> Result<Config, String> { + expanduser(default_config_path()) + .map_err(|e| e.to_string()) + .and_then(|p| fs::read_to_string(p).map_err(|e| e.to_string())) + .and_then(|contents| serde_json::from_str(&contents).map_err(|e| e.to_string())) + } + + pub fn try_save(&self) -> Result<(), String> { + let path = expanduser(default_config_path()).map_err(|e| e.to_string())?; + _ = path.parent().map(std::fs::create_dir_all); + let serialised = serde_json::to_string_pretty(self) + .map_err(|e| format!("Failed to serialize config: {}", e))?; + fs::write(&path, serialised).map_err(|e| format!("Failed to write config file: {}", e))?; + Ok(()) + } } impl Default for Config { @@ -39,10 +57,11 @@ impl Default for Config { .map(String::from) .collect(), skip_window_secs: 60, - learning_rate: 0.1, - decay_rate: 0.001, - min_score_thresh: 1e-5, temperature: 0.6, + low_weight: 0.25, + mid_weight: 0.5, + high_weight: 0.75, + max_weight: 1.0, } } } diff --git a/src/learner.rs b/src/learner.rs index 91b24fe..59db567 100644 --- a/src/learner.rs +++ b/src/learner.rs @@ -5,24 +5,15 @@ 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) } -} - -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 compute_weight(past: &HashSet<String>, now: &HashSet<&String>) -> f32 { + let inter = now.iter().filter(|&&p| past.contains(p)).count(); + if inter == 0 { + return 0.0; + } + (inter as f32) / ((now.len() as f32).sqrt() * (past.len() as f32).sqrt()) } -fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, f32)> { +fn softmax_sample(items: &[(String, f32)], temperature: f32) -> Option<(String, f32)> { if items.is_empty() { return None; } @@ -35,7 +26,7 @@ fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, let exp_values: Vec<_> = items .iter() - .map(|&(k, val)| (k, ((val - max_val) / temperature).exp())) + .map(|(k, val)| (k, ((val - max_val) / temperature).exp())) .collect(); let sum: f32 = exp_values.iter().map(|(_, v)| *v).sum(); @@ -48,80 +39,79 @@ fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, .ok() } +#[derive(Serialize, Deserialize)] +enum Episode { + Continue { + group: HashSet<String>, + avoid: Option<String>, + committed: String, + }, + Escape { + from: HashSet<String>, + to: String, + committed: String, + }, +} + +impl Episode { + fn prune_mut(&mut self, valid: &HashSet<String>) -> bool { + match self { + Episode::Continue { group, avoid, .. } => { + if let Some(a) = avoid + && valid.contains(a) + { + } else { + *avoid = None; + }; + group.retain(|g| valid.contains(g)); + !group.is_empty() + } + Episode::Escape { from, to, .. } => { + if !valid.contains(to) { + return false; + }; + from.retain(|f| valid.contains(f)); + !from.is_empty() + } + } + } +} + #[derive(Serialize, Deserialize, Default)] pub struct Learner { - similar: HashMap<(String, String), f32>, - different: HashMap<(String, String), f32>, + history: Vec<Episode>, } -#[derive(Debug)] pub enum Learning { - SkipExtend, - MoreExtend(Vec<String>, String), SkipToMore(Vec<String>, String), MoreToSkip(Vec<String>, String), } 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)); + self.history.retain_mut(|e| e.prune_mut(valid)); } +} - 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 +impl Learner { + pub fn learn(&mut self, learning: &Learning) { + let stamp = chrono::Utc::now().format("U%Y%m%d-%H%M%S").to_string(); 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. - } - 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() { - 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; - } - } 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; - } + self.history.push(Episode::Escape { + from: trajectory.iter().map(String::clone).collect(), + to: new.clone(), + committed: stamp, + }); } 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. - 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; - } + self.history.push(Episode::Continue { + group: trajectory.iter().map(String::clone).collect(), + avoid: Some(new.clone()), + committed: stamp, + }); } }; - // 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(); } pub fn sample( @@ -129,10 +119,9 @@ impl Learner { trajectory: &[String], action: &Action, candidates: &HashSet<String>, - temperature: f32, + config: &Config, ) -> 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() @@ -143,19 +132,57 @@ impl Learner { return None; } - let items = 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, - }; - (c, score / normalisation) - }) - .collect::<Vec<_>>(); + let mut items: HashMap<String, f32> = + candidates.into_iter().map(|c| (c.clone(), 0.0)).collect(); + + for episode in self.history.iter() { + match (episode, action) { + (Episode::Escape { from, to, .. }, Action::Skip) => { + let w = compute_weight(from, &trajectory); + if let Some(to_w) = items.get_mut(to) { + *to_w += config.max_weight * w; + } + for f in from { + if let Some(from_weight) = items.get_mut(f) { + *from_weight -= config.mid_weight * w; + } + } + } + (Episode::Escape { from, to, .. }, Action::More) => { + let w = compute_weight(from, &trajectory); + if trajectory.contains(to) { + for f in from { + if let Some(f_w) = items.get_mut(f) { + *f_w -= config.low_weight * w; + } + } + } + } + (Episode::Continue { group, avoid, .. }, Action::More) => { + let w = compute_weight(group, &trajectory); + for g in group { + if let Some(g_w) = items.get_mut(g) { + *g_w += config.max_weight * w; + } + } + if let Some(a) = avoid + && let Some(a_w) = items.get_mut(a) + { + *a_w -= config.low_weight * w; + } + } + (Episode::Continue { group, .. }, Action::Skip) => { + let w = compute_weight(group, &trajectory); + for g in group { + if let Some(v) = items.get_mut(g) { + *v -= config.high_weight * w; + } + } + } + } + } - softmax_sample(&items, temperature) + let pairs: Vec<_> = items.into_iter().collect(); + softmax_sample(&pairs, config.temperature) } } diff --git a/src/main.rs b/src/main.rs index 74f90c8..86c6341 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,61 +1,135 @@ -use std::process::Command; +use std::{process::Command, process::exit}; +mod cli; mod config; mod learner; mod state; mod trajectory; -use crate::{config::default_state_path, state::State}; +use crate::cli::{Action, parse}; +use crate::config::Config; +use crate::state::State; fn mpc_load_and_play(path: &str) { - _ = Command::new("mpc") - .arg("clear") - .output() - .expect("failed to execute mpc clear"); - - _ = Command::new("mpc") - .args(["add", path]) - .output() - .expect("failed to execute mpc add"); - - _ = Command::new("mpc") - .arg("play") - .output() - .expect("failed to execute mpc play"); + _ = Command::new("mpc").arg("clear").output(); + + _ = Command::new("mpc").args(["add", path]).output(); + + _ = Command::new("mpc").arg("play").output(); } fn notify(text: &str) { if let Err(e) = Command::new("notify-send") - .args(["-t", "5000", "-a", "mood", &text]) + .args(["-t", "5000", "-a", "mood", text]) .status() { eprintln!("Failed to show notification: {e}"); } } +fn load_state_no_default() -> State { + State::try_load().unwrap_or_else(|e| { + eprintln!("Failed to load state: {e}, nothing to dump."); + exit(1); + }) +} + fn main() { - let state = State::load(&default_state_path()); - let mut state = match state { - Err(e) => { - eprintln!("Failed to load state: {e}"); - State::default() - } - Ok(s) => s, + let action = parse(); + if action.is_none() { + return; }; - let ur = state.update_albums(); - println!("Update result: {ur}"); + let action = action.unwrap(); - if let Some((next_album, prob)) = state.next() { - let text = format!("Now playing: {next_album} ({:.2}%)", prob * 100.0); - println!("{text}"); - mpc_load_and_play(&next_album); - notify(&text); - } else { - println!("No next album available."); - } + match action { + Action::DumpAlbums => { + let state = load_state_no_default(); + println!( + "{}", + state.dump_albums().unwrap_or_else(|e| { + eprintln!("Failed to serialize state (dump-albums): {e}"); + exit(1); + }) + ); + } + + Action::DumpLearner => { + let state = load_state_no_default(); + println!( + "{}", + state.dump_learner().unwrap_or_else(|e| { + eprintln!("Failed to serialize state (dump-learner): {e}"); + exit(1); + }) + ); + } + + Action::DumpTrajectory => { + let state = load_state_no_default(); + println!( + "{}", + state.dump_trajectory().unwrap_or_else(|e| { + eprintln!("Failed to serialize state (dump-trajectory): {e}"); + exit(1); + }) + ); + } + + Action::NewState => { + let default = State::default(); + if let Err(e) = default.try_save() { + eprintln!("Failed to write new state: {e}"); + exit(1); + } + println!("New state written."); + } + + Action::NewConfig => { + let default = Config::default(); + if let Err(e) = default.try_save() { + eprintln!("Failed to write new config: {e}"); + exit(1); + } + println!("New config written."); + } - if let Err(e) = state.try_save() { - eprintln!("Failed to save state: {e}"); + Action::Run => { + let config = match Config::try_load() { + Ok(c) => c, + Err(e) => { + eprintln!("Failed to load config: {e}, writing defaults."); + let default = Config::default(); + if let Err(se) = default.try_save() { + panic!("Failed to save default config: {se}"); + } + default + } + }; + + let mut state = match State::try_load() { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to load state: {e}"); + State::default() + } + }; + + let ur = state.update_albums(&config); + println!("Update result: {ur}"); + + if let Some((next_album, prob)) = state.next(&config) { + let text = format!("Now playing: {next_album} ({:.2}%)", prob * 100.0); + println!("{text}"); + mpc_load_and_play(&next_album); + notify(&text); + } else { + println!("No next album available."); + } + + if let Err(e) = state.try_save() { + panic!("Failed to save state: {e}"); + } + } } } diff --git a/src/state.rs b/src/state.rs index b9a6618..a04412f 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,36 +1,36 @@ -use std::{collections::HashSet, fs}; - -use derive_more::Display; +use std::{collections::HashSet, fmt, fs}; +use directories::ProjectDirs; use expanduser::expanduser; use serde::{Deserialize, Serialize}; use walkdir::WalkDir; -use crate::{ - config::{Config, default_state_path}, - learner::Learner, - trajectory::Trajectory, -}; +use crate::{config::Config, learner::Learner, trajectory::Trajectory}; + +fn default_state_path() -> String { + let name = "mood.bin"; + ProjectDirs::from("qualifier", "organisation", "mood") + .and_then(|pd| pd.config_dir().join(name).to_str().map(|x| x.to_string())) + .unwrap_or(format!("~/.config/mood/{name}")) +} #[derive(Serialize, Deserialize, Default)] pub struct State { - pub config: Config, trajectory: Trajectory, learner: Learner, albums: HashSet<String>, } impl State { - pub fn next(&mut self) -> Option<(String, f32)> { - let learning = self.trajectory.step(self.config.skip_window_secs); + pub fn next(&mut self, config: &Config) -> Option<(String, f32)> { + let learning = self.trajectory.step(config.skip_window_secs); if let Some(ref learning) = learning { - println!("Obtained learning: {learning:?}"); - self.learner.learn(&self.config, learning); + self.learner.learn(learning); } if let Some((new_album, prob)) = self.learner.sample( self.trajectory.slice(), &learning.as_ref().into(), &self.albums, - self.config.temperature, + config, ) { self.trajectory.log(&new_album); Some((new_album, prob)) @@ -40,36 +40,39 @@ impl State { } } -#[derive(Display)] -#[display("+{added}/-{removed}, total {new_total}")] pub struct UpdateResult { added: usize, removed: usize, new_total: usize, } +impl fmt::Display for UpdateResult { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "+{}/-{}, total {}", self.added, self.removed, self.new_total) + } +} + impl State { - pub fn load(path: &str) -> Result<State, String> { - expanduser(path) + pub fn try_load() -> Result<State, String> { + expanduser(default_state_path()) .map_err(|e| e.to_string()) - .and_then(|p| fs::read_to_string(p).map_err(|e| e.to_string())) - .and_then(|contents| ron::from_str(&contents).map_err(|e| e.to_string())) + .and_then(|p| fs::read(p).map_err(|e| e.to_string())) + .and_then(|contents| bincode::deserialize(&contents).map_err(|e| e.to_string())) } pub fn try_save(&self) -> Result<(), String> { let str_path = default_state_path(); let path = expanduser(&str_path).map_err(|e| e.to_string())?; - _ = path.parent().map(|parent| std::fs::create_dir_all(parent)); - let serialised = ron::ser::to_string_pretty(self, ron::ser::PrettyConfig::default()) - .map_err(|e| format!("Failed to serialize state: {}", e))?; + _ = path.parent().map(std::fs::create_dir_all); + let serialised = + bincode::serialize(self).map_err(|e| format!("Failed to serialize state: {}", e))?; fs::write(&path, serialised) .map_err(|e| format!("Failed to write state file at {}: {}", str_path, e))?; Ok(()) } - pub fn update_albums(&mut self) -> UpdateResult { - let roots: Vec<_> = self - .config + pub fn update_albums(&mut self, config: &Config) -> UpdateResult { + let roots: Vec<_> = config .music_roots .iter() .filter_map(|x| expanduser(x).ok().filter(|p| p.is_dir())) @@ -77,14 +80,14 @@ impl State { let mut new_albums: HashSet<String> = HashSet::new(); - for root_path in roots { + for root_path in &roots { for entry in WalkDir::new(&root_path) { let entry = match entry { Ok(e) if e.path().is_dir() => e, _ => continue, }; - if !self.contains_audio(&entry) { + if !Self::contains_audio(&entry, config) { continue; } @@ -113,21 +116,35 @@ impl State { } } - fn contains_audio(&self, entry: &walkdir::DirEntry) -> bool { + fn contains_audio(entry: &walkdir::DirEntry, config: &Config) -> bool { fs::read_dir(entry.path()).ok().is_some_and(|mut entries| { entries.any(|e| { e.as_ref() - .is_ok_and(|entry| self.is_audio_file(&entry.path())) + .is_ok_and(|entry| Self::is_audio_file(&entry.path(), config)) }) }) } - fn is_audio_file(&self, path: &std::path::Path) -> bool { + fn is_audio_file(path: &std::path::Path, config: &Config) -> bool { !path.is_dir() && path .extension() .and_then(|ext| ext.to_str()) .map(|s| s.to_lowercase()) - .is_some_and(|ext| self.config.audio_exts.contains(&format!(".{ext}"))) + .is_some_and(|ext| config.audio_exts.contains(&format!(".{ext}"))) + } +} + +impl State { + pub fn dump_albums(&self) -> Result<String, String> { + serde_json::to_string_pretty(&self.albums).map_err(|e| e.to_string()) + } + + pub fn dump_learner(&self) -> Result<String, String> { + serde_json::to_string_pretty(&self.learner).map_err(|e| e.to_string()) + } + + pub fn dump_trajectory(&self) -> Result<String, String> { + serde_json::to_string_pretty(&self.trajectory).map_err(|e| e.to_string()) } } diff --git a/src/trajectory.rs b/src/trajectory.rs index 214e16c..cbb0661 100644 --- a/src/trajectory.rs +++ b/src/trajectory.rs @@ -3,7 +3,7 @@ use std::{collections::HashSet, time::SystemTime}; use crate::learner::Learning; -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize)] struct LastData { timestamp: SystemTime, album: String, @@ -16,7 +16,7 @@ pub struct Trajectory { history: Vec<String>, } -#[derive(Serialize, Deserialize, Clone, PartialEq, Debug)] +#[derive(Serialize, Deserialize, Clone, PartialEq)] pub enum Action { Skip, More, @@ -26,8 +26,8 @@ impl From<Option<&Learning>> for Action { fn from(learning: Option<&Learning>) -> Action { if let Some(ts) = learning { match ts { - Learning::SkipExtend | Learning::MoreToSkip(_, _) => Action::Skip, - Learning::MoreExtend(_, _) | Learning::SkipToMore(_, _) => Action::More, + Learning::MoreToSkip(_, _) => Action::Skip, + Learning::SkipToMore(_, _) => Action::More, } } else { Action::Skip @@ -71,26 +71,29 @@ impl Trajectory { let last_album = last_data.album.clone(); let current_streak = self.history.clone(); + let learnt_nothing = self + .streak_kind + .as_ref() + .map(|kind| *kind == action) + .unwrap_or(false); - let learning = if let Some(ref kind) = self.streak_kind - && kind == &action - { - self.history.push(last_album.clone()); - match action { - Action::Skip => Learning::SkipExtend, - Action::More => Learning::MoreExtend(current_streak, last_album), + let learning = match (learnt_nothing, &action) { + (false, Action::Skip) => { + Some(Learning::MoreToSkip(current_streak, last_album.clone())) } - } else { - self.history.clear(); - self.history.push(last_album.clone()); - match action { - Action::Skip => Learning::MoreToSkip(current_streak, last_album), - Action::More => Learning::SkipToMore(current_streak, last_album), + (false, Action::More) => { + Some(Learning::SkipToMore(current_streak, last_album.clone())) } + _ => None, + }; + + if !learnt_nothing { + self.history.clear(); }; + self.history.push(last_album); self.streak_kind = Some(action); - Some(learning) + learning } else { None } |