markov_text/
lib.rs

1use std::{collections::HashMap, hash::Hash};
2
3use rand::Rng;
4
5#[derive(Hash, Clone, PartialEq, Eq, Debug)]
6enum Token {
7    Start,
8    Text(String),
9    End,
10}
11
12#[derive(Debug, Clone)]
13struct TokenSampler {
14    token_count: usize,
15    token_to_token_count: HashMap<Token, usize>,
16}
17
18impl TokenSampler {
19    fn new() -> TokenSampler {
20        TokenSampler {
21            token_count: 0,
22            token_to_token_count: HashMap::new(),
23        }
24    }
25
26    fn add_token(&mut self, token: Token) {
27        self.token_count += 1;
28        *self.token_to_token_count.entry(token).or_insert(0) += 1;
29    }
30
31    fn sample(&self) -> Token {
32        let mut random_index = rand::thread_rng().gen_range(0..self.token_count);
33        for (token, count) in &self.token_to_token_count {
34            if &random_index >= count {
35                random_index -= count;
36            } else {
37                return token.clone();
38            }
39        }
40        unreachable!("The token count does not match the tokens in the map");
41    }
42}
43
44#[derive(Debug, Clone)]
45pub struct MarkovTextModel<const CONTEXT_LENGTH: usize = 1> {
46    token_to_token_sampler: HashMap<[Token; CONTEXT_LENGTH], TokenSampler>,
47}
48
49impl<const CONTEXT_LENGTH: usize> MarkovTextModel<CONTEXT_LENGTH> {
50    /// Creates an empty character based MarkovTextModel, where the `CONTEXT_LENGTH` parameter is the context length in characters.
51    pub fn new() -> MarkovTextModel<CONTEXT_LENGTH> {
52        assert!(
53            CONTEXT_LENGTH > 0,
54            "The context length must be greater than 0"
55        );
56        MarkovTextModel::<CONTEXT_LENGTH> {
57            token_to_token_sampler: HashMap::new(),
58        }
59    }
60
61    pub fn add_sample_text(&mut self, text: &str) {
62        self.add_tokenized_sample_text(&self.tokenize(text));
63    }
64
65    pub fn add_sample_texts(&mut self, texts: &Vec<String>) {
66        for text in texts {
67            self.add_sample_text(text);
68        }
69    }
70
71    pub fn generage_text(&self) -> String {
72        let mut current_context = std::array::from_fn::<_, CONTEXT_LENGTH, _>(|_| Token::Start);
73        let mut cumulative_text = "".to_owned();
74
75        loop {
76            let next_token = self.token_to_token_sampler[&current_context].sample();
77
78            match &next_token {
79                Token::Start => {
80                    unreachable!("Start token should not be reachable after Start token")
81                }
82                Token::End => return cumulative_text,
83                Token::Text(current_token_text) => {
84                    cumulative_text = cumulative_text + current_token_text
85                }
86            }
87
88            for i in 1..CONTEXT_LENGTH {
89                current_context[i - 1] = current_context[i].clone();
90            }
91            current_context[CONTEXT_LENGTH - 1] = next_token;
92        }
93    }
94
95    // TODO: Make this more user definable. Maybe with an optional closure when constructing the model?
96    fn tokenize(&self, text: &str) -> Vec<Token> {
97        let mut tokens = vec![Token::Start; CONTEXT_LENGTH];
98        tokens.extend(text.chars().map(|letter| Token::Text(letter.to_string())));
99        // tokens.extend(text.split_whitespace().map(|word| Token::Word(word.to_owned())));
100        tokens.push(Token::End);
101        tokens
102    }
103
104    fn add_tokenized_sample_text(&mut self, tokenized_text: &[Token]) {
105        for window in tokenized_text.windows(CONTEXT_LENGTH + 1) {
106            let mut context = std::array::from_fn::<_, CONTEXT_LENGTH, _>(|_| Token::Start);
107
108            context[..CONTEXT_LENGTH].clone_from_slice(&window[..CONTEXT_LENGTH]);
109
110            let next_token = &window[CONTEXT_LENGTH];
111
112            self.token_to_token_sampler
113                .entry(context)
114                .or_insert(TokenSampler::new())
115                .add_token(next_token.clone());
116        }
117    }
118}
119
120impl<const CONTEXT_LENGTH: usize> Default for MarkovTextModel<CONTEXT_LENGTH> {
121    fn default() -> Self {
122        Self::new()
123    }
124}
markov_text/lib.rs

markov_text/
lib.rs
