interpolize 1.0.0

a rust program that scrapes discord, learns how your friends talk, and generates new messages in their collective voice. yes, this is what we're doing with our lives.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use anyhow::Result;
use clap::{Parser, Subcommand};

mod config;
mod scraper;
mod storage;
mod embed;
mod markov;
mod interpolate;
mod shell;

#[derive(Parser)]
#[command(name = "interpolize")]
struct Cli {
    #[arg(short, long, default_value = "config.toml")]
    config: String,

    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    Scrape,
    Train,
    Chat,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let config = config::Config::load(&cli.config)?;

    match cli.command {
        Command::Scrape => {
            let scraper = scraper::Scraper::new(&config)?;
            let incoming = scraper.scrape_all(&config).await?;
            let mut existing = storage::load("messages.toml")?;
            storage::merge(&mut existing, incoming);
            storage::save("messages.toml", &existing)?;
            eprintln!("saved {} total messages", existing.len());
        }

        Command::Train => {
            let messages = storage::load("messages.toml")?;
            if messages.is_empty() {
                anyhow::bail!("no messages found, run scrape first");
            }

            let embeddings = embed::Embeddings::train(
                &messages,
                config.embeddings.vector_dim,
                config.embeddings.window_size,
            );
            embeddings.save(&config.embeddings.storage_path)?;
            eprintln!("embeddings saved to {}", config.embeddings.storage_path);
        }

        Command::Chat => {
            let messages = storage::load("messages.toml")?;
            if messages.is_empty() {
                anyhow::bail!("no messages found, run scrape first");
            }

            let embeddings = embed::Embeddings::load(&config.embeddings.storage_path)
                .unwrap_or_else(|_| {
                    eprintln!("no embeddings found, training now...");
                    let e = embed::Embeddings::train(
                        &messages,
                        config.embeddings.vector_dim,
                        config.embeddings.window_size,
                    );
                    let _ = e.save(&config.embeddings.storage_path);
                    e
                });

            let texts: Vec<&str> = messages.iter().map(|m| m.content.as_str()).collect();
            let markov = markov::Markov::train(&texts, 5);

            let shell = shell::Shell::new(&config, &embeddings, &messages, markov);
            shell.run()?;
        }
    }

    Ok(())
}