tokengeex 1.1.0

TokenGeeX is an efficient tokenizer for code based on UnigramLM and TokenMonster.
Documentation 
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use rayon::current_num_threads;
use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed},
        Arc,
    },
    time::Instant,
};

#[derive(Clone)]
pub struct Task {
    inner: Arc<TaskInner>,
}

struct TaskInner {
    pub start: Instant,
    pub desc: String,
    pub num_samples: usize,
    pub num_samples_processed: AtomicUsize,
    pub num_bytes_processed: AtomicUsize,
    pub finished: AtomicBool,
}

pub struct LocalTask {
    task: Arc<TaskInner>,
    num_chunk_samples: usize,
    num_bytes_processed: usize,
}

impl Task {
    pub fn new(desc: &str, num_samples: usize, chunk_size: usize) -> Self {
        log::info!(
            "{} | {} samples | {} threads | {} chunks | {} chunk size",
            desc,
            num_samples,
            current_num_threads(),
            (num_samples + chunk_size) / chunk_size,
            chunk_size
        );

        Self {
            inner: Arc::new(TaskInner {
                start: Instant::now(),
                desc: desc.to_string(),
                num_samples,
                num_samples_processed: AtomicUsize::new(0),
                num_bytes_processed: AtomicUsize::new(0),
                finished: AtomicBool::new(false),
            }),
        }
    }

    pub fn start(&self) {
        let task = self.inner.clone();

        std::thread::spawn(move || {
            loop {
                std::thread::sleep(std::time::Duration::from_secs(1));

                if task.finished.load(Relaxed) {
                    break;
                }

                let num_samples_processed = task.num_samples_processed.load(Relaxed);
                let num_bytes_processed = task.num_bytes_processed.load(Relaxed);

                if num_samples_processed >= task.num_samples {
                    break;
                }

                let percent_done = (num_samples_processed as f64 / task.num_samples as f64) * 100.0;

                if percent_done == 0.0 {
                    continue;
                }

                let eta =
                    (task.start.elapsed().as_secs_f64() / percent_done) * (100.0 - percent_done);

                log::debug!(
                    "{} | {:>6.2}% | ETA {:>5}s | {:>5.2}MB/s | {:>5.2}MB/s per thread",
                    task.desc,
                    percent_done,
                    eta.round(),
                    mb_per_sec(num_bytes_processed, task.start),
                    mb_per_sec(num_bytes_processed, task.start) / current_num_threads() as f64,
                );
            }

            let num_bytes_processed = task.num_bytes_processed.load(Relaxed);

            log::info!(
                "FINISHED {} | {} samples | {:.2}MB/s | {:.2}s",
                task.desc,
                task.num_samples,
                mb_per_sec(num_bytes_processed, task.start),
                task.start.elapsed().as_secs_f64(),
            );
        });
    }

    pub fn local(&self, num_samples: usize) -> LocalTask {
        LocalTask {
            task: self.inner.clone(),
            num_chunk_samples: num_samples,
            num_bytes_processed: 0,
        }
    }

    pub fn finish(&self) {
        self.inner.finished.store(true, Relaxed);
    }
}

impl LocalTask {
    pub fn record(&mut self, num_bytes: usize) {
        self.num_bytes_processed += num_bytes;
    }

    pub fn finish(&self) {
        self.task
            .num_samples_processed
            .fetch_add(self.num_chunk_samples, Relaxed);
        self.task
            .num_bytes_processed
            .fetch_add(self.num_bytes_processed, Relaxed);
    }
}

// We chunk up samples into chunks that are at most f/10th of the
// per-thread workload because too large chunks can cause some threads
// to be idle while others are still working. We also prevent
// chunks from being too small to avoid too much overhead.
pub fn par_chunk_size(num_samples: usize, f: usize) -> usize {
    let chunk_size = num_samples / current_num_threads() / f;
    std::cmp::max(1, chunk_size)
}

pub fn mb_per_sec(n: usize, since: std::time::Instant) -> f64 {
    (n as f64 / 1024.0 / 1024.0) / since.elapsed().as_secs_f64()
}