914 lines
30 KiB
Rust
914 lines
30 KiB
Rust
//! HuggingFace-style Model Cards for model documentation and metadata.
|
|
|
|
use crate::{HubError, HubResult, ModelId};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use std::collections::HashMap;
|
|
|
|
/// Demo widget example for model card.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct Widget {
|
|
/// Example input text
|
|
pub text: String,
|
|
/// Example parameters
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub parameters: Option<HashMap<String, Value>>,
|
|
/// Output example
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub output: Option<Value>,
|
|
}
|
|
|
|
impl Widget {
|
|
/// Create a new widget with text.
|
|
pub fn new(text: impl Into<String>) -> Self {
|
|
Self {
|
|
text: text.into(),
|
|
parameters: None,
|
|
output: None,
|
|
}
|
|
}
|
|
|
|
/// Add parameters to the widget.
|
|
pub fn with_parameters(mut self, parameters: HashMap<String, Value>) -> Self {
|
|
self.parameters = Some(parameters);
|
|
self
|
|
}
|
|
|
|
/// Add output to the widget.
|
|
pub fn with_output(mut self, output: Value) -> Self {
|
|
self.output = Some(output);
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Evaluation metric result.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct MetricResult {
|
|
/// Metric name (e.g., "accuracy", "f1", "perplexity")
|
|
pub name: String,
|
|
/// Metric type (e.g., "accuracy", "f1_score")
|
|
#[serde(rename = "type")]
|
|
pub metric_type: String,
|
|
/// Metric value
|
|
pub value: f64,
|
|
/// Dataset used for evaluation
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub dataset: Option<String>,
|
|
/// Split used (e.g., "test", "validation")
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub split: Option<String>,
|
|
/// Additional metadata
|
|
#[serde(skip_serializing_if = "HashMap::is_empty", default)]
|
|
pub metadata: HashMap<String, Value>,
|
|
}
|
|
|
|
impl MetricResult {
|
|
/// Create a new metric result.
|
|
pub fn new(name: impl Into<String>, metric_type: impl Into<String>, value: f64) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
metric_type: metric_type.into(),
|
|
value,
|
|
dataset: None,
|
|
split: None,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Set the dataset.
|
|
pub fn with_dataset(mut self, dataset: impl Into<String>) -> Self {
|
|
self.dataset = Some(dataset.into());
|
|
self
|
|
}
|
|
|
|
/// Set the split.
|
|
pub fn with_split(mut self, split: impl Into<String>) -> Self {
|
|
self.split = Some(split.into());
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Model benchmark result for model index.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ModelIndex {
|
|
/// Model name
|
|
pub name: String,
|
|
/// Benchmark results
|
|
pub results: Vec<MetricResult>,
|
|
}
|
|
|
|
impl ModelIndex {
|
|
/// Create a new model index.
|
|
pub fn new(name: impl Into<String>) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
results: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Add a metric result.
|
|
pub fn add_result(mut self, result: MetricResult) -> Self {
|
|
self.results.push(result);
|
|
self
|
|
}
|
|
}
|
|
|
|
/// CO2 emissions information.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct CO2Emissions {
|
|
/// Emissions in grams of CO2
|
|
pub emissions: f64,
|
|
/// Source of the calculation
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub source: Option<String>,
|
|
/// Training type (e.g., "pretraining", "fine-tuning")
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub training_type: Option<String>,
|
|
/// Geographical location
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub geographical_location: Option<String>,
|
|
/// Hardware used
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub hardware_used: Option<String>,
|
|
}
|
|
|
|
impl CO2Emissions {
|
|
/// Create new CO2 emissions info.
|
|
pub fn new(emissions: f64) -> Self {
|
|
Self {
|
|
emissions,
|
|
source: None,
|
|
training_type: None,
|
|
geographical_location: None,
|
|
hardware_used: None,
|
|
}
|
|
}
|
|
|
|
/// Set the source.
|
|
pub fn with_source(mut self, source: impl Into<String>) -> Self {
|
|
self.source = Some(source.into());
|
|
self
|
|
}
|
|
|
|
/// Set the training type.
|
|
pub fn with_training_type(mut self, training_type: impl Into<String>) -> Self {
|
|
self.training_type = Some(training_type.into());
|
|
self
|
|
}
|
|
}
|
|
|
|
/// HuggingFace-style Model Card documentation.
|
|
///
|
|
/// This is a comprehensive model card for documentation purposes, similar to
|
|
/// HuggingFace's model cards with YAML front matter. It is distinct from the
|
|
/// simpler `ModelCard` in the discovery module.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ModelCardDoc {
|
|
/// Model identifier
|
|
pub model_id: ModelId,
|
|
/// Languages supported (ISO 639-1 codes)
|
|
#[serde(skip_serializing_if = "Vec::is_empty", default)]
|
|
pub language: Vec<String>,
|
|
/// License identifier
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub license: Option<String>,
|
|
/// Library name (e.g., "rustytorch", "transformers")
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub library_name: Option<String>,
|
|
/// Model tags for categorization
|
|
#[serde(skip_serializing_if = "Vec::is_empty", default)]
|
|
pub tags: Vec<String>,
|
|
/// Pipeline tag (e.g., "text-generation", "image-classification")
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub pipeline_tag: Option<String>,
|
|
/// Demo widget examples
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub widget: Option<Vec<Widget>>,
|
|
/// Training datasets
|
|
#[serde(skip_serializing_if = "Vec::is_empty", default)]
|
|
pub datasets: Vec<String>,
|
|
/// Evaluation metrics
|
|
#[serde(skip_serializing_if = "Vec::is_empty", default)]
|
|
pub metrics: Vec<MetricResult>,
|
|
/// Model benchmark index
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub model_index: Option<Vec<ModelIndex>>,
|
|
/// CO2 emissions information
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub co2_eq_emissions: Option<CO2Emissions>,
|
|
/// Custom YAML front matter data
|
|
#[serde(skip_serializing_if = "HashMap::is_empty", default)]
|
|
pub card_data: HashMap<String, Value>,
|
|
/// Markdown content
|
|
pub content: String,
|
|
}
|
|
|
|
impl ModelCardDoc {
|
|
/// Create a new model card documentation.
|
|
pub fn new(model_id: ModelId, content: impl Into<String>) -> Self {
|
|
Self {
|
|
model_id,
|
|
language: Vec::new(),
|
|
license: None,
|
|
library_name: Some("rustytorch".to_string()),
|
|
tags: Vec::new(),
|
|
pipeline_tag: None,
|
|
widget: None,
|
|
datasets: Vec::new(),
|
|
metrics: Vec::new(),
|
|
model_index: None,
|
|
co2_eq_emissions: None,
|
|
card_data: HashMap::new(),
|
|
content: content.into(),
|
|
}
|
|
}
|
|
|
|
/// Parse a model card documentation from README.md format with YAML front matter.
|
|
pub fn from_readme(model_id: ModelId, readme_content: &str) -> HubResult<Self> {
|
|
let readme_content = readme_content.trim();
|
|
|
|
// Check if the content starts with YAML front matter (---)
|
|
if readme_content.starts_with("---") {
|
|
// Find the closing ---
|
|
let lines: Vec<&str> = readme_content.lines().collect();
|
|
let mut yaml_end = None;
|
|
|
|
for (i, line) in lines.iter().enumerate().skip(1) {
|
|
if line.trim() == "---" {
|
|
yaml_end = Some(i);
|
|
break;
|
|
}
|
|
}
|
|
|
|
if let Some(end_idx) = yaml_end {
|
|
// Extract YAML and markdown
|
|
let yaml_lines = &lines[1..end_idx];
|
|
let yaml_str = yaml_lines.join("\n");
|
|
let markdown_lines = &lines[end_idx + 1..];
|
|
let markdown_content = markdown_lines.join("\n").trim().to_string();
|
|
|
|
// Parse YAML
|
|
let yaml_value: serde_yaml::Value =
|
|
serde_yaml::from_str(&yaml_str).map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Invalid YAML front matter: {}", e),
|
|
})?;
|
|
|
|
// Convert to ModelCardDoc
|
|
let mut card = ModelCardDoc::new(model_id, markdown_content);
|
|
|
|
if let serde_yaml::Value::Mapping(map) = yaml_value {
|
|
for (key, value) in map {
|
|
if let Some(key_str) = key.as_str() {
|
|
match key_str {
|
|
"language" => {
|
|
if let Some(lang_seq) = value.as_sequence() {
|
|
card.language = lang_seq
|
|
.iter()
|
|
.filter_map(|v| v.as_str().map(String::from))
|
|
.collect();
|
|
}
|
|
}
|
|
"license" => {
|
|
card.license = value.as_str().map(String::from);
|
|
}
|
|
"library_name" => {
|
|
card.library_name = value.as_str().map(String::from);
|
|
}
|
|
"tags" => {
|
|
if let Some(tags_seq) = value.as_sequence() {
|
|
card.tags = tags_seq
|
|
.iter()
|
|
.filter_map(|v| v.as_str().map(String::from))
|
|
.collect();
|
|
}
|
|
}
|
|
"pipeline_tag" => {
|
|
card.pipeline_tag = value.as_str().map(String::from);
|
|
}
|
|
"datasets" => {
|
|
if let Some(datasets_seq) = value.as_sequence() {
|
|
card.datasets = datasets_seq
|
|
.iter()
|
|
.filter_map(|v| v.as_str().map(String::from))
|
|
.collect();
|
|
}
|
|
}
|
|
"widget" => {
|
|
if let Some(widget_seq) = value.as_sequence() {
|
|
let widgets: Vec<Widget> = widget_seq
|
|
.iter()
|
|
.filter_map(|v| {
|
|
let json = serde_json::to_value(v).ok()?;
|
|
serde_json::from_value(json).ok()
|
|
})
|
|
.collect();
|
|
if !widgets.is_empty() {
|
|
card.widget = Some(widgets);
|
|
}
|
|
}
|
|
}
|
|
"metrics" => {
|
|
if let Some(metrics_seq) = value.as_sequence() {
|
|
card.metrics = metrics_seq
|
|
.iter()
|
|
.filter_map(|v| {
|
|
let json = serde_json::to_value(v).ok()?;
|
|
serde_json::from_value(json).ok()
|
|
})
|
|
.collect();
|
|
}
|
|
}
|
|
"model-index" => {
|
|
if let Some(index_seq) = value.as_sequence() {
|
|
let indices: Vec<ModelIndex> = index_seq
|
|
.iter()
|
|
.filter_map(|v| {
|
|
let json = serde_json::to_value(v).ok()?;
|
|
serde_json::from_value(json).ok()
|
|
})
|
|
.collect();
|
|
if !indices.is_empty() {
|
|
card.model_index = Some(indices);
|
|
}
|
|
}
|
|
}
|
|
"co2_eq_emissions" => {
|
|
let json = serde_json::to_value(&value).ok();
|
|
if let Some(json_val) = json {
|
|
card.co2_eq_emissions =
|
|
serde_json::from_value(json_val).ok();
|
|
}
|
|
}
|
|
_ => {
|
|
// Store in card_data
|
|
if let Ok(json_val) = serde_json::to_value(&value) {
|
|
card.card_data.insert(key_str.to_string(), json_val);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(card)
|
|
} else {
|
|
// No closing ---, treat as plain markdown
|
|
Ok(ModelCardDoc::new(model_id, readme_content.to_string()))
|
|
}
|
|
} else {
|
|
// No front matter, just markdown
|
|
Ok(ModelCardDoc::new(model_id, readme_content.to_string()))
|
|
}
|
|
}
|
|
|
|
/// Serialize model card documentation to README.md format with YAML front matter.
|
|
pub fn to_readme(&self) -> String {
|
|
let mut yaml_map = serde_yaml::Mapping::new();
|
|
|
|
// Add fields to YAML
|
|
if !self.language.is_empty() {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("language".to_string()),
|
|
serde_yaml::to_value(&self.language).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if let Some(ref license) = self.license {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("license".to_string()),
|
|
serde_yaml::Value::String(license.clone()),
|
|
);
|
|
}
|
|
|
|
if let Some(ref library_name) = self.library_name {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("library_name".to_string()),
|
|
serde_yaml::Value::String(library_name.clone()),
|
|
);
|
|
}
|
|
|
|
if !self.tags.is_empty() {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("tags".to_string()),
|
|
serde_yaml::to_value(&self.tags).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if let Some(ref pipeline_tag) = self.pipeline_tag {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("pipeline_tag".to_string()),
|
|
serde_yaml::Value::String(pipeline_tag.clone()),
|
|
);
|
|
}
|
|
|
|
if !self.datasets.is_empty() {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("datasets".to_string()),
|
|
serde_yaml::to_value(&self.datasets).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if let Some(ref widget) = self.widget {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("widget".to_string()),
|
|
serde_yaml::to_value(widget).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if !self.metrics.is_empty() {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("metrics".to_string()),
|
|
serde_yaml::to_value(&self.metrics).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if let Some(ref model_index) = self.model_index {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("model-index".to_string()),
|
|
serde_yaml::to_value(model_index).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
if let Some(ref co2) = self.co2_eq_emissions {
|
|
yaml_map.insert(
|
|
serde_yaml::Value::String("co2_eq_emissions".to_string()),
|
|
serde_yaml::to_value(co2).unwrap_or(serde_yaml::Value::Null),
|
|
);
|
|
}
|
|
|
|
// Add custom card_data
|
|
for (key, value) in &self.card_data {
|
|
if let Ok(yaml_value) = serde_yaml::to_value(value) {
|
|
yaml_map.insert(serde_yaml::Value::String(key.clone()), yaml_value);
|
|
}
|
|
}
|
|
|
|
// Serialize YAML
|
|
let yaml_content =
|
|
serde_yaml::to_string(&serde_yaml::Value::Mapping(yaml_map)).unwrap_or_default();
|
|
|
|
// Combine with markdown
|
|
if yaml_content.trim().is_empty() || yaml_content.trim() == "{}" {
|
|
// No YAML data, just return markdown
|
|
self.content.clone()
|
|
} else {
|
|
format!("---\n{}---\n\n{}", yaml_content, self.content)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Builder for creating model card documentation ergonomically.
|
|
pub struct ModelCardBuilder {
|
|
card: ModelCardDoc,
|
|
}
|
|
|
|
impl ModelCardBuilder {
|
|
/// Create a new builder.
|
|
pub fn new(model_id: ModelId) -> Self {
|
|
Self {
|
|
card: ModelCardDoc::new(model_id, String::new()),
|
|
}
|
|
}
|
|
|
|
/// Set content.
|
|
pub fn content(mut self, content: impl Into<String>) -> Self {
|
|
self.card.content = content.into();
|
|
self
|
|
}
|
|
|
|
/// Add a language.
|
|
pub fn language(mut self, lang: impl Into<String>) -> Self {
|
|
self.card.language.push(lang.into());
|
|
self
|
|
}
|
|
|
|
/// Add multiple languages.
|
|
pub fn languages(mut self, langs: Vec<String>) -> Self {
|
|
self.card.language = langs;
|
|
self
|
|
}
|
|
|
|
/// Set license.
|
|
pub fn license(mut self, license: impl Into<String>) -> Self {
|
|
self.card.license = Some(license.into());
|
|
self
|
|
}
|
|
|
|
/// Set library name.
|
|
pub fn library_name(mut self, name: impl Into<String>) -> Self {
|
|
self.card.library_name = Some(name.into());
|
|
self
|
|
}
|
|
|
|
/// Add a tag.
|
|
pub fn tag(mut self, tag: impl Into<String>) -> Self {
|
|
self.card.tags.push(tag.into());
|
|
self
|
|
}
|
|
|
|
/// Add multiple tags.
|
|
pub fn tags(mut self, tags: Vec<String>) -> Self {
|
|
self.card.tags = tags;
|
|
self
|
|
}
|
|
|
|
/// Set pipeline tag.
|
|
pub fn pipeline_tag(mut self, tag: impl Into<String>) -> Self {
|
|
self.card.pipeline_tag = Some(tag.into());
|
|
self
|
|
}
|
|
|
|
/// Add a widget.
|
|
pub fn widget(mut self, widget: Widget) -> Self {
|
|
self.card.widget.get_or_insert_with(Vec::new).push(widget);
|
|
self
|
|
}
|
|
|
|
/// Add multiple widgets.
|
|
pub fn widgets(mut self, widgets: Vec<Widget>) -> Self {
|
|
self.card.widget = Some(widgets);
|
|
self
|
|
}
|
|
|
|
/// Add a dataset.
|
|
pub fn dataset(mut self, dataset: impl Into<String>) -> Self {
|
|
self.card.datasets.push(dataset.into());
|
|
self
|
|
}
|
|
|
|
/// Add multiple datasets.
|
|
pub fn datasets(mut self, datasets: Vec<String>) -> Self {
|
|
self.card.datasets = datasets;
|
|
self
|
|
}
|
|
|
|
/// Add a metric.
|
|
pub fn metric(mut self, metric: MetricResult) -> Self {
|
|
self.card.metrics.push(metric);
|
|
self
|
|
}
|
|
|
|
/// Add multiple metrics.
|
|
pub fn metrics(mut self, metrics: Vec<MetricResult>) -> Self {
|
|
self.card.metrics = metrics;
|
|
self
|
|
}
|
|
|
|
/// Add a model index.
|
|
pub fn model_index(mut self, index: ModelIndex) -> Self {
|
|
self.card
|
|
.model_index
|
|
.get_or_insert_with(Vec::new)
|
|
.push(index);
|
|
self
|
|
}
|
|
|
|
/// Set CO2 emissions.
|
|
pub fn co2_emissions(mut self, emissions: CO2Emissions) -> Self {
|
|
self.card.co2_eq_emissions = Some(emissions);
|
|
self
|
|
}
|
|
|
|
/// Add custom card data.
|
|
pub fn card_data(mut self, key: impl Into<String>, value: Value) -> Self {
|
|
self.card.card_data.insert(key.into(), value);
|
|
self
|
|
}
|
|
|
|
/// Build the model card documentation.
|
|
pub fn build(self) -> ModelCardDoc {
|
|
self.card
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
fn test_model_id() -> ModelId {
|
|
ModelId::new("rustytorch", "test-model")
|
|
}
|
|
|
|
#[test]
|
|
fn test_widget_creation() {
|
|
let widget = Widget::new("Hello, world!");
|
|
assert_eq!(widget.text, "Hello, world!");
|
|
assert!(widget.parameters.is_none());
|
|
assert!(widget.output.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_widget_with_parameters() {
|
|
let mut params = HashMap::new();
|
|
params.insert("max_length".to_string(), json!(50));
|
|
|
|
let widget = Widget::new("Test input")
|
|
.with_parameters(params.clone())
|
|
.with_output(json!("Test output"));
|
|
|
|
assert_eq!(widget.text, "Test input");
|
|
assert_eq!(widget.parameters, Some(params));
|
|
assert_eq!(widget.output, Some(json!("Test output")));
|
|
}
|
|
|
|
#[test]
|
|
fn test_metric_result_creation() {
|
|
let metric = MetricResult::new("Accuracy", "accuracy", 0.95);
|
|
assert_eq!(metric.name, "Accuracy");
|
|
assert_eq!(metric.metric_type, "accuracy");
|
|
assert_eq!(metric.value, 0.95);
|
|
assert!(metric.dataset.is_none());
|
|
assert!(metric.split.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_metric_result_with_dataset() {
|
|
let metric = MetricResult::new("F1 Score", "f1", 0.88)
|
|
.with_dataset("squad")
|
|
.with_split("test");
|
|
|
|
assert_eq!(metric.dataset, Some("squad".to_string()));
|
|
assert_eq!(metric.split, Some("test".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_index_creation() {
|
|
let metric1 = MetricResult::new("Accuracy", "accuracy", 0.92);
|
|
let metric2 = MetricResult::new("F1", "f1", 0.89);
|
|
|
|
let index = ModelIndex::new("benchmark-1")
|
|
.add_result(metric1)
|
|
.add_result(metric2);
|
|
|
|
assert_eq!(index.name, "benchmark-1");
|
|
assert_eq!(index.results.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_co2_emissions_creation() {
|
|
let co2 = CO2Emissions::new(125.5)
|
|
.with_source("codecarbon")
|
|
.with_training_type("pretraining");
|
|
|
|
assert_eq!(co2.emissions, 125.5);
|
|
assert_eq!(co2.source, Some("codecarbon".to_string()));
|
|
assert_eq!(co2.training_type, Some("pretraining".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_creation() {
|
|
let card = ModelCardDoc::new(test_model_id(), "# Test Model\n\nThis is a test.");
|
|
assert_eq!(card.model_id, test_model_id());
|
|
assert_eq!(card.content, "# Test Model\n\nThis is a test.");
|
|
assert_eq!(card.library_name, Some("rustytorch".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_builder() {
|
|
let card = ModelCardBuilder::new(test_model_id())
|
|
.content("# My Model\n\nA great model.")
|
|
.language("en")
|
|
.license("MIT")
|
|
.tag("nlp")
|
|
.tag("transformers")
|
|
.pipeline_tag("text-generation")
|
|
.dataset("wikitext")
|
|
.build();
|
|
|
|
assert_eq!(card.content, "# My Model\n\nA great model.");
|
|
assert_eq!(card.language, vec!["en"]);
|
|
assert_eq!(card.license, Some("MIT".to_string()));
|
|
assert_eq!(card.tags, vec!["nlp", "transformers"]);
|
|
assert_eq!(card.pipeline_tag, Some("text-generation".to_string()));
|
|
assert_eq!(card.datasets, vec!["wikitext"]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_from_readme_plain_markdown() {
|
|
let readme = "# Test Model\n\nThis is a test model.";
|
|
let card = ModelCardDoc::from_readme(test_model_id(), readme).unwrap();
|
|
|
|
assert_eq!(card.content, "# Test Model\n\nThis is a test model.");
|
|
assert!(card.language.is_empty());
|
|
assert!(card.license.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_from_readme_with_yaml() {
|
|
let readme = r#"---
|
|
language:
|
|
- en
|
|
- zh
|
|
license: MIT
|
|
library_name: rustytorch
|
|
tags:
|
|
- nlp
|
|
- transformers
|
|
pipeline_tag: text-generation
|
|
datasets:
|
|
- wikitext
|
|
- bookcorpus
|
|
---
|
|
|
|
# Test Model
|
|
|
|
This is a test model with YAML front matter."#;
|
|
|
|
let card = ModelCardDoc::from_readme(test_model_id(), readme).unwrap();
|
|
|
|
assert_eq!(card.language, vec!["en", "zh"]);
|
|
assert_eq!(card.license, Some("MIT".to_string()));
|
|
assert_eq!(card.library_name, Some("rustytorch".to_string()));
|
|
assert_eq!(card.tags, vec!["nlp", "transformers"]);
|
|
assert_eq!(card.pipeline_tag, Some("text-generation".to_string()));
|
|
assert_eq!(card.datasets, vec!["wikitext", "bookcorpus"]);
|
|
assert!(card.content.contains("# Test Model"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_from_readme_with_widgets() {
|
|
let readme = r#"---
|
|
widget:
|
|
- text: "Hello, world!"
|
|
- text: "Translate this"
|
|
parameters:
|
|
max_length: 50
|
|
---
|
|
|
|
# Model with Widgets"#;
|
|
|
|
let card = ModelCardDoc::from_readme(test_model_id(), readme).unwrap();
|
|
|
|
assert!(card.widget.is_some());
|
|
let widgets = card.widget.unwrap();
|
|
assert_eq!(widgets.len(), 2);
|
|
assert_eq!(widgets[0].text, "Hello, world!");
|
|
assert_eq!(widgets[1].text, "Translate this");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_from_readme_with_metrics() {
|
|
let readme = r#"---
|
|
metrics:
|
|
- name: Accuracy
|
|
type: accuracy
|
|
value: 0.95
|
|
dataset: squad
|
|
split: test
|
|
- name: F1 Score
|
|
type: f1
|
|
value: 0.88
|
|
---
|
|
|
|
# Evaluated Model"#;
|
|
|
|
let card = ModelCardDoc::from_readme(test_model_id(), readme).unwrap();
|
|
|
|
assert_eq!(card.metrics.len(), 2);
|
|
assert_eq!(card.metrics[0].name, "Accuracy");
|
|
assert_eq!(card.metrics[0].value, 0.95);
|
|
assert_eq!(card.metrics[1].name, "F1 Score");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_from_readme_with_co2() {
|
|
let readme = r#"---
|
|
co2_eq_emissions:
|
|
emissions: 125.5
|
|
source: codecarbon
|
|
training_type: pretraining
|
|
---
|
|
|
|
# Green Model"#;
|
|
|
|
let card = ModelCardDoc::from_readme(test_model_id(), readme).unwrap();
|
|
|
|
assert!(card.co2_eq_emissions.is_some());
|
|
let co2 = card.co2_eq_emissions.unwrap();
|
|
assert_eq!(co2.emissions, 125.5);
|
|
assert_eq!(co2.source, Some("codecarbon".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_to_readme() {
|
|
let card = ModelCardBuilder::new(test_model_id())
|
|
.content("# Test Model\n\nContent here.")
|
|
.language("en")
|
|
.license("MIT")
|
|
.tag("nlp")
|
|
.pipeline_tag("text-generation")
|
|
.build();
|
|
|
|
let readme = card.to_readme();
|
|
|
|
assert!(readme.contains("---"));
|
|
assert!(readme.contains("language"));
|
|
assert!(readme.contains("- en"));
|
|
assert!(readme.contains("license: MIT"));
|
|
assert!(readme.contains("# Test Model"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_roundtrip() {
|
|
let original = ModelCardBuilder::new(test_model_id())
|
|
.content("# Test Model\n\nRoundtrip test.")
|
|
.languages(vec!["en".to_string(), "fr".to_string()])
|
|
.license("Apache-2.0")
|
|
.tags(vec!["nlp".to_string(), "transformer".to_string()])
|
|
.pipeline_tag("text-classification")
|
|
.datasets(vec!["glue".to_string()])
|
|
.metric(MetricResult::new("Accuracy", "accuracy", 0.93))
|
|
.build();
|
|
|
|
let readme = original.to_readme();
|
|
let parsed = ModelCardDoc::from_readme(test_model_id(), &readme).unwrap();
|
|
|
|
assert_eq!(parsed.language, original.language);
|
|
assert_eq!(parsed.license, original.license);
|
|
assert_eq!(parsed.tags, original.tags);
|
|
assert_eq!(parsed.pipeline_tag, original.pipeline_tag);
|
|
assert_eq!(parsed.datasets, original.datasets);
|
|
assert_eq!(parsed.metrics.len(), original.metrics.len());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_with_custom_data() {
|
|
let mut card = ModelCardBuilder::new(test_model_id())
|
|
.content("# Custom Data Model")
|
|
.card_data("custom_field", json!("custom_value"))
|
|
.card_data("version", json!("1.0.0"))
|
|
.build();
|
|
|
|
assert_eq!(
|
|
card.card_data.get("custom_field"),
|
|
Some(&json!("custom_value"))
|
|
);
|
|
assert_eq!(card.card_data.get("version"), Some(&json!("1.0.0")));
|
|
|
|
// Test roundtrip
|
|
let readme = card.to_readme();
|
|
let parsed = ModelCardDoc::from_readme(test_model_id(), &readme).unwrap();
|
|
assert_eq!(
|
|
parsed.card_data.get("custom_field"),
|
|
Some(&json!("custom_value"))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_empty_yaml() {
|
|
let mut card = ModelCardDoc::new(test_model_id(), "# Simple Model\n\nNo metadata.");
|
|
// Remove the default library_name to get truly empty YAML
|
|
card.library_name = None;
|
|
let readme = card.to_readme();
|
|
|
|
// Should not have YAML front matter if no metadata
|
|
assert_eq!(readme, "# Simple Model\n\nNo metadata.");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_card_doc_serialization() {
|
|
let card = ModelCardBuilder::new(test_model_id())
|
|
.content("Test")
|
|
.language("en")
|
|
.license("MIT")
|
|
.build();
|
|
|
|
let json = serde_json::to_string(&card).unwrap();
|
|
let deserialized: ModelCardDoc = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(deserialized.model_id, card.model_id);
|
|
assert_eq!(deserialized.language, card.language);
|
|
assert_eq!(deserialized.license, card.license);
|
|
}
|
|
|
|
#[test]
|
|
fn test_widget_serialization() {
|
|
let widget = Widget::new("test").with_parameters({
|
|
let mut map = HashMap::new();
|
|
map.insert("key".to_string(), json!("value"));
|
|
map
|
|
});
|
|
|
|
let json = serde_json::to_string(&widget).unwrap();
|
|
let deserialized: Widget = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(deserialized.text, widget.text);
|
|
assert_eq!(deserialized.parameters, widget.parameters);
|
|
}
|
|
|
|
#[test]
|
|
fn test_metric_result_serialization() {
|
|
let metric = MetricResult::new("Test", "test", 0.5)
|
|
.with_dataset("test_dataset")
|
|
.with_split("test");
|
|
|
|
let json = serde_json::to_string(&metric).unwrap();
|
|
let deserialized: MetricResult = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(deserialized.name, metric.name);
|
|
assert_eq!(deserialized.value, metric.value);
|
|
assert_eq!(deserialized.dataset, metric.dataset);
|
|
}
|
|
}
|