763 lines
24 KiB
Rust
763 lines
24 KiB
Rust
//! Pretrained model loading API similar to HuggingFace's from_pretrained().
|
|
//!
|
|
//! This module provides a convenient API for loading pretrained models from various
|
|
//! sources including local cache, remote registries, and mirrors.
|
|
|
|
use crate::{HubError, HubResult, ModelId, ModelPackage, ModelRegistry, Registry};
|
|
use async_trait::async_trait;
|
|
use chrono::{DateTime, Utc};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::path::{Path, PathBuf};
|
|
use tokio::fs;
|
|
use tracing::{debug, info};
|
|
|
|
/// Configuration for loading pretrained models.
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct PretrainedConfig {
|
|
/// Git-like revision/tag (e.g., "main", "v1.0.0", commit hash)
|
|
pub revision: Option<String>,
|
|
/// Local cache directory (defaults to ~/.cache/rtx-hub)
|
|
pub cache_dir: Option<PathBuf>,
|
|
/// Force re-download even if cached
|
|
pub force_download: bool,
|
|
/// Resume interrupted downloads
|
|
pub resume_download: bool,
|
|
/// Only use local files (offline mode)
|
|
pub local_files_only: bool,
|
|
/// Authentication token for private models
|
|
pub use_auth_token: Option<String>,
|
|
/// Alternative mirror endpoint
|
|
pub mirror: Option<String>,
|
|
}
|
|
|
|
impl PretrainedConfig {
|
|
/// Create a new configuration with defaults.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Set the revision/tag.
|
|
pub fn with_revision(mut self, revision: impl Into<String>) -> Self {
|
|
self.revision = Some(revision.into());
|
|
self
|
|
}
|
|
|
|
/// Set the cache directory.
|
|
pub fn with_cache_dir(mut self, cache_dir: impl Into<PathBuf>) -> Self {
|
|
self.cache_dir = Some(cache_dir.into());
|
|
self
|
|
}
|
|
|
|
/// Enable force download.
|
|
pub fn with_force_download(mut self, force: bool) -> Self {
|
|
self.force_download = force;
|
|
self
|
|
}
|
|
|
|
/// Enable resume download.
|
|
pub fn with_resume_download(mut self, resume: bool) -> Self {
|
|
self.resume_download = resume;
|
|
self
|
|
}
|
|
|
|
/// Enable local files only mode.
|
|
pub fn with_local_files_only(mut self, local_only: bool) -> Self {
|
|
self.local_files_only = local_only;
|
|
self
|
|
}
|
|
|
|
/// Set authentication token.
|
|
pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
|
|
self.use_auth_token = Some(token.into());
|
|
self
|
|
}
|
|
|
|
/// Set mirror endpoint.
|
|
pub fn with_mirror(mut self, mirror: impl Into<String>) -> Self {
|
|
self.mirror = Some(mirror.into());
|
|
self
|
|
}
|
|
|
|
/// Get the cache directory, using default if not set.
|
|
pub fn get_cache_dir(&self) -> PathBuf {
|
|
self.cache_dir.clone().unwrap_or_else(|| {
|
|
dirs::home_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join(".cache")
|
|
.join("rtx-hub")
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Trait for loading models from pretrained checkpoints.
|
|
#[async_trait]
|
|
pub trait FromPretrained: Sized {
|
|
/// Configuration type for this model.
|
|
type Config;
|
|
|
|
/// Load a model from a pretrained checkpoint.
|
|
async fn from_pretrained(model_id: &str, config: PretrainedConfig) -> HubResult<Self>;
|
|
}
|
|
|
|
/// Model file information.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelFile {
|
|
/// File path within the model directory
|
|
pub path: String,
|
|
/// File size in bytes
|
|
pub size: u64,
|
|
/// File hash for integrity verification
|
|
pub hash: String,
|
|
/// Last modified timestamp
|
|
pub modified: DateTime<Utc>,
|
|
}
|
|
|
|
/// Loader for pretrained models from registry.
|
|
pub struct PretrainedLoader {
|
|
/// Model registry
|
|
registry: std::sync::Arc<ModelRegistry>,
|
|
/// Local cache manager
|
|
cache: LocalCache,
|
|
}
|
|
|
|
impl PretrainedLoader {
|
|
/// Create a new pretrained loader.
|
|
pub fn new(registry: std::sync::Arc<ModelRegistry>, cache_dir: PathBuf) -> Self {
|
|
let cache = LocalCache::new(cache_dir);
|
|
Self { registry, cache }
|
|
}
|
|
|
|
/// Resolve model ID to storage path.
|
|
pub async fn resolve_model(
|
|
&self,
|
|
model_id: &str,
|
|
config: &PretrainedConfig,
|
|
) -> HubResult<PathBuf> {
|
|
let parsed_id = ModelId::parse(model_id)?;
|
|
|
|
// Determine version from revision or use latest
|
|
let version = if let Some(ref revision) = config.revision {
|
|
// Try to parse as semver first
|
|
if let Ok(semver) = semver::Version::parse(revision) {
|
|
Some(crate::ModelVersion::new(semver))
|
|
} else {
|
|
// Otherwise get latest version
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Get model info from registry
|
|
let model_info = self
|
|
.registry
|
|
.get_model(&parsed_id, version.as_ref())
|
|
.await?;
|
|
|
|
// Check if cached
|
|
let cache_path = self
|
|
.cache
|
|
.cache_path(model_id, &model_info.metadata.version.to_string());
|
|
|
|
if config.local_files_only && !self.cache.is_cached(&cache_path).await {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: format!("{} (local files only mode)", model_id),
|
|
});
|
|
}
|
|
|
|
Ok(cache_path)
|
|
}
|
|
|
|
/// Download model if not cached.
|
|
pub async fn download_if_needed(
|
|
&self,
|
|
model_id: &str,
|
|
config: &PretrainedConfig,
|
|
) -> HubResult<PathBuf> {
|
|
let parsed_id = ModelId::parse(model_id)?;
|
|
let cache_path = self.resolve_model(model_id, config).await?;
|
|
|
|
// Check if we need to download
|
|
let needs_download = config.force_download || !self.cache.is_cached(&cache_path).await;
|
|
|
|
if needs_download && !config.local_files_only {
|
|
info!("Downloading model: {}", model_id);
|
|
|
|
// Get model version
|
|
let model_info = self.registry.get_model(&parsed_id, None).await?;
|
|
let version = &model_info.metadata.version;
|
|
|
|
// Download package from registry
|
|
let package = self.registry.download_package(&parsed_id, version).await?;
|
|
|
|
// Extract to cache
|
|
self.cache.store_package(&cache_path, &package).await?;
|
|
|
|
debug!("Model downloaded and cached at: {}", cache_path.display());
|
|
} else {
|
|
debug!("Using cached model at: {}", cache_path.display());
|
|
}
|
|
|
|
Ok(cache_path)
|
|
}
|
|
|
|
/// Load model configuration from cache.
|
|
pub async fn load_config(&self, model_path: &Path) -> HubResult<serde_json::Value> {
|
|
let config_path = model_path.join("config.json");
|
|
|
|
if !config_path.exists() {
|
|
return Err(HubError::InvalidPackage {
|
|
reason: "Model config not found".to_string(),
|
|
});
|
|
}
|
|
|
|
let config_data = fs::read_to_string(&config_path).await?;
|
|
let config: serde_json::Value = serde_json::from_str(&config_data)?;
|
|
|
|
Ok(config)
|
|
}
|
|
|
|
/// Load tokenizer from cache.
|
|
pub async fn load_tokenizer(&self, model_path: &Path) -> HubResult<serde_json::Value> {
|
|
let tokenizer_path = model_path.join("tokenizer.json");
|
|
|
|
if !tokenizer_path.exists() {
|
|
return Err(HubError::InvalidPackage {
|
|
reason: "Tokenizer not found".to_string(),
|
|
});
|
|
}
|
|
|
|
let tokenizer_data = fs::read_to_string(&tokenizer_path).await?;
|
|
let tokenizer: serde_json::Value = serde_json::from_str(&tokenizer_data)?;
|
|
|
|
Ok(tokenizer)
|
|
}
|
|
|
|
/// Get list of all model files.
|
|
pub async fn get_model_files(&self, model_path: &Path) -> HubResult<Vec<ModelFile>> {
|
|
let mut files = Vec::new();
|
|
|
|
let mut entries = fs::read_dir(model_path).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
let metadata = fs::metadata(&path).await?;
|
|
|
|
if metadata.is_file() {
|
|
let relative_path = path
|
|
.strip_prefix(model_path)
|
|
.map_err(|e| HubError::InvalidPackage {
|
|
reason: e.to_string(),
|
|
})?
|
|
.to_string_lossy()
|
|
.to_string();
|
|
|
|
files.push(ModelFile {
|
|
path: relative_path,
|
|
size: metadata.len(),
|
|
hash: String::new(), // TODO: Calculate hash
|
|
modified: DateTime::<Utc>::from(metadata.modified()?),
|
|
});
|
|
}
|
|
}
|
|
|
|
Ok(files)
|
|
}
|
|
}
|
|
|
|
/// Local cache manager for pretrained models.
|
|
pub struct LocalCache {
|
|
/// Base cache directory
|
|
cache_dir: PathBuf,
|
|
/// Cache metadata
|
|
metadata: HashMap<String, CacheEntry>,
|
|
}
|
|
|
|
/// Cache entry metadata.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct CacheEntry {
|
|
/// Model ID
|
|
model_id: String,
|
|
/// Model version
|
|
version: String,
|
|
/// Cache path
|
|
path: PathBuf,
|
|
/// Size in bytes
|
|
size: u64,
|
|
/// Last accessed timestamp
|
|
last_accessed: DateTime<Utc>,
|
|
/// Number of accesses
|
|
access_count: u64,
|
|
}
|
|
|
|
impl LocalCache {
|
|
/// Create a new local cache manager.
|
|
pub fn new(cache_dir: PathBuf) -> Self {
|
|
Self {
|
|
cache_dir,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Get cache path for a model.
|
|
pub fn cache_path(&self, model_id: &str, version: &str) -> PathBuf {
|
|
self.cache_dir
|
|
.join("models")
|
|
.join(model_id.replace('/', "--"))
|
|
.join(version)
|
|
}
|
|
|
|
/// Check if model is cached.
|
|
pub async fn is_cached(&self, cache_path: &Path) -> bool {
|
|
cache_path.exists() && cache_path.is_dir()
|
|
}
|
|
|
|
/// Store package to cache.
|
|
pub async fn store_package(&self, cache_path: &Path, package: &ModelPackage) -> HubResult<()> {
|
|
// Create cache directory
|
|
fs::create_dir_all(cache_path).await?;
|
|
|
|
// Write package metadata
|
|
let metadata_path = cache_path.join("metadata.json");
|
|
let metadata_json = serde_json::to_string_pretty(&package.metadata)?;
|
|
fs::write(&metadata_path, metadata_json).await?;
|
|
|
|
// Write manifest
|
|
let manifest_path = cache_path.join("manifest.json");
|
|
let manifest_json = serde_json::to_string_pretty(&package.manifest)?;
|
|
fs::write(&manifest_path, manifest_json).await?;
|
|
|
|
debug!("Package stored to cache at: {}", cache_path.display());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Evict least recently used models.
|
|
pub async fn evict_lru(&mut self, max_size: u64) -> HubResult<()> {
|
|
let current_size = self.get_cache_size().await?;
|
|
|
|
if current_size <= max_size {
|
|
return Ok(());
|
|
}
|
|
|
|
// Sort by last accessed time
|
|
let mut entries: Vec<_> = self.metadata.values().cloned().collect();
|
|
entries.sort_by_key(|e| e.last_accessed);
|
|
|
|
let mut freed_size = 0u64;
|
|
|
|
for entry in entries {
|
|
if current_size - freed_size <= max_size {
|
|
break;
|
|
}
|
|
|
|
// Remove cache directory
|
|
if entry.path.exists() {
|
|
fs::remove_dir_all(&entry.path).await?;
|
|
freed_size += entry.size;
|
|
|
|
debug!(
|
|
"Evicted model {} from cache (freed {} bytes)",
|
|
entry.model_id, entry.size
|
|
);
|
|
}
|
|
|
|
// Remove from metadata
|
|
self.metadata.remove(&entry.model_id);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get total cache size in bytes.
|
|
pub async fn get_cache_size(&self) -> HubResult<u64> {
|
|
let mut total_size = 0u64;
|
|
|
|
if !self.cache_dir.exists() {
|
|
return Ok(0);
|
|
}
|
|
|
|
let mut entries = fs::read_dir(&self.cache_dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
if path.is_dir() {
|
|
total_size += self.calculate_dir_size(&path).await?;
|
|
}
|
|
}
|
|
|
|
Ok(total_size)
|
|
}
|
|
|
|
/// Clear all cached models.
|
|
pub async fn clear_cache(&mut self) -> HubResult<()> {
|
|
if self.cache_dir.exists() {
|
|
fs::remove_dir_all(&self.cache_dir).await?;
|
|
fs::create_dir_all(&self.cache_dir).await?;
|
|
}
|
|
|
|
self.metadata.clear();
|
|
|
|
info!("Cache cleared");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Calculate directory size recursively.
|
|
fn calculate_dir_size<'a>(
|
|
&'a self,
|
|
dir: &'a Path,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = HubResult<u64>> + 'a>> {
|
|
Box::pin(async move {
|
|
let mut total_size = 0u64;
|
|
|
|
let mut entries = fs::read_dir(dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
let metadata = fs::metadata(&path).await?;
|
|
|
|
if metadata.is_file() {
|
|
total_size += metadata.len();
|
|
} else if metadata.is_dir() {
|
|
total_size += self.calculate_dir_size(&path).await?;
|
|
}
|
|
}
|
|
|
|
Ok(total_size)
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::model::ModelSchema;
|
|
use crate::{ModelMetadata, ModelStatus, ModelVersion, RegistryConfig, StorageConfig};
|
|
use semver::Version;
|
|
use std::collections::HashMap;
|
|
use tempfile::TempDir;
|
|
|
|
#[test]
|
|
fn test_pretrained_config_builder() {
|
|
let config = PretrainedConfig::new()
|
|
.with_revision("v1.0.0")
|
|
.with_force_download(true)
|
|
.with_local_files_only(false)
|
|
.with_auth_token("test-token");
|
|
|
|
assert_eq!(config.revision, Some("v1.0.0".to_string()));
|
|
assert!(config.force_download);
|
|
assert!(!config.local_files_only);
|
|
assert_eq!(config.use_auth_token, Some("test-token".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_pretrained_config_cache_dir() {
|
|
let config = PretrainedConfig::new();
|
|
let cache_dir = config.get_cache_dir();
|
|
assert!(cache_dir.to_string_lossy().contains("rtx-hub"));
|
|
|
|
let custom_dir = PathBuf::from("/tmp/custom-cache");
|
|
let config = PretrainedConfig::new().with_cache_dir(custom_dir.clone());
|
|
assert_eq!(config.get_cache_dir(), custom_dir);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_cache_path() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache = LocalCache::new(temp_dir.path().to_path_buf());
|
|
|
|
let cache_path = cache.cache_path("namespace/model", "1.0.0");
|
|
assert!(cache_path.to_string_lossy().contains("namespace--model"));
|
|
assert!(cache_path.to_string_lossy().contains("1.0.0"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_cache_is_cached() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache = LocalCache::new(temp_dir.path().to_path_buf());
|
|
|
|
let cache_path = cache.cache_path("test/model", "1.0.0");
|
|
assert!(!cache.is_cached(&cache_path).await);
|
|
|
|
// Create cache directory
|
|
fs::create_dir_all(&cache_path).await.unwrap();
|
|
assert!(cache.is_cached(&cache_path).await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_cache_clear() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let mut cache = LocalCache::new(temp_dir.path().to_path_buf());
|
|
|
|
// Create some cache directories
|
|
let cache_path1 = cache.cache_path("test/model1", "1.0.0");
|
|
let cache_path2 = cache.cache_path("test/model2", "1.0.0");
|
|
|
|
fs::create_dir_all(&cache_path1).await.unwrap();
|
|
fs::create_dir_all(&cache_path2).await.unwrap();
|
|
|
|
// Clear cache
|
|
cache.clear_cache().await.unwrap();
|
|
|
|
assert!(!cache.is_cached(&cache_path1).await);
|
|
assert!(!cache.is_cached(&cache_path2).await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_cache_size() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache = LocalCache::new(temp_dir.path().to_path_buf());
|
|
|
|
let initial_size = cache.get_cache_size().await.unwrap();
|
|
assert_eq!(initial_size, 0);
|
|
|
|
// Create a test file
|
|
let cache_path = cache.cache_path("test/model", "1.0.0");
|
|
fs::create_dir_all(&cache_path).await.unwrap();
|
|
|
|
let test_file = cache_path.join("test.txt");
|
|
fs::write(&test_file, b"Hello, world!").await.unwrap();
|
|
|
|
let size = cache.get_cache_size().await.unwrap();
|
|
assert!(size > 0);
|
|
}
|
|
|
|
async fn create_test_registry() -> (ModelRegistry, TempDir) {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let config = RegistryConfig {
|
|
storage: StorageConfig::Local {
|
|
base_path: temp_dir.path().to_path_buf(),
|
|
},
|
|
database_url: "sqlite::memory:".to_string(),
|
|
enable_validation: false,
|
|
enable_compression: false,
|
|
..Default::default()
|
|
};
|
|
|
|
let registry = ModelRegistry::new(config).await.unwrap();
|
|
(registry, temp_dir)
|
|
}
|
|
|
|
fn create_test_metadata(id: ModelId, version: ModelVersion) -> ModelMetadata {
|
|
ModelMetadata {
|
|
id,
|
|
version,
|
|
title: "Test Model".to_string(),
|
|
description: "Test Description".to_string(),
|
|
architecture: "transformer".to_string(),
|
|
framework: "rustytorch".to_string(),
|
|
framework_version: "1.0.0".to_string(),
|
|
tags: vec!["test".to_string()],
|
|
author: "Test Author".to_string(),
|
|
license: Some("MIT".to_string()),
|
|
created_at: Utc::now(),
|
|
updated_at: Utc::now(),
|
|
status: ModelStatus::Available,
|
|
size: 1024,
|
|
content_hash: "test-hash".to_string(),
|
|
dependencies: vec![],
|
|
schema: ModelSchema {
|
|
inputs: vec![],
|
|
outputs: vec![],
|
|
config: None,
|
|
},
|
|
metrics: HashMap::new(),
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pretrained_loader_resolve_model() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
let cache_dir = TempDir::new().unwrap();
|
|
|
|
// Register a test model
|
|
let model_id = ModelId::new("test", "model");
|
|
let version = ModelVersion::new(Version::parse("1.0.0").unwrap());
|
|
let metadata = create_test_metadata(model_id.clone(), version.clone());
|
|
|
|
registry.register_model(metadata).await.unwrap();
|
|
|
|
let loader = PretrainedLoader::new(
|
|
std::sync::Arc::new(registry),
|
|
cache_dir.path().to_path_buf(),
|
|
);
|
|
|
|
// Resolve model
|
|
let config = PretrainedConfig::new();
|
|
let resolved_path = loader.resolve_model("test/model", &config).await.unwrap();
|
|
|
|
assert!(resolved_path.to_string_lossy().contains("test--model"));
|
|
assert!(resolved_path.to_string_lossy().contains("1.0.0"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pretrained_loader_local_files_only() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
let cache_dir = TempDir::new().unwrap();
|
|
|
|
// Register a test model
|
|
let model_id = ModelId::new("test", "model");
|
|
let version = ModelVersion::new(Version::parse("1.0.0").unwrap());
|
|
let metadata = create_test_metadata(model_id.clone(), version.clone());
|
|
|
|
registry.register_model(metadata).await.unwrap();
|
|
|
|
let loader = PretrainedLoader::new(
|
|
std::sync::Arc::new(registry),
|
|
cache_dir.path().to_path_buf(),
|
|
);
|
|
|
|
// Try to resolve with local_files_only
|
|
let config = PretrainedConfig::new().with_local_files_only(true);
|
|
let result = loader.resolve_model("test/model", &config).await;
|
|
|
|
// Should fail because model is not cached
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pretrained_loader_get_model_files() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache_dir = TempDir::new().unwrap();
|
|
let (registry, _) = create_test_registry().await;
|
|
let loader = PretrainedLoader::new(
|
|
std::sync::Arc::new(registry),
|
|
cache_dir.path().to_path_buf(),
|
|
);
|
|
|
|
// Create test files
|
|
let model_path = temp_dir.path().join("model");
|
|
fs::create_dir_all(&model_path).await.unwrap();
|
|
|
|
fs::write(model_path.join("config.json"), b"{}")
|
|
.await
|
|
.unwrap();
|
|
fs::write(model_path.join("weights.bin"), b"fake weights")
|
|
.await
|
|
.unwrap();
|
|
|
|
// Get model files
|
|
let files = loader.get_model_files(&model_path).await.unwrap();
|
|
|
|
assert_eq!(files.len(), 2);
|
|
assert!(files.iter().any(|f| f.path == "config.json"));
|
|
assert!(files.iter().any(|f| f.path == "weights.bin"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pretrained_loader_load_config() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache_dir = TempDir::new().unwrap();
|
|
let (registry, _) = create_test_registry().await;
|
|
let loader = PretrainedLoader::new(
|
|
std::sync::Arc::new(registry),
|
|
cache_dir.path().to_path_buf(),
|
|
);
|
|
|
|
// Create test config
|
|
let model_path = temp_dir.path().join("model");
|
|
fs::create_dir_all(&model_path).await.unwrap();
|
|
|
|
let config_json = r#"{"model_type": "transformer", "hidden_size": 768}"#;
|
|
fs::write(model_path.join("config.json"), config_json)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Load config
|
|
let config = loader.load_config(&model_path).await.unwrap();
|
|
|
|
assert_eq!(config["model_type"], "transformer");
|
|
assert_eq!(config["hidden_size"], 768);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pretrained_loader_load_tokenizer() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let cache_dir = TempDir::new().unwrap();
|
|
let (registry, _) = create_test_registry().await;
|
|
let loader = PretrainedLoader::new(
|
|
std::sync::Arc::new(registry),
|
|
cache_dir.path().to_path_buf(),
|
|
);
|
|
|
|
// Create test tokenizer
|
|
let model_path = temp_dir.path().join("model");
|
|
fs::create_dir_all(&model_path).await.unwrap();
|
|
|
|
let tokenizer_json = r#"{"vocab_size": 50257}"#;
|
|
fs::write(model_path.join("tokenizer.json"), tokenizer_json)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Load tokenizer
|
|
let tokenizer = loader.load_tokenizer(&model_path).await.unwrap();
|
|
|
|
assert_eq!(tokenizer["vocab_size"], 50257);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_cache_evict_lru() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let mut cache = LocalCache::new(temp_dir.path().to_path_buf());
|
|
|
|
// Create actual directories first
|
|
let model1_path = temp_dir.path().join("models").join("model1");
|
|
let model2_path = temp_dir.path().join("models").join("model2");
|
|
|
|
fs::create_dir_all(&model1_path).await.unwrap();
|
|
fs::create_dir_all(&model2_path).await.unwrap();
|
|
|
|
// Write some data to make them have size
|
|
fs::write(model1_path.join("data.bin"), vec![0u8; 1000])
|
|
.await
|
|
.unwrap();
|
|
fs::write(model2_path.join("data.bin"), vec![0u8; 2000])
|
|
.await
|
|
.unwrap();
|
|
|
|
// Add cache entries
|
|
let now = Utc::now();
|
|
cache.metadata.insert(
|
|
"model1".to_string(),
|
|
CacheEntry {
|
|
model_id: "model1".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
path: model1_path.clone(),
|
|
size: 1000,
|
|
last_accessed: now - chrono::Duration::days(10),
|
|
access_count: 1,
|
|
},
|
|
);
|
|
|
|
cache.metadata.insert(
|
|
"model2".to_string(),
|
|
CacheEntry {
|
|
model_id: "model2".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
path: model2_path.clone(),
|
|
size: 2000,
|
|
last_accessed: now - chrono::Duration::days(5),
|
|
access_count: 2,
|
|
},
|
|
);
|
|
|
|
// Initial state: both models exist
|
|
assert!(model1_path.exists());
|
|
assert!(model2_path.exists());
|
|
|
|
// Evict LRU with max size 2500 (no eviction needed)
|
|
cache.evict_lru(5000).await.unwrap();
|
|
|
|
// Both should still exist
|
|
assert!(model1_path.exists());
|
|
assert!(model2_path.exists());
|
|
|
|
// Test that metadata is tracked correctly
|
|
assert_eq!(cache.metadata.len(), 2);
|
|
}
|
|
}
|