trustformers/
hub_model_card.rs

1//! Automatic model card (README.md) generation for HuggingFace Hub.
2//!
3//! Generates markdown following the HuggingFace model card specification,
4//! including YAML front matter metadata and structured sections.
5
6use crate::error::{Result, TrustformersError};
7use serde::{Deserialize, Serialize};
8use std::path::Path;
9use tracing::debug;
10
11// ─── Metadata ─────────────────────────────────────────────────────────────────
12
13/// Model card metadata stored as YAML front matter
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct ModelCardMetadata {
16    /// Languages the model supports, e.g. `["en", "fr"]`
17    pub language: Vec<String>,
18    /// SPDX license identifier, e.g. `"apache-2.0"`
19    pub license: String,
20    /// Library name, typically `"trustformers"`
21    pub library_name: String,
22    /// Arbitrary tags, e.g. `["text-generation", "causal-lm"]`
23    pub tags: Vec<String>,
24    /// Training datasets referenced by the model
25    pub datasets: Vec<String>,
26    /// Evaluation metrics reported for this model
27    pub metrics: Vec<String>,
28    /// Model architecture type, e.g. `"bert"` or `"gpt2"`
29    pub model_type: Option<String>,
30    /// HuggingFace pipeline tag, e.g. `"text-generation"`
31    pub pipeline_tag: Option<String>,
32}
33
34impl Default for ModelCardMetadata {
35    fn default() -> Self {
36        Self {
37            language: vec!["en".to_string()],
38            license: "apache-2.0".to_string(),
39            library_name: "trustformers".to_string(),
40            tags: Vec::new(),
41            datasets: Vec::new(),
42            metrics: Vec::new(),
43            model_type: None,
44            pipeline_tag: None,
45        }
46    }
47}
48
49// ─── BenchmarkResult ──────────────────────────────────────────────────────────
50
51/// A single benchmark measurement for a model
52#[derive(Debug, Clone)]
53pub struct BenchmarkResult {
54    /// Task type, e.g. `"text-classification"`
55    pub task: String,
56    /// Benchmark dataset, e.g. `"glue/sst2"`
57    pub dataset: String,
58    /// Metric name, e.g. `"accuracy"`
59    pub metric: String,
60    /// Numeric value of the metric
61    pub value: f64,
62}
63
64impl BenchmarkResult {
65    /// Create a new benchmark result
66    pub fn new(
67        task: impl Into<String>,
68        dataset: impl Into<String>,
69        metric: impl Into<String>,
70        value: f64,
71    ) -> Self {
72        Self {
73            task: task.into(),
74            dataset: dataset.into(),
75            metric: metric.into(),
76            value,
77        }
78    }
79}
80
81// ─── TrainingInfo ─────────────────────────────────────────────────────────────
82
83/// Information about the model training process
84#[derive(Debug, Clone)]
85pub struct TrainingInfo {
86    /// Training framework name
87    pub framework: String,
88    /// Total number of trainable parameters
89    pub num_parameters: Option<u64>,
90    /// List of training data sources
91    pub training_data: Vec<String>,
92    /// Optimizer name, e.g. `"AdamW"`
93    pub optimizer: Option<String>,
94    /// Peak learning rate
95    pub learning_rate: Option<f64>,
96    /// Per-device batch size
97    pub batch_size: Option<usize>,
98    /// Number of training epochs
99    pub num_epochs: Option<usize>,
100    /// Hardware used for training, e.g. `"4x A100 80GB"`
101    pub hardware: Option<String>,
102}
103
104impl Default for TrainingInfo {
105    fn default() -> Self {
106        Self {
107            framework: "TrustformeRS".to_string(),
108            num_parameters: None,
109            training_data: Vec::new(),
110            optimizer: None,
111            learning_rate: None,
112            batch_size: None,
113            num_epochs: None,
114            hardware: None,
115        }
116    }
117}
118
119// ─── ModelCard ────────────────────────────────────────────────────────────────
120
121/// Full model card content combining metadata and markdown sections
122#[derive(Debug, Clone)]
123pub struct ModelCard {
124    /// YAML front matter metadata
125    pub metadata: ModelCardMetadata,
126    /// Human-readable model name
127    pub model_name: String,
128    /// Short description of what the model does
129    pub model_description: String,
130    /// List of intended use cases
131    pub intended_uses: Vec<String>,
132    /// Known limitations of the model
133    pub limitations: Vec<String>,
134    /// Training details
135    pub training_info: TrainingInfo,
136    /// Benchmark results
137    pub benchmarks: Vec<BenchmarkResult>,
138    /// BibTeX or plain-text citation
139    pub citation: Option<String>,
140    /// Author or organisation name
141    pub author: Option<String>,
142}
143
144impl ModelCard {
145    /// Create a minimal model card from a name and description
146    pub fn new(model_name: impl Into<String>, description: impl Into<String>) -> Self {
147        Self {
148            metadata: ModelCardMetadata::default(),
149            model_name: model_name.into(),
150            model_description: description.into(),
151            intended_uses: Vec::new(),
152            limitations: Vec::new(),
153            training_info: TrainingInfo::default(),
154            benchmarks: Vec::new(),
155            citation: None,
156            author: None,
157        }
158    }
159
160    /// Generate the full markdown string for the model card
161    pub fn to_markdown(&self) -> String {
162        let mut md = String::new();
163
164        // ── YAML front matter ──────────────────────────────────────────────
165        md.push_str("---\n");
166        if !self.metadata.language.is_empty() {
167            md.push_str("language:\n");
168            for lang in &self.metadata.language {
169                md.push_str(&format!("- {lang}\n"));
170            }
171        }
172        md.push_str(&format!("license: {}\n", self.metadata.license));
173        md.push_str(&format!("library_name: {}\n", self.metadata.library_name));
174        if !self.metadata.tags.is_empty() {
175            md.push_str("tags:\n");
176            for tag in &self.metadata.tags {
177                md.push_str(&format!("- {tag}\n"));
178            }
179        }
180        if !self.metadata.datasets.is_empty() {
181            md.push_str("datasets:\n");
182            for ds in &self.metadata.datasets {
183                md.push_str(&format!("- {ds}\n"));
184            }
185        }
186        if !self.metadata.metrics.is_empty() {
187            md.push_str("metrics:\n");
188            for m in &self.metadata.metrics {
189                md.push_str(&format!("- {m}\n"));
190            }
191        }
192        if let Some(ref mt) = self.metadata.model_type {
193            md.push_str(&format!("model_type: {mt}\n"));
194        }
195        if let Some(ref pt) = self.metadata.pipeline_tag {
196            md.push_str(&format!("pipeline_tag: {pt}\n"));
197        }
198        md.push_str("---\n\n");
199
200        // ── Title ──────────────────────────────────────────────────────────
201        md.push_str(&format!("# {}\n\n", self.model_name));
202
203        // ── Author ─────────────────────────────────────────────────────────
204        if let Some(ref author) = self.author {
205            md.push_str(&format!("*Author: {author}*\n\n"));
206        }
207
208        // ── Model Description ──────────────────────────────────────────────
209        md.push_str("## Model Description\n\n");
210        md.push_str(&self.model_description);
211        md.push_str("\n\n");
212
213        // ── Intended Uses ──────────────────────────────────────────────────
214        if !self.intended_uses.is_empty() {
215            md.push_str("## Intended Uses\n\n");
216            for use_case in &self.intended_uses {
217                md.push_str(&format!("- {use_case}\n"));
218            }
219            md.push('\n');
220        }
221
222        // ── Limitations ────────────────────────────────────────────────────
223        if !self.limitations.is_empty() {
224            md.push_str("## Limitations\n\n");
225            for lim in &self.limitations {
226                md.push_str(&format!("- {lim}\n"));
227            }
228            md.push('\n');
229        }
230
231        // ── Training Details ───────────────────────────────────────────────
232        md.push_str("## Training Details\n\n");
233        md.push_str(&format!(
234            "- **Framework:** {}\n",
235            self.training_info.framework
236        ));
237        if let Some(n) = self.training_info.num_parameters {
238            md.push_str(&format!("- **Parameters:** {n}\n"));
239        }
240        if !self.training_info.training_data.is_empty() {
241            md.push_str(&format!(
242                "- **Training Data:** {}\n",
243                self.training_info.training_data.join(", ")
244            ));
245        }
246        if let Some(ref opt) = self.training_info.optimizer {
247            md.push_str(&format!("- **Optimizer:** {opt}\n"));
248        }
249        if let Some(lr) = self.training_info.learning_rate {
250            md.push_str(&format!("- **Learning Rate:** {lr}\n"));
251        }
252        if let Some(bs) = self.training_info.batch_size {
253            md.push_str(&format!("- **Batch Size:** {bs}\n"));
254        }
255        if let Some(ep) = self.training_info.num_epochs {
256            md.push_str(&format!("- **Epochs:** {ep}\n"));
257        }
258        if let Some(ref hw) = self.training_info.hardware {
259            md.push_str(&format!("- **Hardware:** {hw}\n"));
260        }
261        md.push('\n');
262
263        // ── Benchmarks ─────────────────────────────────────────────────────
264        if !self.benchmarks.is_empty() {
265            md.push_str("## Evaluation Results\n\n");
266            md.push_str("| Task | Dataset | Metric | Value |\n");
267            md.push_str("|------|---------|--------|-------|\n");
268            for b in &self.benchmarks {
269                md.push_str(&format!(
270                    "| {} | {} | {} | {:.4} |\n",
271                    b.task, b.dataset, b.metric, b.value
272                ));
273            }
274            md.push('\n');
275        }
276
277        // ── Citation ───────────────────────────────────────────────────────
278        if let Some(ref citation) = self.citation {
279            md.push_str("## Citation\n\n");
280            md.push_str("```bibtex\n");
281            md.push_str(citation);
282            md.push_str("\n```\n\n");
283        }
284
285        // ── Footer ─────────────────────────────────────────────────────────
286        md.push_str("---\n");
287        md.push_str("*Generated by [TrustformeRS](https://github.com/cool-japan/trustformers)*\n");
288
289        md
290    }
291
292    /// Parse a model card from a markdown string.
293    ///
294    /// Reads the YAML front matter between the first pair of `---` delimiters
295    /// and extracts the `## Model Description` section body.
296    pub fn from_markdown(content: &str) -> Result<Self> {
297        let mut metadata = ModelCardMetadata::default();
298        let mut model_name = String::new();
299        let mut model_description = String::new();
300        let mut author: Option<String> = None;
301
302        // Extract YAML front matter
303        if content.starts_with("---") {
304            let rest = &content[3..];
305            if let Some(end) = rest.find("\n---") {
306                let yaml_str = &rest[..end];
307                // Parse just the fields we care about
308                for line in yaml_str.lines() {
309                    if let Some(val) = line.strip_prefix("license: ") {
310                        metadata.license = val.trim().to_string();
311                    } else if let Some(val) = line.strip_prefix("library_name: ") {
312                        metadata.library_name = val.trim().to_string();
313                    } else if let Some(val) = line.strip_prefix("model_type: ") {
314                        metadata.model_type = Some(val.trim().to_string());
315                    } else if let Some(val) = line.strip_prefix("pipeline_tag: ") {
316                        metadata.pipeline_tag = Some(val.trim().to_string());
317                    } else if line.trim_start().starts_with("- ") {
318                        // handled implicitly; full YAML parse is below
319                    }
320                }
321
322                // Use serde_yaml_ng for a richer parse of the front matter
323                let parsed: serde_yaml_ng::Value =
324                    serde_yaml_ng::from_str(yaml_str).unwrap_or(serde_yaml_ng::Value::Null);
325
326                if let serde_yaml_ng::Value::Mapping(ref map) = parsed {
327                    if let Some(serde_yaml_ng::Value::Sequence(seq)) = map.get("language") {
328                        metadata.language =
329                            seq.iter().filter_map(|x| x.as_str().map(String::from)).collect();
330                    }
331                    if let Some(serde_yaml_ng::Value::Sequence(seq)) = map.get("tags") {
332                        metadata.tags =
333                            seq.iter().filter_map(|x| x.as_str().map(String::from)).collect();
334                    }
335                    if let Some(serde_yaml_ng::Value::Sequence(seq)) = map.get("datasets") {
336                        metadata.datasets =
337                            seq.iter().filter_map(|x| x.as_str().map(String::from)).collect();
338                    }
339                    if let Some(serde_yaml_ng::Value::Sequence(seq)) = map.get("metrics") {
340                        metadata.metrics =
341                            seq.iter().filter_map(|x| x.as_str().map(String::from)).collect();
342                    }
343                }
344            }
345        }
346
347        // Extract title (first `# ` heading)
348        for line in content.lines() {
349            if let Some(name) = line.strip_prefix("# ") {
350                model_name = name.trim().to_string();
351                break;
352            }
353        }
354
355        // Extract author
356        for line in content.lines() {
357            if let Some(rest) = line.strip_prefix("*Author: ") {
358                author = Some(rest.trim_end_matches('*').to_string());
359                break;
360            }
361        }
362
363        // Extract ## Model Description section
364        let mut in_description = false;
365        for line in content.lines() {
366            if line.starts_with("## Model Description") {
367                in_description = true;
368                continue;
369            }
370            if in_description {
371                if line.starts_with("## ") {
372                    break;
373                }
374                if !model_description.is_empty() || !line.is_empty() {
375                    model_description.push_str(line);
376                    model_description.push('\n');
377                }
378            }
379        }
380        let model_description = model_description.trim().to_string();
381
382        if model_name.is_empty() {
383            return Err(TrustformersError::InvalidInput {
384                message: "Model card must contain a top-level heading (# Model Name)".to_string(),
385                parameter: Some("model_name".to_string()),
386                expected: Some("A line starting with '# '".to_string()),
387                received: None,
388                suggestion: Some("Add '# Your Model Name' to the markdown".to_string()),
389            });
390        }
391
392        debug!(
393            model_name = %model_name,
394            "Parsed model card from markdown"
395        );
396
397        Ok(Self {
398            metadata,
399            model_name,
400            model_description,
401            intended_uses: Vec::new(),
402            limitations: Vec::new(),
403            training_info: TrainingInfo::default(),
404            benchmarks: Vec::new(),
405            citation: None,
406            author,
407        })
408    }
409
410    /// Save the model card to a file as markdown
411    pub fn save(&self, path: &Path) -> Result<()> {
412        let content = self.to_markdown();
413        if let Some(parent) = path.parent() {
414            std::fs::create_dir_all(parent).map_err(|e| TrustformersError::Io {
415                message: format!("Cannot create parent directory: {e}"),
416                path: Some(parent.display().to_string()),
417                suggestion: None,
418            })?;
419        }
420        std::fs::write(path, content).map_err(|e| TrustformersError::Io {
421            message: format!("Cannot write model card: {e}"),
422            path: Some(path.display().to_string()),
423            suggestion: None,
424        })?;
425        Ok(())
426    }
427
428    /// Load a model card from a markdown file
429    pub fn load(path: &Path) -> Result<Self> {
430        let content = std::fs::read_to_string(path).map_err(|e| TrustformersError::Io {
431            message: format!("Cannot read model card file: {e}"),
432            path: Some(path.display().to_string()),
433            suggestion: Some("Ensure the file exists and is readable".to_string()),
434        })?;
435        Self::from_markdown(&content)
436    }
437
438    /// Add a benchmark result (fluent mutating method)
439    pub fn add_benchmark(&mut self, result: BenchmarkResult) -> &mut Self {
440        self.benchmarks.push(result);
441        self
442    }
443
444    /// Replace the training info (builder-style consuming method)
445    pub fn with_training_info(mut self, info: TrainingInfo) -> Self {
446        self.training_info = info;
447        self
448    }
449
450    /// Validate the model card for completeness.
451    ///
452    /// Returns a list of warning strings for missing or recommended fields.
453    pub fn validate(&self) -> Vec<String> {
454        let mut warnings = Vec::new();
455
456        if self.model_name.is_empty() {
457            warnings.push("model_name is empty".to_string());
458        }
459        if self.model_description.is_empty() {
460            warnings.push("model_description is empty — add a description".to_string());
461        }
462        if self.metadata.language.is_empty() {
463            warnings.push("metadata.language is empty — specify supported languages".to_string());
464        }
465        if self.metadata.license.is_empty() {
466            warnings.push("metadata.license is empty — specify a license".to_string());
467        }
468        if self.intended_uses.is_empty() {
469            warnings.push("intended_uses is empty — document intended uses".to_string());
470        }
471        if self.limitations.is_empty() {
472            warnings.push("limitations is empty — document model limitations".to_string());
473        }
474        if self.benchmarks.is_empty() {
475            warnings.push("benchmarks is empty — add evaluation results".to_string());
476        }
477        if self.metadata.pipeline_tag.is_none() {
478            warnings.push("metadata.pipeline_tag is not set".to_string());
479        }
480        if self.author.is_none() {
481            warnings.push("author is not set".to_string());
482        }
483
484        warnings
485    }
486}
487
488// ─── ModelCardGenerator ───────────────────────────────────────────────────────
489
490/// Utility for auto-generating model cards from structured metadata
491pub struct ModelCardGenerator;
492
493impl ModelCardGenerator {
494    /// Generate a model card from basic model info and training details
495    pub fn generate(
496        model_type: &str,
497        model_name: &str,
498        training_info: TrainingInfo,
499        pipeline_tag: Option<&str>,
500    ) -> ModelCard {
501        let mut metadata = ModelCardMetadata {
502            model_type: Some(model_type.to_string()),
503            pipeline_tag: pipeline_tag.map(String::from),
504            tags: vec![model_type.to_string()],
505            ..Default::default()
506        };
507
508        if let Some(pt) = pipeline_tag {
509            if !metadata.tags.contains(&pt.to_string()) {
510                metadata.tags.push(pt.to_string());
511            }
512        }
513
514        let description = format!(
515            "This is a {model_type} model trained with TrustformeRS. \
516             It was trained using the {} framework.",
517            training_info.framework
518        );
519
520        ModelCard {
521            metadata,
522            model_name: model_name.to_string(),
523            model_description: description,
524            intended_uses: vec![format!("This model can be used for {model_type} tasks.")],
525            limitations: vec![
526                "Model performance may degrade on out-of-distribution data.".to_string(),
527                "The model has not been evaluated for all use cases.".to_string(),
528            ],
529            training_info,
530            benchmarks: Vec::new(),
531            citation: None,
532            author: None,
533        }
534    }
535
536    /// Generate a model card with pre-populated benchmark results
537    pub fn generate_with_benchmarks(
538        model_type: &str,
539        model_name: &str,
540        training_info: TrainingInfo,
541        benchmarks: Vec<BenchmarkResult>,
542    ) -> ModelCard {
543        let mut card = Self::generate(model_type, model_name, training_info, None);
544        card.benchmarks = benchmarks;
545
546        // Auto-populate metrics in metadata from benchmarks
547        card.metadata.metrics = card
548            .benchmarks
549            .iter()
550            .map(|b| b.metric.clone())
551            .collect::<std::collections::HashSet<_>>()
552            .into_iter()
553            .collect();
554
555        card
556    }
557}
558
559// ─── ModelCard additional methods ─────────────────────────────────────────────
560
561impl ModelCard {
562    /// Render only the YAML front-matter block (between the `---` delimiters).
563    pub fn to_yaml_frontmatter(&self) -> String {
564        let mut yaml = String::from("---\n");
565        if !self.metadata.language.is_empty() {
566            yaml.push_str("language:\n");
567            for lang in &self.metadata.language {
568                yaml.push_str(&format!("- {lang}\n"));
569            }
570        }
571        yaml.push_str(&format!("license: {}\n", self.metadata.license));
572        yaml.push_str(&format!("library_name: {}\n", self.metadata.library_name));
573        if !self.metadata.tags.is_empty() {
574            yaml.push_str("tags:\n");
575            for tag in &self.metadata.tags {
576                yaml.push_str(&format!("- {tag}\n"));
577            }
578        }
579        if !self.metadata.datasets.is_empty() {
580            yaml.push_str("datasets:\n");
581            for ds in &self.metadata.datasets {
582                yaml.push_str(&format!("- {ds}\n"));
583            }
584        }
585        if !self.metadata.metrics.is_empty() {
586            yaml.push_str("metrics:\n");
587            for m in &self.metadata.metrics {
588                yaml.push_str(&format!("- {m}\n"));
589            }
590        }
591        if let Some(ref mt) = self.metadata.model_type {
592            yaml.push_str(&format!("model_type: {mt}\n"));
593        }
594        if let Some(ref pt) = self.metadata.pipeline_tag {
595            yaml.push_str(&format!("pipeline_tag: {pt}\n"));
596        }
597        if let Some(n) = self.training_info.num_parameters {
598            yaml.push_str(&format!("num_parameters: {n}\n"));
599        }
600        yaml.push_str("---\n");
601        yaml
602    }
603}
604
605// ─── ModelCardError ───────────────────────────────────────────────────────────
606
607/// Errors that can occur when building or parsing a model card.
608#[derive(Debug, Clone, PartialEq, Eq)]
609pub enum ModelCardError {
610    /// A required field was not provided.
611    MissingField(String),
612    /// A field value was invalid.
613    InvalidField { field: String, reason: String },
614    /// The markdown could not be parsed.
615    ParseError(String),
616}
617
618impl std::fmt::Display for ModelCardError {
619    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
620        match self {
621            ModelCardError::MissingField(name) => write!(f, "Missing required field: {name}"),
622            ModelCardError::InvalidField { field, reason } => {
623                write!(f, "Invalid field '{field}': {reason}")
624            },
625            ModelCardError::ParseError(msg) => write!(f, "Parse error: {msg}"),
626        }
627    }
628}
629
630impl std::error::Error for ModelCardError {}
631
632// ─── ModelCardBuilder ─────────────────────────────────────────────────────────
633
634/// Builder for constructing a [`ModelCard`] with a fluent interface.
635///
636/// Every setter returns `Self` for chaining.  Call [`build`][ModelCardBuilder::build]
637/// when all required fields are provided.
638#[derive(Debug, Clone, Default)]
639pub struct ModelCardBuilder {
640    model_id: Option<String>,
641    architecture: Option<String>,
642    parameters: Option<u64>,
643    languages: Vec<String>,
644    license: Option<String>,
645    datasets: Vec<String>,
646    metrics: Vec<(String, f32)>,
647    tags: Vec<String>,
648    limitations: Option<String>,
649    bias_risks: Option<String>,
650    description: Option<String>,
651    pipeline_tag: Option<String>,
652    author: Option<String>,
653}
654
655impl ModelCardBuilder {
656    /// Start a new builder.
657    pub fn new() -> Self {
658        Self::default()
659    }
660
661    /// Set the model identifier (used as the model name).
662    pub fn with_model_id(mut self, id: &str) -> Self {
663        self.model_id = Some(id.to_string());
664        self
665    }
666
667    /// Set the model architecture (e.g. "bert", "llama").
668    pub fn with_architecture(mut self, arch: &str) -> Self {
669        self.architecture = Some(arch.to_string());
670        self
671    }
672
673    /// Set the approximate number of trainable parameters.
674    pub fn with_parameters(mut self, params: u64) -> Self {
675        self.parameters = Some(params);
676        self
677    }
678
679    /// Add a supported language (BCP-47 code, e.g. "en").
680    /// Multiple calls accumulate languages.
681    pub fn with_language(mut self, lang: &str) -> Self {
682        self.languages.push(lang.to_string());
683        self
684    }
685
686    /// Set the SPDX license identifier (e.g. "apache-2.0").
687    pub fn with_license(mut self, license: &str) -> Self {
688        self.license = Some(license.to_string());
689        self
690    }
691
692    /// Add a training dataset name.
693    pub fn with_dataset(mut self, dataset: &str) -> Self {
694        self.datasets.push(dataset.to_string());
695        self
696    }
697
698    /// Add an evaluation metric name and its numeric value.
699    pub fn with_metrics(mut self, name: &str, value: f32) -> Self {
700        self.metrics.push((name.to_string(), value));
701        self
702    }
703
704    /// Replace the tag list.
705    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
706        self.tags = tags;
707        self
708    }
709
710    /// Set the limitations section text.
711    pub fn with_limitations(mut self, text: &str) -> Self {
712        self.limitations = Some(text.to_string());
713        self
714    }
715
716    /// Set the bias & risks section text.
717    pub fn with_bias_risks(mut self, text: &str) -> Self {
718        self.bias_risks = Some(text.to_string());
719        self
720    }
721
722    /// Set a short description of the model.
723    pub fn with_description(mut self, desc: &str) -> Self {
724        self.description = Some(desc.to_string());
725        self
726    }
727
728    /// Set the HuggingFace pipeline tag.
729    pub fn with_pipeline_tag(mut self, tag: &str) -> Self {
730        self.pipeline_tag = Some(tag.to_string());
731        self
732    }
733
734    /// Set the author name.
735    pub fn with_author(mut self, author: &str) -> Self {
736        self.author = Some(author.to_string());
737        self
738    }
739
740    /// Build the [`ModelCard`].
741    ///
742    /// Returns `Err` if required fields (`model_id`) are missing.
743    pub fn build(self) -> std::result::Result<ModelCard, ModelCardError> {
744        let model_id = self
745            .model_id
746            .ok_or_else(|| ModelCardError::MissingField("model_id".to_string()))?;
747
748        if model_id.is_empty() {
749            return Err(ModelCardError::InvalidField {
750                field: "model_id".to_string(),
751                reason: "must not be empty".to_string(),
752            });
753        }
754
755        // Resolve metadata.
756        let metadata = ModelCardMetadata {
757            language: if self.languages.is_empty() {
758                vec!["en".to_string()]
759            } else {
760                self.languages
761            },
762            license: self.license.unwrap_or_else(|| "apache-2.0".to_string()),
763            library_name: "trustformers".to_string(),
764            tags: {
765                let mut tags = self.tags;
766                if let Some(ref arch) = self.architecture {
767                    if !tags.contains(arch) {
768                        tags.push(arch.clone());
769                    }
770                }
771                tags
772            },
773            datasets: self.datasets,
774            metrics: self.metrics.iter().map(|(n, _)| n.clone()).collect(),
775            model_type: self.architecture.clone(),
776            pipeline_tag: self.pipeline_tag,
777        };
778
779        let description = self.description.unwrap_or_else(|| {
780            self.architecture
781                .as_deref()
782                .map(|a| format!("A {a} model trained with TrustformeRS."))
783                .unwrap_or_else(|| "A model trained with TrustformeRS.".to_string())
784        });
785
786        // Build limitations list.
787        let mut limitations = Vec::new();
788        if let Some(lim_text) = self.limitations {
789            limitations.push(lim_text);
790        }
791        if let Some(bias_text) = self.bias_risks {
792            limitations.push(format!("[Bias/Risks] {bias_text}"));
793        }
794        if limitations.is_empty() {
795            limitations
796                .push("Model performance may degrade on out-of-distribution data.".to_string());
797        }
798
799        // Training info.
800        let training_info = TrainingInfo {
801            num_parameters: self.parameters,
802            ..Default::default()
803        };
804
805        // Add metrics as benchmarks.
806        let benchmarks: Vec<BenchmarkResult> = self
807            .metrics
808            .iter()
809            .map(|(name, value)| {
810                let value = *value;
811                BenchmarkResult::new("evaluation", "unknown", name.as_str(), value as f64)
812            })
813            .collect();
814
815        Ok(ModelCard {
816            metadata,
817            model_name: model_id,
818            model_description: description,
819            intended_uses: Vec::new(),
820            limitations,
821            training_info,
822            benchmarks,
823            citation: None,
824            author: self.author,
825        })
826    }
827}
828
829// ─── ModelCardTemplate ────────────────────────────────────────────────────────
830
831/// Pre-defined card templates for common model types.
832#[derive(Debug, Clone, Copy, PartialEq, Eq)]
833pub enum ModelCardTemplate {
834    /// General-purpose language model.
835    General,
836    /// Classification model (sequence-level label).
837    Classification,
838    /// Text-generation / language-generation model.
839    Generation,
840    /// Multilingual model supporting several languages.
841    Multilingual,
842}
843
844impl ModelCardTemplate {
845    /// Apply the template, returning a pre-populated [`ModelCardBuilder`].
846    pub fn apply(self, model_id: &str) -> ModelCardBuilder {
847        match self {
848            ModelCardTemplate::General => ModelCardBuilder::new()
849                .with_model_id(model_id)
850                .with_language("en")
851                .with_license("apache-2.0")
852                .with_tags(vec![
853                    "trustformers".to_string(),
854                    "transformer".to_string(),
855                ])
856                .with_limitations(
857                    "Performance may vary on out-of-distribution inputs.",
858                )
859                .with_bias_risks(
860                    "The model may reflect biases present in its training data.",
861                ),
862
863            ModelCardTemplate::Classification => ModelCardBuilder::new()
864                .with_model_id(model_id)
865                .with_language("en")
866                .with_license("apache-2.0")
867                .with_pipeline_tag("text-classification")
868                .with_tags(vec![
869                    "text-classification".to_string(),
870                    "trustformers".to_string(),
871                ])
872                .with_limitations(
873                    "Classification accuracy may degrade on out-of-distribution data.",
874                )
875                .with_bias_risks(
876                    "Classifier may exhibit label bias if training data is imbalanced.",
877                ),
878
879            ModelCardTemplate::Generation => ModelCardBuilder::new()
880                .with_model_id(model_id)
881                .with_language("en")
882                .with_license("apache-2.0")
883                .with_pipeline_tag("text-generation")
884                .with_tags(vec![
885                    "text-generation".to_string(),
886                    "causal-lm".to_string(),
887                    "trustformers".to_string(),
888                ])
889                .with_limitations(
890                    "Generated text may be factually incorrect or harmful.",
891                )
892                .with_bias_risks(
893                    "The model may generate biased, offensive, or misleading content.",
894                ),
895
896            ModelCardTemplate::Multilingual => ModelCardBuilder::new()
897                .with_model_id(model_id)
898                .with_language("en")
899                .with_language("fr")
900                .with_language("de")
901                .with_language("es")
902                .with_language("zh")
903                .with_license("apache-2.0")
904                .with_tags(vec![
905                    "multilingual".to_string(),
906                    "trustformers".to_string(),
907                ])
908                .with_limitations(
909                    "Performance varies across languages; low-resource languages may perform worse.",
910                )
911                .with_bias_risks(
912                    "Multilingual models can exhibit cross-lingual bias.",
913                ),
914        }
915    }
916
917    /// Return a string label for this template.
918    pub fn label(self) -> &'static str {
919        match self {
920            ModelCardTemplate::General => "general",
921            ModelCardTemplate::Classification => "classification",
922            ModelCardTemplate::Generation => "generation",
923            ModelCardTemplate::Multilingual => "multilingual",
924        }
925    }
926}
927
928// ─── Tests ────────────────────────────────────────────────────────────────────
929
930#[cfg(test)]
931mod tests {
932    use super::*;
933    use std::path::PathBuf;
934
935    fn temp_dir() -> PathBuf {
936        std::env::temp_dir()
937    }
938
939    #[test]
940    fn test_model_card_new() {
941        let card = ModelCard::new("MyBERT", "A BERT-based classification model");
942        assert_eq!(card.model_name, "MyBERT");
943        assert!(!card.model_description.is_empty());
944    }
945
946    #[test]
947    fn test_model_card_metadata_default() {
948        let meta = ModelCardMetadata::default();
949        assert_eq!(meta.license, "apache-2.0");
950        assert_eq!(meta.library_name, "trustformers");
951        assert!(!meta.language.is_empty());
952    }
953
954    #[test]
955    fn test_to_markdown_contains_title() {
956        let card = ModelCard::new("TestModel", "Test description");
957        let md = card.to_markdown();
958        assert!(md.contains("# TestModel"));
959    }
960
961    #[test]
962    fn test_to_markdown_contains_yaml_front_matter() {
963        let card = ModelCard::new("TestModel", "desc");
964        let md = card.to_markdown();
965        assert!(md.starts_with("---\n"));
966        assert!(md.contains("license: apache-2.0"));
967        assert!(md.contains("library_name: trustformers"));
968    }
969
970    #[test]
971    fn test_to_markdown_with_benchmarks() {
972        let mut card = ModelCard::new("BenchModel", "desc");
973        card.add_benchmark(BenchmarkResult::new(
974            "text-classification",
975            "glue/sst2",
976            "accuracy",
977            0.9234,
978        ));
979        let md = card.to_markdown();
980        assert!(md.contains("## Evaluation Results"));
981        assert!(md.contains("0.9234"));
982    }
983
984    #[test]
985    fn test_from_markdown_roundtrip() {
986        let mut original = ModelCard::new("RoundtripModel", "A test model for roundtrip testing");
987        original.author = Some("Test Author".to_string());
988        original.metadata.model_type = Some("bert".to_string());
989        original.metadata.pipeline_tag = Some("text-classification".to_string());
990
991        let md = original.to_markdown();
992        let parsed = ModelCard::from_markdown(&md).unwrap();
993
994        assert_eq!(parsed.model_name, "RoundtripModel");
995        assert!(parsed.model_description.contains("roundtrip testing"));
996        assert_eq!(parsed.metadata.model_type, Some("bert".to_string()));
997        assert_eq!(
998            parsed.metadata.pipeline_tag,
999            Some("text-classification".to_string())
1000        );
1001    }
1002
1003    #[test]
1004    fn test_from_markdown_missing_title_error() {
1005        let md = "No title here\nSome content";
1006        assert!(ModelCard::from_markdown(md).is_err());
1007    }
1008
1009    #[test]
1010    fn test_save_and_load() {
1011        let dir = temp_dir().join("trustformers_model_card_test");
1012        std::fs::create_dir_all(&dir).unwrap();
1013        let path = dir.join("README.md");
1014
1015        let card = ModelCard::new("SaveLoadModel", "Testing save/load functionality");
1016        card.save(&path).unwrap();
1017
1018        let loaded = ModelCard::load(&path).unwrap();
1019        assert_eq!(loaded.model_name, "SaveLoadModel");
1020
1021        std::fs::remove_dir_all(&dir).ok();
1022    }
1023
1024    #[test]
1025    fn test_add_benchmark_chaining() {
1026        let mut card = ModelCard::new("BenchModel", "desc");
1027        card.add_benchmark(BenchmarkResult::new("task1", "ds1", "f1", 0.85))
1028            .add_benchmark(BenchmarkResult::new("task2", "ds2", "acc", 0.92));
1029        assert_eq!(card.benchmarks.len(), 2);
1030    }
1031
1032    #[test]
1033    fn test_with_training_info() {
1034        let info = TrainingInfo {
1035            num_parameters: Some(110_000_000),
1036            optimizer: Some("AdamW".to_string()),
1037            learning_rate: Some(2e-5),
1038            num_epochs: Some(3),
1039            hardware: Some("1x A100".to_string()),
1040            ..Default::default()
1041        };
1042        let card = ModelCard::new("TrainModel", "desc").with_training_info(info);
1043        assert_eq!(card.training_info.num_parameters, Some(110_000_000));
1044        assert_eq!(card.training_info.optimizer, Some("AdamW".to_string()));
1045    }
1046
1047    #[test]
1048    fn test_validate_warns_on_empty_card() {
1049        let card = ModelCard::new("", "");
1050        let warnings = card.validate();
1051        assert!(!warnings.is_empty());
1052        // Should warn about model_name, description, uses, limitations, benchmarks
1053        assert!(warnings.iter().any(|w| w.contains("model_name")));
1054        assert!(warnings.iter().any(|w| w.contains("model_description")));
1055    }
1056
1057    #[test]
1058    fn test_validate_clean_card() {
1059        let mut card = ModelCard::new("CleanModel", "A clean well-documented model");
1060        card.author = Some("Author".to_string());
1061        card.intended_uses = vec!["classification".to_string()];
1062        card.limitations = vec!["limited training data".to_string()];
1063        card.metadata.pipeline_tag = Some("text-classification".to_string());
1064        card.add_benchmark(BenchmarkResult::new("tc", "sst2", "acc", 0.9));
1065
1066        let warnings = card.validate();
1067        // model_name, description, language, license should all be fine
1068        assert!(!warnings.iter().any(|w| w.contains("model_name")));
1069        assert!(!warnings.iter().any(|w| w.contains("model_description")));
1070        assert!(!warnings.iter().any(|w| w.contains("pipeline_tag")));
1071    }
1072
1073    #[test]
1074    fn test_generator_generate() {
1075        let info = TrainingInfo {
1076            num_epochs: Some(5),
1077            ..Default::default()
1078        };
1079        let card = ModelCardGenerator::generate("gpt2", "MyGPT", info, Some("text-generation"));
1080        assert_eq!(card.model_name, "MyGPT");
1081        assert_eq!(card.metadata.model_type, Some("gpt2".to_string()));
1082        assert_eq!(
1083            card.metadata.pipeline_tag,
1084            Some("text-generation".to_string())
1085        );
1086        assert!(!card.intended_uses.is_empty());
1087        assert!(!card.limitations.is_empty());
1088    }
1089
1090    #[test]
1091    fn test_generator_with_benchmarks() {
1092        let benchmarks = vec![
1093            BenchmarkResult::new("lm", "wikitext", "perplexity", 15.3),
1094            BenchmarkResult::new("lm", "ptb", "perplexity", 22.1),
1095        ];
1096        let card = ModelCardGenerator::generate_with_benchmarks(
1097            "gpt2",
1098            "BenchGPT",
1099            TrainingInfo::default(),
1100            benchmarks,
1101        );
1102        assert_eq!(card.benchmarks.len(), 2);
1103        assert!(!card.metadata.metrics.is_empty());
1104    }
1105
1106    #[test]
1107    fn test_markdown_includes_training_details() {
1108        let info = TrainingInfo {
1109            num_parameters: Some(340_000_000),
1110            optimizer: Some("Adam".to_string()),
1111            learning_rate: Some(1e-4),
1112            batch_size: Some(32),
1113            num_epochs: Some(10),
1114            hardware: Some("8x V100".to_string()),
1115            ..Default::default()
1116        };
1117        let card = ModelCard::new("DetailedModel", "desc").with_training_info(info);
1118        let md = card.to_markdown();
1119        assert!(md.contains("340000000"));
1120        assert!(md.contains("Adam"));
1121        assert!(md.contains("0.0001"));
1122        assert!(md.contains("8x V100"));
1123    }
1124
1125    #[test]
1126    fn test_benchmark_result_new() {
1127        let b = BenchmarkResult::new("ner", "conll2003", "f1", 0.93);
1128        assert_eq!(b.task, "ner");
1129        assert_eq!(b.dataset, "conll2003");
1130        assert_eq!(b.metric, "f1");
1131        assert!((b.value - 0.93).abs() < 1e-9);
1132    }
1133
1134    // ── ModelCardBuilder tests ────────────────────────────────────────────────
1135
1136    #[test]
1137    fn test_builder_minimal() {
1138        let card = ModelCardBuilder::new().with_model_id("my-model").build().unwrap();
1139        assert_eq!(card.model_name, "my-model");
1140        assert_eq!(card.metadata.license, "apache-2.0");
1141        assert!(!card.metadata.language.is_empty());
1142    }
1143
1144    #[test]
1145    fn test_builder_missing_model_id() {
1146        let err = ModelCardBuilder::new().build().unwrap_err();
1147        assert!(matches!(err, ModelCardError::MissingField(_)));
1148    }
1149
1150    #[test]
1151    fn test_builder_empty_model_id() {
1152        let err = ModelCardBuilder::new().with_model_id("").build().unwrap_err();
1153        assert!(matches!(err, ModelCardError::InvalidField { .. }));
1154    }
1155
1156    #[test]
1157    fn test_builder_full() {
1158        let card = ModelCardBuilder::new()
1159            .with_model_id("my-bert")
1160            .with_architecture("bert")
1161            .with_parameters(110_000_000)
1162            .with_language("en")
1163            .with_language("de")
1164            .with_license("mit")
1165            .with_dataset("glue")
1166            .with_metrics("accuracy", 0.94)
1167            .with_metrics("f1", 0.91)
1168            .with_tags(vec!["nlp".to_string(), "bert".to_string()])
1169            .with_limitations("Only handles English and German well.")
1170            .with_bias_risks("May reflect biases in training data.")
1171            .with_description("A BERT model fine-tuned for NER.")
1172            .with_author("COOLJAPAN")
1173            .build()
1174            .unwrap();
1175
1176        assert_eq!(card.model_name, "my-bert");
1177        assert_eq!(card.training_info.num_parameters, Some(110_000_000));
1178        assert!(card.metadata.language.contains(&"en".to_string()));
1179        assert!(card.metadata.language.contains(&"de".to_string()));
1180        assert_eq!(card.metadata.license, "mit");
1181        assert!(card.metadata.datasets.contains(&"glue".to_string()));
1182        assert_eq!(card.benchmarks.len(), 2);
1183        assert!(card.author.as_deref() == Some("COOLJAPAN"));
1184        assert!(!card.limitations.is_empty());
1185    }
1186
1187    #[test]
1188    fn test_builder_arch_tag_added_automatically() {
1189        let card = ModelCardBuilder::new()
1190            .with_model_id("llama-model")
1191            .with_architecture("llama")
1192            .build()
1193            .unwrap();
1194        assert!(card.metadata.tags.contains(&"llama".to_string()));
1195        assert_eq!(card.metadata.model_type, Some("llama".to_string()));
1196    }
1197
1198    #[test]
1199    fn test_builder_default_language_fallback() {
1200        let card = ModelCardBuilder::new().with_model_id("no-lang-model").build().unwrap();
1201        // Should default to ["en"].
1202        assert_eq!(card.metadata.language, vec!["en"]);
1203    }
1204
1205    #[test]
1206    fn test_builder_multiple_languages() {
1207        let card = ModelCardBuilder::new()
1208            .with_model_id("multi-lang")
1209            .with_language("en")
1210            .with_language("fr")
1211            .with_language("es")
1212            .build()
1213            .unwrap();
1214        assert_eq!(card.metadata.language.len(), 3);
1215        assert!(card.metadata.language.contains(&"fr".to_string()));
1216    }
1217
1218    #[test]
1219    fn test_builder_metrics_populate_benchmarks_and_metadata() {
1220        let card = ModelCardBuilder::new()
1221            .with_model_id("bench-model")
1222            .with_metrics("perplexity", 12.5)
1223            .with_metrics("bleu", 0.45)
1224            .build()
1225            .unwrap();
1226        assert_eq!(card.benchmarks.len(), 2);
1227        assert!(card.metadata.metrics.contains(&"perplexity".to_string()));
1228        assert!(card.metadata.metrics.contains(&"bleu".to_string()));
1229    }
1230
1231    // ── to_yaml_frontmatter tests ─────────────────────────────────────────────
1232
1233    #[test]
1234    fn test_to_yaml_frontmatter_structure() {
1235        let card = ModelCard::new("TestModel", "desc");
1236        let yaml = card.to_yaml_frontmatter();
1237        assert!(yaml.starts_with("---\n"));
1238        assert!(yaml.ends_with("---\n"));
1239        assert!(yaml.contains("license: apache-2.0"));
1240        assert!(yaml.contains("library_name: trustformers"));
1241    }
1242
1243    #[test]
1244    fn test_to_yaml_frontmatter_with_model_type() {
1245        let mut card = ModelCard::new("BertModel", "desc");
1246        card.metadata.model_type = Some("bert".to_string());
1247        card.metadata.pipeline_tag = Some("text-classification".to_string());
1248        let yaml = card.to_yaml_frontmatter();
1249        assert!(yaml.contains("model_type: bert"));
1250        assert!(yaml.contains("pipeline_tag: text-classification"));
1251    }
1252
1253    #[test]
1254    fn test_to_yaml_frontmatter_languages() {
1255        let mut card = ModelCard::new("MultiLang", "desc");
1256        card.metadata.language = vec!["en".to_string(), "fr".to_string()];
1257        let yaml = card.to_yaml_frontmatter();
1258        assert!(yaml.contains("language:"));
1259        assert!(yaml.contains("- en"));
1260        assert!(yaml.contains("- fr"));
1261    }
1262
1263    #[test]
1264    fn test_to_yaml_frontmatter_num_parameters() {
1265        let mut card = ModelCard::new("BigModel", "desc");
1266        card.training_info.num_parameters = Some(7_000_000_000);
1267        let yaml = card.to_yaml_frontmatter();
1268        assert!(yaml.contains("num_parameters: 7000000000"));
1269    }
1270
1271    // ── ModelCardTemplate tests ───────────────────────────────────────────────
1272
1273    #[test]
1274    fn test_template_general() {
1275        let card = ModelCardTemplate::General.apply("my-general-model").build().unwrap();
1276        assert_eq!(card.model_name, "my-general-model");
1277        assert!(!card.limitations.is_empty());
1278    }
1279
1280    #[test]
1281    fn test_template_classification() {
1282        let card = ModelCardTemplate::Classification.apply("my-classifier").build().unwrap();
1283        assert_eq!(
1284            card.metadata.pipeline_tag,
1285            Some("text-classification".to_string())
1286        );
1287        assert!(card.metadata.tags.contains(&"text-classification".to_string()));
1288    }
1289
1290    #[test]
1291    fn test_template_generation() {
1292        let card = ModelCardTemplate::Generation.apply("my-gpt").build().unwrap();
1293        assert_eq!(
1294            card.metadata.pipeline_tag,
1295            Some("text-generation".to_string())
1296        );
1297        assert!(card.metadata.tags.contains(&"causal-lm".to_string()));
1298    }
1299
1300    #[test]
1301    fn test_template_multilingual() {
1302        let card = ModelCardTemplate::Multilingual.apply("my-multi").build().unwrap();
1303        assert!(card.metadata.language.len() >= 5);
1304        assert!(card.metadata.language.contains(&"zh".to_string()));
1305    }
1306
1307    #[test]
1308    fn test_template_labels() {
1309        assert_eq!(ModelCardTemplate::General.label(), "general");
1310        assert_eq!(ModelCardTemplate::Classification.label(), "classification");
1311        assert_eq!(ModelCardTemplate::Generation.label(), "generation");
1312        assert_eq!(ModelCardTemplate::Multilingual.label(), "multilingual");
1313    }
1314
1315    #[test]
1316    fn test_template_further_customisation() {
1317        let card = ModelCardTemplate::Generation
1318            .apply("base-model")
1319            .with_parameters(1_000_000_000)
1320            .with_dataset("c4")
1321            .with_author("COOLJAPAN")
1322            .build()
1323            .unwrap();
1324        assert_eq!(card.training_info.num_parameters, Some(1_000_000_000));
1325        assert!(card.metadata.datasets.contains(&"c4".to_string()));
1326        assert_eq!(card.author.as_deref(), Some("COOLJAPAN"));
1327    }
1328
1329    #[test]
1330    fn test_from_markdown_roundtrip_builder() {
1331        let original = ModelCardBuilder::new()
1332            .with_model_id("RoundtripBuilt")
1333            .with_architecture("roberta")
1334            .with_language("en")
1335            .with_description("A RoBERTa model built with the builder.")
1336            .build()
1337            .unwrap();
1338
1339        let md = original.to_markdown();
1340        let parsed = ModelCard::from_markdown(&md).unwrap();
1341        assert_eq!(parsed.model_name, "RoundtripBuilt");
1342        assert!(parsed.model_description.contains("RoBERTa"));
1343    }
1344
1345    // ── ModelCardError tests ──────────────────────────────────────────────────
1346
1347    #[test]
1348    fn test_model_card_error_display_missing_field() {
1349        let err = ModelCardError::MissingField("model_id".to_string());
1350        assert!(err.to_string().contains("model_id"));
1351    }
1352
1353    #[test]
1354    fn test_model_card_error_display_invalid_field() {
1355        let err = ModelCardError::InvalidField {
1356            field: "license".to_string(),
1357            reason: "unknown identifier".to_string(),
1358        };
1359        assert!(err.to_string().contains("license"));
1360        assert!(err.to_string().contains("unknown identifier"));
1361    }
1362
1363    #[test]
1364    fn test_model_card_error_display_parse_error() {
1365        let err = ModelCardError::ParseError("unexpected end of YAML".to_string());
1366        assert!(err.to_string().contains("Parse error"));
1367        assert!(err.to_string().contains("unexpected end"));
1368    }
1369}
trustformers/hub_model_card.rs

trustformers/
hub_model_card.rs
