35 lines
761 B
Rust
35 lines
761 B
Rust
use std::time::Instant;
|
|
|
|
/// Records a single generation run and computes TPS.
|
|
#[derive(Debug)]
|
|
pub struct GenerationRecord {
|
|
pub start: Instant,
|
|
pub prompt_tokens: u32,
|
|
pub completion_tokens: u32,
|
|
}
|
|
|
|
impl GenerationRecord {
|
|
pub fn new(prompt_tokens: u32) -> Self {
|
|
Self {
|
|
start: Instant::now(),
|
|
prompt_tokens,
|
|
completion_tokens: 0,
|
|
}
|
|
}
|
|
|
|
pub fn tick(&mut self) {
|
|
self.completion_tokens += 1;
|
|
}
|
|
|
|
pub fn elapsed_secs(&self) -> f64 {
|
|
self.start.elapsed().as_secs_f64()
|
|
}
|
|
|
|
pub fn tokens_per_second(&self) -> f64 {
|
|
if self.completion_tokens == 0 {
|
|
return 0.0;
|
|
}
|
|
self.completion_tokens as f64 / self.elapsed_secs()
|
|
}
|
|
}
|