Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,913 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user