llm_token_visualizer/
data.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize)]
4pub struct TokenInfo {
5    pub text: String,
6    pub confidence: f64,
7}
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct TokenFlag {
11    pub start: usize,
12    pub end: usize,
13    pub flag: String,
14    pub description: Option<String>,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct TokenAnalysis {
19    pub tokens: Vec<TokenInfo>,
20    pub flags: Vec<TokenFlag>,
21}
22
23#[derive(Debug, Clone)]
24pub struct VisualizationConfig {
25    pub verbose: bool,
26    pub show_confidence_scores: bool,
27    pub show_flags: bool,
28}
29
30impl Default for VisualizationConfig {
31    fn default() -> Self {
32        Self {
33            verbose: false,
34            show_confidence_scores: true,
35            show_flags: true,
36        }
37    }
38}
39
40#[derive(Debug, Clone, PartialEq)]
41pub enum ConfidenceLevel {
42    VeryLow,  // 0.0 - 0.3
43    Low,      // 0.3 - 0.5
44    Medium,   // 0.5 - 0.7
45    High,     // 0.7 - 0.9
46    VeryHigh, // 0.9 - 1.0
47}
48
49impl From<f64> for ConfidenceLevel {
50    fn from(confidence: f64) -> Self {
51        match confidence {
52            c if c < 0.3 => ConfidenceLevel::VeryLow,
53            c if c < 0.5 => ConfidenceLevel::Low,
54            c if c < 0.7 => ConfidenceLevel::Medium,
55            c if c < 0.9 => ConfidenceLevel::High,
56            _ => ConfidenceLevel::VeryHigh,
57        }
58    }
59}
60
61#[derive(Debug, Clone, PartialEq)]
62pub enum FlagType {
63    Fact,
64    Uncertain,
65    Overconfident,
66    Hallucination,
67    Other(String),
68}
69
70impl From<&str> for FlagType {
71    fn from(s: &str) -> Self {
72        match s.to_lowercase().as_str() {
73            "fact" => FlagType::Fact,
74            "uncertain" => FlagType::Uncertain,
75            "overconfident" => FlagType::Overconfident,
76            "hallucination" => FlagType::Hallucination,
77            _ => FlagType::Other(s.to_string()),
78        }
79    }
80}
81
82impl TokenAnalysis {
83    pub fn validate(&self) -> Result<(), String> {
84        if self.tokens.is_empty() {
85            return Err("No tokens provided".to_string());
86        }
87
88        for (i, token) in self.tokens.iter().enumerate() {
89            if token.confidence < 0.0 || token.confidence > 1.0 {
90                return Err(format!(
91                    "Invalid confidence score for token {}: {}",
92                    i, token.confidence
93                ));
94            }
95        }
96
97        for flag in &self.flags {
98            if flag.start >= self.tokens.len() || flag.end > self.tokens.len() {
99                return Err(format!(
100                    "Flag span out of bounds: {} to {}",
101                    flag.start, flag.end
102                ));
103            }
104            if flag.start >= flag.end {
105                return Err(format!("Invalid flag span: {} to {}", flag.start, flag.end));
106            }
107        }
108
109        Ok(())
110    }
111
112    pub fn get_flags_for_token(&self, token_index: usize) -> Vec<&TokenFlag> {
113        self.flags
114            .iter()
115            .filter(|flag| token_index >= flag.start && token_index < flag.end)
116            .collect()
117    }
118
119    pub fn get_confidence_stats(&self) -> (f64, f64, f64) {
120        let confidences: Vec<f64> = self.tokens.iter().map(|t| t.confidence).collect();
121        let sum: f64 = confidences.iter().sum();
122        let avg = sum / confidences.len() as f64;
123        let min = confidences.iter().fold(f64::INFINITY, |a, &b| a.min(b));
124        let max = confidences.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
125        (min, max, avg)
126    }
127}
llm_token_visualizer/data.rs

llm_token_visualizer/
data.rs
