Skip to main content

entrenar/hf_pipeline/publish/
model_card.rs

1//! Model card generation for HuggingFace Hub
2//!
3//! Generates HF-compatible README.md with YAML front matter for model
4//! metadata, metrics, and training details.
5
6use crate::eval::evaluator::{EvalResult, Metric};
7
8/// Model card for HuggingFace Hub (README.md with YAML front matter)
9#[derive(Clone, Debug)]
10pub struct ModelCard {
11    /// Model name
12    pub model_name: String,
13    /// Model description
14    pub description: String,
15    /// License identifier (SPDX)
16    pub license: Option<String>,
17    /// Language codes (e.g., "en", "multilingual")
18    pub language: Vec<String>,
19    /// Tags for discoverability
20    pub tags: Vec<String>,
21    /// Evaluation metrics (name → value)
22    pub metrics: Vec<(String, f64)>,
23    /// Training details (free-form markdown)
24    pub training_details: Option<String>,
25    /// Base model repository ID
26    pub base_model: Option<String>,
27}
28
29impl ModelCard {
30    /// Create a model card from an `EvalResult`
31    #[must_use]
32    pub fn from_eval_result(result: &EvalResult) -> Self {
33        let metrics: Vec<(String, f64)> = result
34            .scores
35            .iter()
36            .map(|(metric, &value)| (metric_to_yaml_key(metric), value))
37            .collect();
38
39        Self {
40            model_name: result.model_name.clone(),
41            description: format!("Model: {}", result.model_name),
42            license: None,
43            language: Vec::new(),
44            tags: vec!["entrenar".to_string()],
45            metrics,
46            training_details: None,
47            base_model: None,
48        }
49    }
50
51    /// Render the model card as markdown with YAML front matter
52    #[must_use]
53    pub fn to_markdown(&self) -> String {
54        let mut out = String::new();
55        out.push_str("---\n");
56
57        // YAML front matter
58        if let Some(license) = &self.license {
59            out.push_str(&format!("license: {license}\n"));
60        }
61
62        if !self.language.is_empty() {
63            out.push_str("language:\n");
64            for lang in &self.language {
65                out.push_str(&format!("- {lang}\n"));
66            }
67        }
68
69        if !self.tags.is_empty() {
70            out.push_str("tags:\n");
71            for tag in &self.tags {
72                out.push_str(&format!("- {tag}\n"));
73            }
74        }
75
76        if let Some(base) = &self.base_model {
77            out.push_str(&format!("base_model: {base}\n"));
78        }
79
80        if !self.metrics.is_empty() {
81            out.push_str("model-index:\n");
82            out.push_str(&format!("- name: {}\n", self.model_name));
83            out.push_str("  results:\n");
84            out.push_str("  - metrics:\n");
85            for (name, value) in &self.metrics {
86                out.push_str(&format!("    - type: {name}\n"));
87                out.push_str(&format!("      value: {value}\n"));
88            }
89        }
90
91        out.push_str("---\n\n");
92
93        // Markdown body
94        out.push_str(&format!("# {}\n\n", self.model_name));
95        out.push_str(&format!("{}\n\n", self.description));
96
97        if !self.metrics.is_empty() {
98            out.push_str("## Evaluation Results\n\n");
99            out.push_str("| Metric | Value |\n");
100            out.push_str("|--------|-------|\n");
101            for (name, value) in &self.metrics {
102                out.push_str(&format!("| {name} | {value:.4} |\n"));
103            }
104            out.push('\n');
105        }
106
107        if let Some(details) = &self.training_details {
108            out.push_str("## Training Details\n\n");
109            out.push_str(details);
110            out.push('\n');
111        }
112
113        out.push_str("\n---\n*Generated by [entrenar](https://github.com/paiml/entrenar)*\n");
114
115        out
116    }
117}
118
119/// Convert a `Metric` to a YAML-friendly key name
120fn metric_to_yaml_key(metric: &Metric) -> String {
121    match metric {
122        Metric::WER => "wer".to_string(),
123        Metric::RTFx => "rtfx".to_string(),
124        Metric::BLEU => "bleu".to_string(),
125        Metric::ROUGE(v) => format!("{v}").to_lowercase().replace('-', "_"),
126        Metric::Perplexity => "perplexity".to_string(),
127        Metric::MMLUAccuracy => "mmlu_accuracy".to_string(),
128        Metric::PassAtK(k) => format!("pass_at_{k}"),
129        Metric::NDCGAtK(k) => format!("ndcg_at_{k}"),
130        other => other.name().to_lowercase(),
131    }
132}