use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ModelCard {
pub model_id: String,
pub name: String,
pub version: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub author: Option<String>,
pub created_at: String,
pub framework_version: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rust_version: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub license: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub training_data: Option<TrainingDataInfo>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub hyperparameters: HashMap<String, serde_json::Value>,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub metrics: HashMap<String, serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub architecture: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub param_count: Option<u64>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub target_hardware: Vec<String>,
#[serde(default, skip_serializing_if = "HashMap::is_empty", flatten)]
pub extra: HashMap<String, serde_json::Value>,
}
impl ModelCard {
#[must_use]
pub fn new(model_id: impl Into<String>, version: impl Into<String>) -> Self {
let model_id = model_id.into();
let name = model_id.clone();
Self {
model_id,
name,
version: version.into(),
author: None,
created_at: Self::now_iso8601(),
framework_version: format!("aprender {}", env!("CARGO_PKG_VERSION")),
rust_version: option_env!("RUSTC_VERSION").map(String::from),
description: None,
license: None,
training_data: None,
hyperparameters: HashMap::new(),
metrics: HashMap::new(),
architecture: None,
param_count: None,
target_hardware: vec!["cpu".to_string()],
extra: HashMap::new(),
}
}
#[must_use]
pub fn with_name(mut self, name: impl Into<String>) -> Self {
self.name = name.into();
self
}
#[must_use]
pub fn with_author(mut self, author: impl Into<String>) -> Self {
self.author = Some(author.into());
self
}
#[must_use]
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = Some(desc.into());
self
}
#[must_use]
pub fn with_license(mut self, license: impl Into<String>) -> Self {
self.license = Some(license.into());
self
}
#[must_use]
pub fn with_architecture(mut self, arch: impl Into<String>) -> Self {
self.architecture = Some(arch.into());
self
}
#[must_use]
pub fn with_param_count(mut self, count: u64) -> Self {
self.param_count = Some(count);
self
}
#[must_use]
pub fn with_training_data(mut self, data: TrainingDataInfo) -> Self {
self.training_data = Some(data);
self
}
#[must_use]
pub fn with_hyperparameter(
mut self,
key: impl Into<String>,
value: impl Into<serde_json::Value>,
) -> Self {
self.hyperparameters.insert(key.into(), value.into());
self
}
#[must_use]
pub fn with_metric(
mut self,
key: impl Into<String>,
value: impl Into<serde_json::Value>,
) -> Self {
self.metrics.insert(key.into(), value.into());
self
}
fn now_iso8601() -> String {
let now = std::time::SystemTime::now();
let duration = now
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
let secs = duration.as_secs();
let days = secs / 86400;
let time_secs = secs % 86400;
let hours = time_secs / 3600;
let minutes = (time_secs % 3600) / 60;
let seconds = time_secs % 60;
let (year, month, day) = days_to_ymd(days);
format!("{year:04}-{month:02}-{day:02}T{hours:02}:{minutes:02}:{seconds:02}Z")
}
#[must_use]
pub fn to_huggingface(&self) -> String {
use std::fmt::Write;
let mut output = String::from("---\n");
if let Some(license) = &self.license {
let _ = writeln!(output, "license: {}", license.to_lowercase());
}
output.push_str("pipeline_tag: text-generation\n");
output.push_str("tags:\n");
if let Some(arch) = &self.architecture {
let _ = writeln!(output, " - {}", arch.to_lowercase());
}
output.push_str(" - aprender\n");
output.push_str(" - rust\n");
output.push_str("model-index:\n");
let _ = writeln!(output, " - name: {}", self.model_id);
if !self.metrics.is_empty() {
output.push_str(" results:\n");
output.push_str(" - task:\n");
output.push_str(" type: text-generation\n");
output.push_str(" metrics:\n");
for (key, value) in &self.metrics {
let _ = writeln!(output, " - name: {key}");
output.push_str(" type: custom\n");
let _ = writeln!(output, " value: {value}");
}
}
output.push_str("---\n\n");
let _ = writeln!(output, "# {}\n", self.name);
if let Some(desc) = &self.description {
let _ = writeln!(output, "{desc}\n");
}
if let Some(data) = &self.training_data {
output.push_str("## Training Data\n\n");
let _ = writeln!(output, "- **Source:** {}", data.name);
if let Some(samples) = data.samples {
let _ = writeln!(output, "- **Samples:** {samples}");
}
if let Some(hash) = &data.hash {
let _ = writeln!(output, "- **Hash:** `{hash}`");
}
output.push('\n');
}
output.push_str("## Framework\n\n");
let _ = writeln!(output, "- **Version:** {}", self.framework_version);
if let Some(rust) = &self.rust_version {
let _ = writeln!(output, "- **Rust:** {rust}");
}
output
}
pub fn to_json(&self) -> Result<String, serde_json::Error> {
serde_json::to_string_pretty(self)
}
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
serde_json::from_str(json)
}
}
impl Default for ModelCard {
fn default() -> Self {
Self::new("unnamed", "0.0.0")
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TrainingDataInfo {
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub samples: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hash: Option<String>,
}
impl TrainingDataInfo {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
samples: None,
hash: None,
}
}
#[must_use]
pub fn with_samples(mut self, count: u64) -> Self {
self.samples = Some(count);
self
}
#[must_use]
pub fn with_hash(mut self, hash: impl Into<String>) -> Self {
self.hash = Some(hash.into());
self
}
}
fn days_to_ymd(days: u64) -> (u32, u32, u32) {
let mut remaining = days as i64;
let mut year = 1970i32;
loop {
let days_in_year = if is_leap_year(year) { 366 } else { 365 };
if remaining < days_in_year {
break;
}
remaining -= days_in_year;
year += 1;
}
let leap = is_leap_year(year);
let months = if leap {
[31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
} else {
[31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
};
let mut month = 1u32;
for days_in_month in months {
if remaining < days_in_month {
break;
}
remaining -= days_in_month;
month += 1;
}
let day = remaining as u32 + 1;
(year as u32, month, day)
}
fn is_leap_year(year: i32) -> bool {
(year % 4 == 0 && year % 100 != 0) || (year % 400 == 0)
}
include!("generating.rs");