778 lines
24 KiB
Rust
778 lines
24 KiB
Rust
//! Pre-trained model management and downloading.
|
|
//!
|
|
//! Provides a registry of available artifact detection models and
|
|
//! utilities for downloading and caching them.
|
|
|
|
use crate::error::{ArtifactError, ArtifactResult};
|
|
use crate::labels::ArtifactType;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::path::{Path, PathBuf};
|
|
|
|
/// Information about a pre-trained model
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelInfo {
|
|
/// Unique model identifier
|
|
pub id: String,
|
|
/// Human-readable name
|
|
pub name: String,
|
|
/// Model description
|
|
pub description: String,
|
|
/// Model version
|
|
pub version: String,
|
|
/// Source URL for downloading
|
|
pub source: ModelSource,
|
|
/// Expected input channels
|
|
pub n_channels: usize,
|
|
/// Expected sampling frequency
|
|
pub sfreq: f64,
|
|
/// Window size the model was trained on
|
|
pub window_size: usize,
|
|
/// Artifact types this model can detect
|
|
pub artifact_types: Vec<ArtifactType>,
|
|
/// File size in bytes
|
|
pub file_size: u64,
|
|
/// SHA256 hash for verification
|
|
pub sha256: Option<String>,
|
|
/// License information
|
|
pub license: String,
|
|
/// Training dataset description
|
|
pub training_data: String,
|
|
}
|
|
|
|
impl ModelInfo {
|
|
/// Create a new model info
|
|
pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
|
|
Self {
|
|
id: id.into(),
|
|
name: name.into(),
|
|
description: String::new(),
|
|
version: "1.0.0".to_string(),
|
|
source: ModelSource::Local(PathBuf::new()),
|
|
n_channels: 64,
|
|
sfreq: 1000.0,
|
|
window_size: 1000,
|
|
artifact_types: ArtifactType::all(),
|
|
file_size: 0,
|
|
sha256: None,
|
|
license: "MIT".to_string(),
|
|
training_data: String::new(),
|
|
}
|
|
}
|
|
|
|
/// Set description
|
|
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
|
|
self.description = desc.into();
|
|
self
|
|
}
|
|
|
|
/// Set source
|
|
pub fn with_source(mut self, source: ModelSource) -> Self {
|
|
self.source = source;
|
|
self
|
|
}
|
|
|
|
/// Set input configuration
|
|
pub fn with_input_config(mut self, n_channels: usize, sfreq: f64, window_size: usize) -> Self {
|
|
self.n_channels = n_channels;
|
|
self.sfreq = sfreq;
|
|
self.window_size = window_size;
|
|
self
|
|
}
|
|
|
|
/// Set supported artifact types
|
|
pub fn with_artifact_types(mut self, types: Vec<ArtifactType>) -> Self {
|
|
self.artifact_types = types;
|
|
self
|
|
}
|
|
|
|
/// Check if model supports a specific artifact type
|
|
pub fn supports(&self, artifact_type: ArtifactType) -> bool {
|
|
self.artifact_types.contains(&artifact_type)
|
|
}
|
|
}
|
|
|
|
/// Source location for a model
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum ModelSource {
|
|
/// Local file path
|
|
Local(PathBuf),
|
|
/// Remote URL
|
|
Url(String),
|
|
/// Hugging Face Hub
|
|
HuggingFace {
|
|
/// Repository identifier (e.g., "username/model-name")
|
|
repo: String,
|
|
/// Name of the model file within the repository
|
|
filename: String,
|
|
},
|
|
/// GitHub release
|
|
GitHub {
|
|
/// Repository owner (username or organization)
|
|
owner: String,
|
|
/// Repository name
|
|
repo: String,
|
|
/// Release tag (e.g., "v1.0.0")
|
|
tag: String,
|
|
/// Release asset filename
|
|
asset: String,
|
|
},
|
|
}
|
|
|
|
impl ModelSource {
|
|
/// Create a local source
|
|
pub fn local(path: impl Into<PathBuf>) -> Self {
|
|
Self::Local(path.into())
|
|
}
|
|
|
|
/// Create a URL source
|
|
pub fn url(url: impl Into<String>) -> Self {
|
|
Self::Url(url.into())
|
|
}
|
|
|
|
/// Create a Hugging Face source
|
|
pub fn huggingface(repo: impl Into<String>, filename: impl Into<String>) -> Self {
|
|
Self::HuggingFace {
|
|
repo: repo.into(),
|
|
filename: filename.into(),
|
|
}
|
|
}
|
|
|
|
/// Create a GitHub release source
|
|
pub fn github(
|
|
owner: impl Into<String>,
|
|
repo: impl Into<String>,
|
|
tag: impl Into<String>,
|
|
asset: impl Into<String>,
|
|
) -> Self {
|
|
Self::GitHub {
|
|
owner: owner.into(),
|
|
repo: repo.into(),
|
|
tag: tag.into(),
|
|
asset: asset.into(),
|
|
}
|
|
}
|
|
|
|
/// Get the download URL
|
|
pub fn to_url(&self) -> Option<String> {
|
|
match self {
|
|
Self::Local(_) => None,
|
|
Self::Url(url) => Some(url.clone()),
|
|
Self::HuggingFace { repo, filename } => Some(format!(
|
|
"https://huggingface.co/{}/resolve/main/{}",
|
|
repo, filename
|
|
)),
|
|
Self::GitHub {
|
|
owner,
|
|
repo,
|
|
tag,
|
|
asset,
|
|
} => Some(format!(
|
|
"https://github.com/{}/{}/releases/download/{}/{}",
|
|
owner, repo, tag, asset
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Check if this is a local source
|
|
pub fn is_local(&self) -> bool {
|
|
matches!(self, Self::Local(_))
|
|
}
|
|
}
|
|
|
|
/// Registry of available pre-trained models
|
|
#[derive(Debug, Clone)]
|
|
pub struct ModelRegistry {
|
|
/// Available models by ID
|
|
models: HashMap<String, ModelInfo>,
|
|
/// Cache directory for downloaded models
|
|
cache_dir: PathBuf,
|
|
}
|
|
|
|
impl Default for ModelRegistry {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl ModelRegistry {
|
|
/// Create a new model registry with default models
|
|
pub fn new() -> Self {
|
|
let cache_dir = Self::default_cache_dir();
|
|
let mut registry = Self {
|
|
models: HashMap::new(),
|
|
cache_dir,
|
|
};
|
|
|
|
// Register built-in models
|
|
registry.register_builtin_models();
|
|
|
|
registry
|
|
}
|
|
|
|
/// Create with custom cache directory
|
|
pub fn with_cache_dir(cache_dir: impl Into<PathBuf>) -> Self {
|
|
let cache_dir = cache_dir.into();
|
|
let mut registry = Self {
|
|
models: HashMap::new(),
|
|
cache_dir,
|
|
};
|
|
registry.register_builtin_models();
|
|
registry
|
|
}
|
|
|
|
/// Get default cache directory
|
|
pub fn default_cache_dir() -> PathBuf {
|
|
dirs::cache_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join("rustytorch")
|
|
.join("models")
|
|
.join("neuro-artifacts")
|
|
}
|
|
|
|
/// Register built-in models
|
|
fn register_builtin_models(&mut self) {
|
|
// EEG artifact detector (general purpose)
|
|
let eeg_general = ModelInfo::new("eeg-artifact-v1", "EEG Artifact Detector v1")
|
|
.with_description(
|
|
"General-purpose EEG artifact detection model trained on diverse datasets",
|
|
)
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"eeg-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(64, 1000.0, 1000)
|
|
.with_artifact_types(vec![
|
|
ArtifactType::EyeBlink,
|
|
ArtifactType::EyeMovement,
|
|
ArtifactType::Muscle,
|
|
ArtifactType::Heartbeat,
|
|
ArtifactType::LineNoise,
|
|
ArtifactType::Movement,
|
|
]);
|
|
self.register(eeg_general);
|
|
|
|
// MEG artifact detector
|
|
let meg_general = ModelInfo::new("meg-artifact-v1", "MEG Artifact Detector v1")
|
|
.with_description(
|
|
"MEG artifact detection model optimized for magnetometer/gradiometer data",
|
|
)
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"meg-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(306, 1000.0, 1000) // 306 channels for Elekta Neuromag
|
|
.with_artifact_types(ArtifactType::all());
|
|
self.register(meg_general);
|
|
|
|
// Eye artifact specialist
|
|
let eye_artifact = ModelInfo::new("eye-artifact-v1", "Eye Artifact Specialist v1")
|
|
.with_description(
|
|
"Specialized model for detecting eye-related artifacts with high precision",
|
|
)
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"eye-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(64, 500.0, 500)
|
|
.with_artifact_types(vec![ArtifactType::EyeBlink, ArtifactType::EyeMovement]);
|
|
self.register(eye_artifact);
|
|
|
|
// Muscle artifact specialist
|
|
let muscle_artifact = ModelInfo::new("muscle-artifact-v1", "Muscle Artifact Specialist v1")
|
|
.with_description("Specialized model for detecting EMG contamination")
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"muscle-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(64, 1000.0, 500)
|
|
.with_artifact_types(vec![ArtifactType::Muscle]);
|
|
self.register(muscle_artifact);
|
|
|
|
// Cardiac artifact specialist
|
|
let cardiac_artifact =
|
|
ModelInfo::new("cardiac-artifact-v1", "Cardiac Artifact Specialist v1")
|
|
.with_description("Specialized model for detecting heartbeat artifacts")
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"cardiac-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(64, 500.0, 1000)
|
|
.with_artifact_types(vec![ArtifactType::Heartbeat]);
|
|
self.register(cardiac_artifact);
|
|
|
|
// Lightweight model for real-time
|
|
let realtime = ModelInfo::new("realtime-artifact-v1", "Real-Time Artifact Detector v1")
|
|
.with_description(
|
|
"Lightweight model optimized for real-time artifact detection with minimal latency",
|
|
)
|
|
.with_source(ModelSource::huggingface(
|
|
"rustytorch/neuro-artifacts",
|
|
"realtime-artifact-v1.onnx",
|
|
))
|
|
.with_input_config(32, 500.0, 250) // Smaller window for low latency
|
|
.with_artifact_types(vec![
|
|
ArtifactType::EyeBlink,
|
|
ArtifactType::Muscle,
|
|
ArtifactType::Heartbeat,
|
|
]);
|
|
self.register(realtime);
|
|
}
|
|
|
|
/// Register a model
|
|
pub fn register(&mut self, model: ModelInfo) {
|
|
self.models.insert(model.id.clone(), model);
|
|
}
|
|
|
|
/// Get model info by ID
|
|
pub fn get(&self, id: &str) -> Option<&ModelInfo> {
|
|
self.models.get(id)
|
|
}
|
|
|
|
/// List all available models
|
|
pub fn list(&self) -> Vec<&ModelInfo> {
|
|
self.models.values().collect()
|
|
}
|
|
|
|
/// List models that support a specific artifact type
|
|
pub fn list_for_artifact(&self, artifact_type: ArtifactType) -> Vec<&ModelInfo> {
|
|
self.models
|
|
.values()
|
|
.filter(|m| m.supports(artifact_type))
|
|
.collect()
|
|
}
|
|
|
|
/// Get cache directory
|
|
pub fn cache_dir(&self) -> &Path {
|
|
&self.cache_dir
|
|
}
|
|
|
|
/// Get cached path for a model
|
|
pub fn cached_path(&self, model_id: &str) -> PathBuf {
|
|
self.cache_dir.join(format!("{}.onnx", model_id))
|
|
}
|
|
|
|
/// Check if model is cached locally
|
|
pub fn is_cached(&self, model_id: &str) -> bool {
|
|
self.cached_path(model_id).exists()
|
|
}
|
|
|
|
/// Get model path (cached or local)
|
|
pub fn get_path(&self, model_id: &str) -> ArtifactResult<PathBuf> {
|
|
let model = self
|
|
.get(model_id)
|
|
.ok_or_else(|| ArtifactError::Model(format!("Unknown model: {}", model_id)))?;
|
|
|
|
if let ModelSource::Local(path) = &model.source {
|
|
if path.exists() {
|
|
Ok(path.clone())
|
|
} else {
|
|
Err(ArtifactError::Model(format!(
|
|
"Model file not found: {}",
|
|
path.display()
|
|
)))
|
|
}
|
|
} else {
|
|
let cached = self.cached_path(model_id);
|
|
if cached.exists() {
|
|
Ok(cached)
|
|
} else {
|
|
Err(ArtifactError::Model(format!(
|
|
"Model not cached. Call download_model first: {}",
|
|
model_id
|
|
)))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Download a model to the cache
|
|
pub async fn download_model(registry: &ModelRegistry, model_id: &str) -> ArtifactResult<PathBuf> {
|
|
let model = registry
|
|
.get(model_id)
|
|
.ok_or_else(|| ArtifactError::Model(format!("Unknown model: {}", model_id)))?;
|
|
|
|
// Check if already cached
|
|
let cache_path = registry.cached_path(model_id);
|
|
if cache_path.exists() {
|
|
return Ok(cache_path);
|
|
}
|
|
|
|
// Get download URL
|
|
let url = model.source.to_url().ok_or_else(|| {
|
|
ArtifactError::Download("Model source is local, cannot download".to_string())
|
|
})?;
|
|
|
|
// Ensure cache directory exists
|
|
if let Some(parent) = cache_path.parent() {
|
|
tokio::fs::create_dir_all(parent).await.map_err(|e| {
|
|
ArtifactError::Download(format!("Failed to create cache directory: {}", e))
|
|
})?;
|
|
}
|
|
|
|
// Download file
|
|
download_file(&url, &cache_path).await?;
|
|
|
|
// Verify hash if provided
|
|
if let Some(ref expected_hash) = model.sha256 {
|
|
let actual_hash = compute_file_hash(&cache_path).await?;
|
|
if actual_hash != *expected_hash {
|
|
// Remove corrupted file
|
|
let _ = tokio::fs::remove_file(&cache_path).await;
|
|
return Err(ArtifactError::Download(format!(
|
|
"Hash mismatch: expected {}, got {}",
|
|
expected_hash, actual_hash
|
|
)));
|
|
}
|
|
}
|
|
|
|
Ok(cache_path)
|
|
}
|
|
|
|
/// Download a file from URL
|
|
async fn download_file(url: &str, path: &Path) -> ArtifactResult<()> {
|
|
// Use reqwest if available, otherwise fall back to simpler method
|
|
#[cfg(feature = "download")]
|
|
{
|
|
let response = reqwest::get(url)
|
|
.await
|
|
.map_err(|e| ArtifactError::Download(format!("Failed to download: {}", e)))?;
|
|
|
|
if !response.status().is_success() {
|
|
return Err(ArtifactError::Download(format!(
|
|
"Download failed with status: {}",
|
|
response.status()
|
|
)));
|
|
}
|
|
|
|
let bytes = response
|
|
.bytes()
|
|
.await
|
|
.map_err(|e| ArtifactError::Download(format!("Failed to read response: {}", e)))?;
|
|
|
|
tokio::fs::write(path, bytes)
|
|
.await
|
|
.map_err(|e| ArtifactError::Download(format!("Failed to write file: {}", e)))?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(not(feature = "download"))]
|
|
{
|
|
// Without download feature, we can't download
|
|
Err(ArtifactError::Download(format!(
|
|
"Download feature not enabled. URL: {}, Path: {}",
|
|
url,
|
|
path.display()
|
|
)))
|
|
}
|
|
}
|
|
|
|
/// Compute SHA256 hash of a file
|
|
async fn compute_file_hash(path: &Path) -> ArtifactResult<String> {
|
|
use std::io::Read;
|
|
|
|
let path = path.to_path_buf();
|
|
tokio::task::spawn_blocking(move || {
|
|
let mut file = std::fs::File::open(&path).map_err(|e| {
|
|
ArtifactError::Download(format!("Failed to open file for hashing: {}", e))
|
|
})?;
|
|
|
|
let mut hasher = Sha256::new();
|
|
let mut buffer = [0u8; 8192];
|
|
|
|
loop {
|
|
let n = file
|
|
.read(&mut buffer)
|
|
.map_err(|e| ArtifactError::Download(format!("Failed to read file: {}", e)))?;
|
|
if n == 0 {
|
|
break;
|
|
}
|
|
hasher.update(&buffer[..n]);
|
|
}
|
|
|
|
let hash_bytes = hasher.finalize_inner();
|
|
let hex: String = hash_bytes.iter().map(|b| format!("{:02x}", b)).collect();
|
|
Ok(hex)
|
|
})
|
|
.await
|
|
.map_err(|e| ArtifactError::Download(format!("Hash computation failed: {}", e)))?
|
|
}
|
|
|
|
/// Simple SHA256 implementation for verification
|
|
struct Sha256 {
|
|
h: [u32; 8],
|
|
block: [u8; 64],
|
|
block_len: usize,
|
|
total_len: u64,
|
|
}
|
|
|
|
impl Sha256 {
|
|
fn new() -> Self {
|
|
Self {
|
|
h: [
|
|
0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab,
|
|
0x5be0cd19,
|
|
],
|
|
block: [0; 64],
|
|
block_len: 0,
|
|
total_len: 0,
|
|
}
|
|
}
|
|
|
|
fn update(&mut self, data: &[u8]) {
|
|
self.total_len += data.len() as u64;
|
|
let mut offset = 0;
|
|
|
|
// Fill current block
|
|
if self.block_len > 0 {
|
|
let space = 64 - self.block_len;
|
|
let to_copy = space.min(data.len());
|
|
self.block[self.block_len..self.block_len + to_copy].copy_from_slice(&data[..to_copy]);
|
|
self.block_len += to_copy;
|
|
offset = to_copy;
|
|
|
|
if self.block_len == 64 {
|
|
self.process_block();
|
|
self.block_len = 0;
|
|
}
|
|
}
|
|
|
|
// Process full blocks
|
|
while offset + 64 <= data.len() {
|
|
self.block.copy_from_slice(&data[offset..offset + 64]);
|
|
self.process_block();
|
|
offset += 64;
|
|
}
|
|
|
|
// Save remaining
|
|
if offset < data.len() {
|
|
let remaining = data.len() - offset;
|
|
self.block[..remaining].copy_from_slice(&data[offset..]);
|
|
self.block_len = remaining;
|
|
}
|
|
}
|
|
|
|
fn process_block(&mut self) {
|
|
const K: [u32; 64] = [
|
|
0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4,
|
|
0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe,
|
|
0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f,
|
|
0x4a7484aa, 0x5cb0a9dc, 0x76f988da, 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
|
|
0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc,
|
|
0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b,
|
|
0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, 0x19a4c116,
|
|
0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
|
|
0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7,
|
|
0xc67178f2,
|
|
];
|
|
|
|
let mut w = [0u32; 64];
|
|
|
|
// Prepare message schedule
|
|
for i in 0..16 {
|
|
w[i] = u32::from_be_bytes([
|
|
self.block[i * 4],
|
|
self.block[i * 4 + 1],
|
|
self.block[i * 4 + 2],
|
|
self.block[i * 4 + 3],
|
|
]);
|
|
}
|
|
|
|
for i in 16..64 {
|
|
let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3);
|
|
let s1 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10);
|
|
w[i] = w[i - 16]
|
|
.wrapping_add(s0)
|
|
.wrapping_add(w[i - 7])
|
|
.wrapping_add(s1);
|
|
}
|
|
|
|
// Working variables
|
|
let [mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h] = self.h;
|
|
|
|
// Compression function
|
|
for i in 0..64 {
|
|
let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
|
|
let ch = (e & f) ^ ((!e) & g);
|
|
let temp1 = h
|
|
.wrapping_add(s1)
|
|
.wrapping_add(ch)
|
|
.wrapping_add(K[i])
|
|
.wrapping_add(w[i]);
|
|
let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
|
|
let maj = (a & b) ^ (a & c) ^ (b & c);
|
|
let temp2 = s0.wrapping_add(maj);
|
|
|
|
h = g;
|
|
g = f;
|
|
f = e;
|
|
e = d.wrapping_add(temp1);
|
|
d = c;
|
|
c = b;
|
|
b = a;
|
|
a = temp1.wrapping_add(temp2);
|
|
}
|
|
|
|
// Update hash values
|
|
self.h[0] = self.h[0].wrapping_add(a);
|
|
self.h[1] = self.h[1].wrapping_add(b);
|
|
self.h[2] = self.h[2].wrapping_add(c);
|
|
self.h[3] = self.h[3].wrapping_add(d);
|
|
self.h[4] = self.h[4].wrapping_add(e);
|
|
self.h[5] = self.h[5].wrapping_add(f);
|
|
self.h[6] = self.h[6].wrapping_add(g);
|
|
self.h[7] = self.h[7].wrapping_add(h);
|
|
}
|
|
}
|
|
|
|
/// Format helper for hex output
|
|
struct HexFormatter([u8; 32]);
|
|
|
|
impl std::fmt::LowerHex for HexFormatter {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
for byte in &self.0 {
|
|
write!(f, "{:02x}", byte)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Sha256 {
|
|
fn finalize(self) -> HexFormatter {
|
|
HexFormatter(self.finalize_inner())
|
|
}
|
|
|
|
fn finalize_inner(mut self) -> [u8; 32] {
|
|
// Padding
|
|
let bit_len = self.total_len * 8;
|
|
self.block[self.block_len] = 0x80;
|
|
self.block_len += 1;
|
|
|
|
if self.block_len > 56 {
|
|
self.block[self.block_len..64].fill(0);
|
|
self.process_block();
|
|
self.block_len = 0;
|
|
}
|
|
|
|
self.block[self.block_len..56].fill(0);
|
|
self.block[56..64].copy_from_slice(&bit_len.to_be_bytes());
|
|
self.process_block();
|
|
|
|
// Output
|
|
let mut result = [0u8; 32];
|
|
for (i, &h) in self.h.iter().enumerate() {
|
|
result[i * 4..(i + 1) * 4].copy_from_slice(&h.to_be_bytes());
|
|
}
|
|
result
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_info_creation() {
|
|
let model = ModelInfo::new("test-model", "Test Model")
|
|
.with_description("A test model")
|
|
.with_input_config(32, 500.0, 500);
|
|
|
|
assert_eq!(model.id, "test-model");
|
|
assert_eq!(model.name, "Test Model");
|
|
assert_eq!(model.n_channels, 32);
|
|
assert_eq!(model.sfreq, 500.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_source_url() {
|
|
let hf = ModelSource::huggingface("user/repo", "model.onnx");
|
|
assert_eq!(
|
|
hf.to_url(),
|
|
Some("https://huggingface.co/user/repo/resolve/main/model.onnx".to_string())
|
|
);
|
|
|
|
let gh = ModelSource::github("owner", "repo", "v1.0.0", "model.onnx");
|
|
assert_eq!(
|
|
gh.to_url(),
|
|
Some("https://github.com/owner/repo/releases/download/v1.0.0/model.onnx".to_string())
|
|
);
|
|
|
|
let local = ModelSource::local("/path/to/model.onnx");
|
|
assert!(local.to_url().is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_registry() {
|
|
let registry = ModelRegistry::new();
|
|
|
|
// Check built-in models are registered
|
|
assert!(registry.get("eeg-artifact-v1").is_some());
|
|
assert!(registry.get("meg-artifact-v1").is_some());
|
|
|
|
// List all models
|
|
let models = registry.list();
|
|
assert!(!models.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_supports() {
|
|
let model = ModelInfo::new("test", "Test")
|
|
.with_artifact_types(vec![ArtifactType::EyeBlink, ArtifactType::Muscle]);
|
|
|
|
assert!(model.supports(ArtifactType::EyeBlink));
|
|
assert!(model.supports(ArtifactType::Muscle));
|
|
assert!(!model.supports(ArtifactType::Heartbeat));
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_for_artifact() {
|
|
let registry = ModelRegistry::new();
|
|
|
|
let eye_models = registry.list_for_artifact(ArtifactType::EyeBlink);
|
|
assert!(!eye_models.is_empty());
|
|
|
|
// All eye models should support EyeBlink
|
|
for model in eye_models {
|
|
assert!(model.supports(ArtifactType::EyeBlink));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_cached_path() {
|
|
let registry = ModelRegistry::new();
|
|
let path = registry.cached_path("test-model");
|
|
assert!(path.to_string_lossy().contains("test-model.onnx"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_sha256() {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(b"hello world");
|
|
let result = format!("{:x}", hasher.finalize());
|
|
assert_eq!(
|
|
result,
|
|
"b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sha256_empty() {
|
|
let hasher = Sha256::new();
|
|
let result = format!("{:x}", hasher.finalize());
|
|
assert_eq!(
|
|
result,
|
|
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sha256_large() {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(&[0u8; 1000]);
|
|
let result = format!("{:x}", hasher.finalize());
|
|
// Just verify it completes without error and produces valid hex
|
|
assert_eq!(result.len(), 64);
|
|
}
|
|
}
|