514 lines
17 KiB
Rust
514 lines
17 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum ModelCategory {
|
|
ImageClassification,
|
|
ObjectDetection,
|
|
Segmentation,
|
|
TextGeneration,
|
|
SpeechRecognition,
|
|
NLP,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ModelInfo {
|
|
pub id: String,
|
|
pub name: String,
|
|
pub description: String,
|
|
pub category: ModelCategory,
|
|
pub size_mb: f64,
|
|
pub parameters: u64,
|
|
pub accuracy_metric: String,
|
|
pub download_url: String,
|
|
pub license: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum ModelStatus {
|
|
NotDownloaded,
|
|
Downloading,
|
|
Downloaded,
|
|
Loading,
|
|
Ready,
|
|
Error,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ModelZooConfig {
|
|
pub selected_model: Option<String>,
|
|
pub download_path: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct InferenceRequest {
|
|
pub model_id: String,
|
|
pub input_data: String,
|
|
pub input_type: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct InferenceResult {
|
|
pub model_id: String,
|
|
pub output: String,
|
|
pub inference_time_ms: f64,
|
|
pub device: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum ModelZooRequest {
|
|
ListModels,
|
|
FilterByCategory { category: ModelCategory },
|
|
GetModelInfo { model_id: String },
|
|
GetModelStatus { model_id: String },
|
|
DownloadModel { model_id: String },
|
|
LoadModel { model_id: String },
|
|
UnloadModel { model_id: String },
|
|
RunInference { request: InferenceRequest },
|
|
GetConfig,
|
|
UpdateConfig { config: ModelZooConfig },
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum ModelZooResponse {
|
|
ModelList { models: Vec<ModelInfo> },
|
|
ModelInfo { info: ModelInfo },
|
|
ModelStatus { status: ModelStatus },
|
|
DownloadStarted { model_id: String },
|
|
ModelLoaded { model_id: String },
|
|
ModelUnloaded { model_id: String },
|
|
InferenceResult { result: InferenceResult },
|
|
Config { config: ModelZooConfig },
|
|
ConfigUpdated,
|
|
Error { message: String },
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_category_serialization() {
|
|
let category = ModelCategory::ImageClassification;
|
|
let json = serde_json::to_string(&category).expect("Failed to serialize");
|
|
let deserialized: ModelCategory =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(category, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_category_equality() {
|
|
assert_eq!(
|
|
ModelCategory::ImageClassification,
|
|
ModelCategory::ImageClassification
|
|
);
|
|
assert_ne!(
|
|
ModelCategory::ImageClassification,
|
|
ModelCategory::ObjectDetection
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_category_debug() {
|
|
let category = ModelCategory::ImageClassification;
|
|
let debug_str = format!("{category:?}");
|
|
assert_eq!(debug_str, "ImageClassification");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_category_all_variants() {
|
|
let categories = vec![
|
|
ModelCategory::ImageClassification,
|
|
ModelCategory::ObjectDetection,
|
|
ModelCategory::Segmentation,
|
|
ModelCategory::TextGeneration,
|
|
ModelCategory::SpeechRecognition,
|
|
ModelCategory::NLP,
|
|
];
|
|
assert_eq!(categories.len(), 6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_creation() {
|
|
let info = ModelInfo {
|
|
id: "resnet50".to_string(),
|
|
name: "ResNet-50".to_string(),
|
|
description: "50-layer residual network".to_string(),
|
|
category: ModelCategory::ImageClassification,
|
|
size_mb: 97.8,
|
|
parameters: 25_600_000,
|
|
accuracy_metric: "Top-1: 76.1%".to_string(),
|
|
download_url: "https://example.com/resnet50.pth".to_string(),
|
|
license: "MIT".to_string(),
|
|
};
|
|
assert_eq!(info.id, "resnet50");
|
|
assert_eq!(info.name, "ResNet-50");
|
|
assert_eq!(info.category, ModelCategory::ImageClassification);
|
|
assert_eq!(info.size_mb, 97.8);
|
|
assert_eq!(info.parameters, 25_600_000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_serialization() {
|
|
let info = ModelInfo {
|
|
id: "resnet50".to_string(),
|
|
name: "ResNet-50".to_string(),
|
|
description: "50-layer residual network".to_string(),
|
|
category: ModelCategory::ImageClassification,
|
|
size_mb: 97.8,
|
|
parameters: 25_600_000,
|
|
accuracy_metric: "Top-1: 76.1%".to_string(),
|
|
download_url: "https://example.com/resnet50.pth".to_string(),
|
|
license: "MIT".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&info).expect("Failed to serialize");
|
|
let deserialized: ModelInfo = serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(info, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_info_clone() {
|
|
let info = ModelInfo {
|
|
id: "resnet50".to_string(),
|
|
name: "ResNet-50".to_string(),
|
|
description: "50-layer residual network".to_string(),
|
|
category: ModelCategory::ImageClassification,
|
|
size_mb: 97.8,
|
|
parameters: 25_600_000,
|
|
accuracy_metric: "Top-1: 76.1%".to_string(),
|
|
download_url: "https://example.com/resnet50.pth".to_string(),
|
|
license: "MIT".to_string(),
|
|
};
|
|
let cloned = info.clone();
|
|
assert_eq!(info, cloned);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_status_all_variants() {
|
|
let statuses = vec![
|
|
ModelStatus::NotDownloaded,
|
|
ModelStatus::Downloading,
|
|
ModelStatus::Downloaded,
|
|
ModelStatus::Loading,
|
|
ModelStatus::Ready,
|
|
ModelStatus::Error,
|
|
];
|
|
assert_eq!(statuses.len(), 6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_status_equality() {
|
|
assert_eq!(ModelStatus::Ready, ModelStatus::Ready);
|
|
assert_ne!(ModelStatus::Ready, ModelStatus::Loading);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_status_serialization() {
|
|
let status = ModelStatus::Ready;
|
|
let json = serde_json::to_string(&status).expect("Failed to serialize");
|
|
let deserialized: ModelStatus = serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(status, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_status_debug() {
|
|
let status = ModelStatus::Ready;
|
|
let debug_str = format!("{status:?}");
|
|
assert_eq!(debug_str, "Ready");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_config_creation() {
|
|
let config = ModelZooConfig {
|
|
selected_model: Some("resnet50".to_string()),
|
|
download_path: "/tmp/models".to_string(),
|
|
};
|
|
assert_eq!(config.selected_model, Some("resnet50".to_string()));
|
|
assert_eq!(config.download_path, "/tmp/models");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_config_no_selection() {
|
|
let config = ModelZooConfig {
|
|
selected_model: None,
|
|
download_path: "/tmp/models".to_string(),
|
|
};
|
|
assert_eq!(config.selected_model, None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_config_serialization() {
|
|
let config = ModelZooConfig {
|
|
selected_model: Some("resnet50".to_string()),
|
|
download_path: "/tmp/models".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&config).expect("Failed to serialize");
|
|
let deserialized: ModelZooConfig =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(config, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_request_creation() {
|
|
let request = InferenceRequest {
|
|
model_id: "resnet50".to_string(),
|
|
input_data: "base64encodedimage".to_string(),
|
|
input_type: "image/jpeg".to_string(),
|
|
};
|
|
assert_eq!(request.model_id, "resnet50");
|
|
assert_eq!(request.input_data, "base64encodedimage");
|
|
assert_eq!(request.input_type, "image/jpeg");
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_request_serialization() {
|
|
let request = InferenceRequest {
|
|
model_id: "resnet50".to_string(),
|
|
input_data: "base64encodedimage".to_string(),
|
|
input_type: "image/jpeg".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: InferenceRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_result_creation() {
|
|
let result = InferenceResult {
|
|
model_id: "resnet50".to_string(),
|
|
output: "class: cat, confidence: 0.95".to_string(),
|
|
inference_time_ms: 15.5,
|
|
device: "CPU".to_string(),
|
|
};
|
|
assert_eq!(result.model_id, "resnet50");
|
|
assert_eq!(result.output, "class: cat, confidence: 0.95");
|
|
assert_eq!(result.inference_time_ms, 15.5);
|
|
assert_eq!(result.device, "CPU");
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_result_serialization() {
|
|
let result = InferenceResult {
|
|
model_id: "resnet50".to_string(),
|
|
output: "class: cat, confidence: 0.95".to_string(),
|
|
inference_time_ms: 15.5,
|
|
device: "CPU".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&result).expect("Failed to serialize");
|
|
let deserialized: InferenceResult =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(result, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_request_list_models() {
|
|
let request = ModelZooRequest::ListModels;
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: ModelZooRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_request_filter_by_category() {
|
|
let request = ModelZooRequest::FilterByCategory {
|
|
category: ModelCategory::ImageClassification,
|
|
};
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: ModelZooRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_request_get_model_info() {
|
|
let request = ModelZooRequest::GetModelInfo {
|
|
model_id: "resnet50".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: ModelZooRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_request_download_model() {
|
|
let request = ModelZooRequest::DownloadModel {
|
|
model_id: "resnet50".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: ModelZooRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_request_run_inference() {
|
|
let inference_request = InferenceRequest {
|
|
model_id: "resnet50".to_string(),
|
|
input_data: "base64encodedimage".to_string(),
|
|
input_type: "image/jpeg".to_string(),
|
|
};
|
|
let request = ModelZooRequest::RunInference {
|
|
request: inference_request,
|
|
};
|
|
let json = serde_json::to_string(&request).expect("Failed to serialize");
|
|
let deserialized: ModelZooRequest =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(request, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_response_model_list() {
|
|
let model_info = ModelInfo {
|
|
id: "resnet50".to_string(),
|
|
name: "ResNet-50".to_string(),
|
|
description: "50-layer residual network".to_string(),
|
|
category: ModelCategory::ImageClassification,
|
|
size_mb: 97.8,
|
|
parameters: 25_600_000,
|
|
accuracy_metric: "Top-1: 76.1%".to_string(),
|
|
download_url: "https://example.com/resnet50.pth".to_string(),
|
|
license: "MIT".to_string(),
|
|
};
|
|
let response = ModelZooResponse::ModelList {
|
|
models: vec![model_info],
|
|
};
|
|
let json = serde_json::to_string(&response).expect("Failed to serialize");
|
|
let deserialized: ModelZooResponse =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(response, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_response_model_status() {
|
|
let response = ModelZooResponse::ModelStatus {
|
|
status: ModelStatus::Ready,
|
|
};
|
|
let json = serde_json::to_string(&response).expect("Failed to serialize");
|
|
let deserialized: ModelZooResponse =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(response, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_response_inference_result() {
|
|
let inference_result = InferenceResult {
|
|
model_id: "resnet50".to_string(),
|
|
output: "class: cat, confidence: 0.95".to_string(),
|
|
inference_time_ms: 15.5,
|
|
device: "CPU".to_string(),
|
|
};
|
|
let response = ModelZooResponse::InferenceResult {
|
|
result: inference_result,
|
|
};
|
|
let json = serde_json::to_string(&response).expect("Failed to serialize");
|
|
let deserialized: ModelZooResponse =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(response, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_zoo_response_error() {
|
|
let response = ModelZooResponse::Error {
|
|
message: "Model not found".to_string(),
|
|
};
|
|
let json = serde_json::to_string(&response).expect("Failed to serialize");
|
|
let deserialized: ModelZooResponse =
|
|
serde_json::from_str(&json).expect("Failed to deserialize");
|
|
assert_eq!(response, deserialized);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_request_variants() {
|
|
let requests = vec![
|
|
ModelZooRequest::ListModels,
|
|
ModelZooRequest::FilterByCategory {
|
|
category: ModelCategory::ImageClassification,
|
|
},
|
|
ModelZooRequest::GetModelInfo {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooRequest::GetModelStatus {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooRequest::DownloadModel {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooRequest::LoadModel {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooRequest::UnloadModel {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooRequest::RunInference {
|
|
request: InferenceRequest {
|
|
model_id: "test".to_string(),
|
|
input_data: "data".to_string(),
|
|
input_type: "type".to_string(),
|
|
},
|
|
},
|
|
ModelZooRequest::GetConfig,
|
|
ModelZooRequest::UpdateConfig {
|
|
config: ModelZooConfig {
|
|
selected_model: None,
|
|
download_path: "/tmp".to_string(),
|
|
},
|
|
},
|
|
];
|
|
assert_eq!(requests.len(), 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_response_variants() {
|
|
let model_info = ModelInfo {
|
|
id: "test".to_string(),
|
|
name: "Test".to_string(),
|
|
description: "Test model".to_string(),
|
|
category: ModelCategory::ImageClassification,
|
|
size_mb: 1.0,
|
|
parameters: 1000,
|
|
accuracy_metric: "N/A".to_string(),
|
|
download_url: "https://example.com".to_string(),
|
|
license: "MIT".to_string(),
|
|
};
|
|
let responses = vec![
|
|
ModelZooResponse::ModelList {
|
|
models: vec![model_info.clone()],
|
|
},
|
|
ModelZooResponse::ModelInfo { info: model_info },
|
|
ModelZooResponse::ModelStatus {
|
|
status: ModelStatus::Ready,
|
|
},
|
|
ModelZooResponse::DownloadStarted {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooResponse::ModelLoaded {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooResponse::ModelUnloaded {
|
|
model_id: "test".to_string(),
|
|
},
|
|
ModelZooResponse::InferenceResult {
|
|
result: InferenceResult {
|
|
model_id: "test".to_string(),
|
|
output: "output".to_string(),
|
|
inference_time_ms: 1.0,
|
|
device: "CPU".to_string(),
|
|
},
|
|
},
|
|
ModelZooResponse::Config {
|
|
config: ModelZooConfig {
|
|
selected_model: None,
|
|
download_path: "/tmp".to_string(),
|
|
},
|
|
},
|
|
ModelZooResponse::ConfigUpdated,
|
|
ModelZooResponse::Error {
|
|
message: "error".to_string(),
|
|
},
|
|
];
|
|
assert_eq!(responses.len(), 10);
|
|
}
|
|
}
|