976 lines
29 KiB
Rust
976 lines
29 KiB
Rust
//! Model search and discovery functionality.
|
|
//!
|
|
//! This module provides comprehensive search and discovery capabilities for models
|
|
//! in the registry, including filtering, sorting, trending analysis, and similarity search.
|
|
|
|
use crate::{HubResult, ModelId, ModelInfo, ModelMetadata, ModelRegistry, QueryOptions, Registry};
|
|
use chrono::{DateTime, Utc};
|
|
use dashmap::DashMap;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use tracing::debug;
|
|
|
|
/// Sort options for search results.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
|
|
pub enum SortBy {
|
|
/// Sort by total downloads
|
|
Downloads,
|
|
/// Sort by number of likes
|
|
Likes,
|
|
/// Sort by last modified date (most recent first)
|
|
Recent,
|
|
/// Sort by trending score
|
|
Trending,
|
|
/// Sort by relevance (for text search)
|
|
Relevance,
|
|
}
|
|
|
|
/// Search query for model discovery.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SearchQuery {
|
|
/// Text search query (searches title, description, tags)
|
|
pub query: Option<String>,
|
|
/// Filter by task type (e.g., "text-classification", "image-generation")
|
|
pub filter_by_task: Option<String>,
|
|
/// Filter by library/framework
|
|
pub filter_by_library: Option<String>,
|
|
/// Filter by supported languages
|
|
pub filter_by_language: Option<Vec<String>>,
|
|
/// Filter by license
|
|
pub filter_by_license: Option<Vec<String>>,
|
|
/// Filter by author/organization
|
|
pub filter_by_author: Option<String>,
|
|
/// Sort by field
|
|
pub sort_by: SortBy,
|
|
/// Maximum number of results
|
|
pub limit: usize,
|
|
/// Offset for pagination
|
|
pub offset: usize,
|
|
}
|
|
|
|
impl Default for SearchQuery {
|
|
fn default() -> Self {
|
|
Self {
|
|
query: None,
|
|
filter_by_task: None,
|
|
filter_by_library: None,
|
|
filter_by_language: None,
|
|
filter_by_license: None,
|
|
filter_by_author: None,
|
|
sort_by: SortBy::Relevance,
|
|
limit: 20,
|
|
offset: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SearchQuery {
|
|
/// Create a new search query.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Set text query.
|
|
pub fn with_query(mut self, query: impl Into<String>) -> Self {
|
|
self.query = Some(query.into());
|
|
self
|
|
}
|
|
|
|
/// Filter by task.
|
|
pub fn with_task(mut self, task: impl Into<String>) -> Self {
|
|
self.filter_by_task = Some(task.into());
|
|
self
|
|
}
|
|
|
|
/// Filter by library.
|
|
pub fn with_library(mut self, library: impl Into<String>) -> Self {
|
|
self.filter_by_library = Some(library.into());
|
|
self
|
|
}
|
|
|
|
/// Filter by languages.
|
|
pub fn with_languages(mut self, languages: Vec<String>) -> Self {
|
|
self.filter_by_language = Some(languages);
|
|
self
|
|
}
|
|
|
|
/// Filter by licenses.
|
|
pub fn with_licenses(mut self, licenses: Vec<String>) -> Self {
|
|
self.filter_by_license = Some(licenses);
|
|
self
|
|
}
|
|
|
|
/// Filter by author.
|
|
pub fn with_author(mut self, author: impl Into<String>) -> Self {
|
|
self.filter_by_author = Some(author.into());
|
|
self
|
|
}
|
|
|
|
/// Set sort order.
|
|
pub fn with_sort_by(mut self, sort_by: SortBy) -> Self {
|
|
self.sort_by = sort_by;
|
|
self
|
|
}
|
|
|
|
/// Set result limit.
|
|
pub fn with_limit(mut self, limit: usize) -> Self {
|
|
self.limit = limit;
|
|
self
|
|
}
|
|
|
|
/// Set offset.
|
|
pub fn with_offset(mut self, offset: usize) -> Self {
|
|
self.offset = offset;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Model summary for search results.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelSummary {
|
|
/// Model ID
|
|
pub model_id: ModelId,
|
|
/// Model title
|
|
pub title: String,
|
|
/// Short description
|
|
pub description: String,
|
|
/// Model author
|
|
pub author: String,
|
|
/// Model tags
|
|
pub tags: Vec<String>,
|
|
/// License
|
|
pub license: Option<String>,
|
|
/// Last modified
|
|
pub last_modified: DateTime<Utc>,
|
|
}
|
|
|
|
impl ModelSummary {
|
|
/// Create a model summary from metadata.
|
|
pub fn from_metadata(metadata: &ModelMetadata) -> Self {
|
|
Self {
|
|
model_id: metadata.id.clone(),
|
|
title: metadata.title.clone(),
|
|
description: metadata.description.clone(),
|
|
author: metadata.author.clone(),
|
|
tags: metadata.tags.clone(),
|
|
license: metadata.license.clone(),
|
|
last_modified: metadata.updated_at,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Search result with ranking information.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SearchResult {
|
|
/// Model identifier
|
|
pub model_id: ModelId,
|
|
/// Model summary
|
|
pub model_summary: ModelSummary,
|
|
/// Total downloads
|
|
pub downloads: u64,
|
|
/// Number of likes
|
|
pub likes: u64,
|
|
/// Trending score (0.0 to 1.0)
|
|
pub trending_score: f64,
|
|
/// Last modified timestamp
|
|
pub last_modified: DateTime<Utc>,
|
|
/// Relevance score (0.0 to 1.0) for text search
|
|
pub relevance_score: Option<f64>,
|
|
}
|
|
|
|
/// Model discovery and search functionality.
|
|
pub struct ModelDiscovery {
|
|
/// Model registry
|
|
registry: Arc<ModelRegistry>,
|
|
/// Model metrics tracking
|
|
metrics: Arc<DashMap<String, ModelMetrics>>,
|
|
}
|
|
|
|
impl ModelDiscovery {
|
|
/// Create a new model discovery instance.
|
|
pub fn new(registry: ModelRegistry) -> Self {
|
|
Self {
|
|
registry: Arc::new(registry),
|
|
metrics: Arc::new(DashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Execute a search query.
|
|
pub async fn search(&self, query: SearchQuery) -> HubResult<Vec<SearchResult>> {
|
|
debug!("Executing search query: {:?}", query);
|
|
|
|
// Build registry query options
|
|
let mut options = QueryOptions::default();
|
|
|
|
if let Some(ref library) = query.filter_by_library {
|
|
options.framework = Some(library.clone());
|
|
}
|
|
|
|
if let Some(ref task) = query.filter_by_task {
|
|
options.tags = vec![task.clone()];
|
|
}
|
|
|
|
// Get models from registry
|
|
let models = if let Some(ref text_query) = query.query {
|
|
self.registry.search_models(text_query, options).await?
|
|
} else {
|
|
self.registry.list_models(options).await?
|
|
};
|
|
|
|
// Convert to search results
|
|
let mut results: Vec<SearchResult> = models
|
|
.iter()
|
|
.filter(|m| self.apply_filters(m, &query))
|
|
.map(|m| self.to_search_result(m, &query))
|
|
.collect();
|
|
|
|
// Sort results
|
|
self.sort_results(&mut results, query.sort_by);
|
|
|
|
// Apply pagination
|
|
let start = query.offset.min(results.len());
|
|
let end = (query.offset + query.limit).min(results.len());
|
|
results = results[start..end].to_vec();
|
|
|
|
debug!("Found {} search results", results.len());
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
/// Get trending models.
|
|
pub async fn get_trending(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
|
let query = SearchQuery::new()
|
|
.with_sort_by(SortBy::Trending)
|
|
.with_limit(limit);
|
|
|
|
self.search(query).await
|
|
}
|
|
|
|
/// Get most downloaded models.
|
|
pub async fn get_popular(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
|
let query = SearchQuery::new()
|
|
.with_sort_by(SortBy::Downloads)
|
|
.with_limit(limit);
|
|
|
|
self.search(query).await
|
|
}
|
|
|
|
/// Get recently updated models.
|
|
pub async fn get_recent(&self, limit: usize) -> HubResult<Vec<SearchResult>> {
|
|
let query = SearchQuery::new()
|
|
.with_sort_by(SortBy::Recent)
|
|
.with_limit(limit);
|
|
|
|
self.search(query).await
|
|
}
|
|
|
|
/// Get models for a specific task.
|
|
pub async fn get_by_task(&self, task: &str, limit: usize) -> HubResult<Vec<SearchResult>> {
|
|
let query = SearchQuery::new().with_task(task).with_limit(limit);
|
|
|
|
self.search(query).await
|
|
}
|
|
|
|
/// Get similar models based on tags and task.
|
|
pub async fn get_similar(
|
|
&self,
|
|
model_id: &ModelId,
|
|
limit: usize,
|
|
) -> HubResult<Vec<SearchResult>> {
|
|
// Get the reference model
|
|
let reference = self.registry.get_model(model_id, None).await?;
|
|
|
|
// Build query based on model tags
|
|
let mut query = SearchQuery::new().with_limit(limit);
|
|
|
|
if !reference.metadata.tags.is_empty() {
|
|
// Use first tag as task filter
|
|
query = query.with_task(reference.metadata.tags[0].clone());
|
|
}
|
|
|
|
// Search and filter out the reference model
|
|
let mut results = self.search(query).await?;
|
|
results.retain(|r| r.model_id != *model_id);
|
|
|
|
// Calculate similarity scores
|
|
for result in &mut results {
|
|
let similarity = self.calculate_similarity(&reference.metadata, &result.model_summary);
|
|
result.relevance_score = Some(similarity);
|
|
}
|
|
|
|
// Sort by similarity
|
|
results.sort_by(|a, b| {
|
|
b.relevance_score
|
|
.unwrap_or(0.0)
|
|
.partial_cmp(&a.relevance_score.unwrap_or(0.0))
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
/// Record a download for metrics.
|
|
pub async fn record_download(&self, model_id: &ModelId) -> HubResult<()> {
|
|
let key = model_id.to_string();
|
|
|
|
self.metrics
|
|
.entry(key)
|
|
.and_modify(|m| {
|
|
m.downloads += 1;
|
|
m.last_downloaded = Some(Utc::now());
|
|
})
|
|
.or_insert_with(|| ModelMetrics {
|
|
model_id: model_id.clone(),
|
|
downloads: 1,
|
|
likes: 0,
|
|
last_downloaded: Some(Utc::now()),
|
|
last_liked: None,
|
|
daily_downloads: HashMap::new(),
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Record a like for metrics.
|
|
pub async fn record_like(&self, model_id: &ModelId) -> HubResult<()> {
|
|
let key = model_id.to_string();
|
|
|
|
self.metrics
|
|
.entry(key)
|
|
.and_modify(|m| {
|
|
m.likes += 1;
|
|
m.last_liked = Some(Utc::now());
|
|
})
|
|
.or_insert_with(|| ModelMetrics {
|
|
model_id: model_id.clone(),
|
|
downloads: 0,
|
|
likes: 1,
|
|
last_downloaded: None,
|
|
last_liked: Some(Utc::now()),
|
|
daily_downloads: HashMap::new(),
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Apply additional filters to model info.
|
|
fn apply_filters(&self, model: &ModelInfo, query: &SearchQuery) -> bool {
|
|
// Filter by author
|
|
if let Some(ref author) = query.filter_by_author {
|
|
if !model
|
|
.metadata
|
|
.author
|
|
.to_lowercase()
|
|
.contains(&author.to_lowercase())
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Filter by license
|
|
if let Some(ref licenses) = query.filter_by_license {
|
|
if let Some(ref model_license) = model.metadata.license {
|
|
if !licenses
|
|
.iter()
|
|
.any(|l| model_license.to_lowercase().contains(&l.to_lowercase()))
|
|
{
|
|
return false;
|
|
}
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Filter by language (using tags as proxy)
|
|
if let Some(ref languages) = query.filter_by_language {
|
|
let has_language = languages.iter().any(|lang| {
|
|
model
|
|
.metadata
|
|
.tags
|
|
.iter()
|
|
.any(|tag| tag.to_lowercase().contains(&lang.to_lowercase()))
|
|
});
|
|
|
|
if !has_language {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
/// Convert ModelInfo to SearchResult.
|
|
fn to_search_result(&self, model: &ModelInfo, query: &SearchQuery) -> SearchResult {
|
|
let metrics = self.get_metrics(&model.metadata.id);
|
|
|
|
let trending_score = self.calculate_trending_score(&metrics, &model.metadata);
|
|
|
|
let relevance_score =
|
|
if query.query.is_some() {
|
|
Some(self.calculate_text_relevance(
|
|
&model.metadata,
|
|
query.query.as_deref().unwrap_or(""),
|
|
))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
SearchResult {
|
|
model_id: model.metadata.id.clone(),
|
|
model_summary: ModelSummary::from_metadata(&model.metadata),
|
|
downloads: metrics.downloads,
|
|
likes: metrics.likes,
|
|
trending_score,
|
|
last_modified: model.metadata.updated_at,
|
|
relevance_score,
|
|
}
|
|
}
|
|
|
|
/// Get metrics for a model.
|
|
fn get_metrics(&self, model_id: &ModelId) -> ModelMetrics {
|
|
self.metrics
|
|
.get(&model_id.to_string())
|
|
.map(|m| m.clone())
|
|
.unwrap_or_else(|| ModelMetrics {
|
|
model_id: model_id.clone(),
|
|
downloads: 0,
|
|
likes: 0,
|
|
last_downloaded: None,
|
|
last_liked: None,
|
|
daily_downloads: HashMap::new(),
|
|
})
|
|
}
|
|
|
|
/// Calculate trending score based on recent activity.
|
|
fn calculate_trending_score(&self, metrics: &ModelMetrics, metadata: &ModelMetadata) -> f64 {
|
|
let now = Utc::now();
|
|
let age_days = (now - metadata.created_at).num_days() as f64;
|
|
|
|
// Avoid division by zero
|
|
let age_factor = if age_days > 0.0 {
|
|
1.0 / (1.0 + age_days.ln())
|
|
} else {
|
|
1.0
|
|
};
|
|
|
|
// Recent activity factor
|
|
let recent_factor = if let Some(last_downloaded) = metrics.last_downloaded {
|
|
let days_since = (now - last_downloaded).num_days() as f64;
|
|
(-days_since / 7.0).exp() // Exponential decay over 7 days
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Combine downloads, likes, and recency
|
|
let popularity = (metrics.downloads as f64).ln_1p() + (metrics.likes as f64 * 2.0).ln_1p();
|
|
|
|
let score = (popularity * age_factor * 0.4) + (recent_factor * 0.6);
|
|
|
|
// Normalize to 0.0-1.0 range
|
|
score.min(1.0).max(0.0)
|
|
}
|
|
|
|
/// Calculate text relevance score.
|
|
fn calculate_text_relevance(&self, metadata: &ModelMetadata, query: &str) -> f64 {
|
|
let query_lower = query.to_lowercase();
|
|
let mut score = 0.0;
|
|
|
|
// Title match (highest weight)
|
|
if metadata.title.to_lowercase().contains(&query_lower) {
|
|
score += 0.5;
|
|
}
|
|
|
|
// Description match
|
|
if metadata.description.to_lowercase().contains(&query_lower) {
|
|
score += 0.3;
|
|
}
|
|
|
|
// Tag matches
|
|
let tag_matches = metadata
|
|
.tags
|
|
.iter()
|
|
.filter(|tag| tag.to_lowercase().contains(&query_lower))
|
|
.count();
|
|
score += (tag_matches as f64 * 0.1).min(0.2);
|
|
|
|
score.min(1.0)
|
|
}
|
|
|
|
/// Calculate similarity between models.
|
|
fn calculate_similarity(&self, reference: &ModelMetadata, candidate: &ModelSummary) -> f64 {
|
|
let mut score = 0.0;
|
|
|
|
// Tag overlap
|
|
let common_tags = reference
|
|
.tags
|
|
.iter()
|
|
.filter(|tag| candidate.tags.contains(tag))
|
|
.count();
|
|
|
|
if !reference.tags.is_empty() {
|
|
score += (common_tags as f64 / reference.tags.len() as f64) * 0.6;
|
|
}
|
|
|
|
// Same author
|
|
if reference.author == candidate.author {
|
|
score += 0.2;
|
|
}
|
|
|
|
// Same license
|
|
if reference.license == candidate.license {
|
|
score += 0.2;
|
|
}
|
|
|
|
score.min(1.0)
|
|
}
|
|
|
|
/// Sort search results.
|
|
fn sort_results(&self, results: &mut [SearchResult], sort_by: SortBy) {
|
|
match sort_by {
|
|
SortBy::Downloads => {
|
|
results.sort_by(|a, b| b.downloads.cmp(&a.downloads));
|
|
}
|
|
SortBy::Likes => {
|
|
results.sort_by(|a, b| b.likes.cmp(&a.likes));
|
|
}
|
|
SortBy::Recent => {
|
|
results.sort_by(|a, b| b.last_modified.cmp(&a.last_modified));
|
|
}
|
|
SortBy::Trending => {
|
|
results.sort_by(|a, b| {
|
|
b.trending_score
|
|
.partial_cmp(&a.trending_score)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
}
|
|
SortBy::Relevance => {
|
|
results.sort_by(|a, b| {
|
|
let score_a = a.relevance_score.unwrap_or(0.0);
|
|
let score_b = b.relevance_score.unwrap_or(0.0);
|
|
score_b
|
|
.partial_cmp(&score_a)
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model metrics for tracking popularity and trends.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelMetrics {
|
|
/// Model ID
|
|
pub model_id: ModelId,
|
|
/// Total downloads
|
|
pub downloads: u64,
|
|
/// Total likes
|
|
pub likes: u64,
|
|
/// Last download timestamp
|
|
pub last_downloaded: Option<DateTime<Utc>>,
|
|
/// Last like timestamp
|
|
pub last_liked: Option<DateTime<Utc>>,
|
|
/// Daily download counts (date -> count)
|
|
pub daily_downloads: HashMap<String, u64>,
|
|
}
|
|
|
|
impl ModelMetrics {
|
|
/// Create new metrics for a model.
|
|
pub fn new(model_id: ModelId) -> Self {
|
|
Self {
|
|
model_id,
|
|
downloads: 0,
|
|
likes: 0,
|
|
last_downloaded: None,
|
|
last_liked: None,
|
|
daily_downloads: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Increment download count.
|
|
pub fn record_download(&mut self) {
|
|
self.downloads += 1;
|
|
self.last_downloaded = Some(Utc::now());
|
|
|
|
// Record daily download
|
|
let date = Utc::now().format("%Y-%m-%d").to_string();
|
|
*self.daily_downloads.entry(date).or_insert(0) += 1;
|
|
}
|
|
|
|
/// Increment like count.
|
|
pub fn record_like(&mut self) {
|
|
self.likes += 1;
|
|
self.last_liked = Some(Utc::now());
|
|
}
|
|
|
|
/// Get downloads for a specific date.
|
|
pub fn get_downloads_for_date(&self, date: &str) -> u64 {
|
|
self.daily_downloads.get(date).copied().unwrap_or(0)
|
|
}
|
|
|
|
/// Get total downloads in the last N days.
|
|
pub fn get_recent_downloads(&self, days: i64) -> u64 {
|
|
let cutoff = Utc::now() - chrono::Duration::days(days);
|
|
|
|
self.daily_downloads
|
|
.iter()
|
|
.filter(|(date_str, _)| {
|
|
if let Ok(date) = chrono::NaiveDate::parse_from_str(date_str, "%Y-%m-%d") {
|
|
let datetime = date.and_hms_opt(0, 0, 0).unwrap();
|
|
datetime.and_utc() >= cutoff
|
|
} else {
|
|
false
|
|
}
|
|
})
|
|
.map(|(_, count)| count)
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::model::ModelSchema;
|
|
use crate::{ModelStatus, ModelVersion, RegistryConfig, StorageConfig};
|
|
use semver::Version;
|
|
use tempfile::TempDir;
|
|
|
|
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, title: &str, tags: Vec<String>) -> ModelMetadata {
|
|
ModelMetadata {
|
|
id,
|
|
version: ModelVersion::new(Version::parse("1.0.0").unwrap()),
|
|
title: title.to_string(),
|
|
description: format!("Description for {}", title),
|
|
architecture: "transformer".to_string(),
|
|
framework: "rustytorch".to_string(),
|
|
framework_version: "1.0.0".to_string(),
|
|
tags,
|
|
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(),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_search_query_builder() {
|
|
let query = SearchQuery::new()
|
|
.with_query("BERT")
|
|
.with_task("text-classification")
|
|
.with_library("rustytorch")
|
|
.with_limit(10)
|
|
.with_offset(5);
|
|
|
|
assert_eq!(query.query, Some("BERT".to_string()));
|
|
assert_eq!(
|
|
query.filter_by_task,
|
|
Some("text-classification".to_string())
|
|
);
|
|
assert_eq!(query.filter_by_library, Some("rustytorch".to_string()));
|
|
assert_eq!(query.limit, 10);
|
|
assert_eq!(query.offset, 5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_summary_from_metadata() {
|
|
let model_id = ModelId::new("test", "model");
|
|
let metadata =
|
|
create_test_metadata(model_id.clone(), "Test Model", vec!["test".to_string()]);
|
|
|
|
let summary = ModelSummary::from_metadata(&metadata);
|
|
|
|
assert_eq!(summary.model_id, model_id);
|
|
assert_eq!(summary.title, "Test Model");
|
|
assert_eq!(summary.tags, vec!["test".to_string()]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_search() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
// Register test models
|
|
let model1 = create_test_metadata(
|
|
ModelId::new("test", "bert-base"),
|
|
"BERT Base Model",
|
|
vec!["nlp".to_string(), "text-classification".to_string()],
|
|
);
|
|
|
|
let model2 = create_test_metadata(
|
|
ModelId::new("test", "gpt2"),
|
|
"GPT-2 Model",
|
|
vec!["nlp".to_string(), "text-generation".to_string()],
|
|
);
|
|
|
|
registry.register_model(model1).await.unwrap();
|
|
registry.register_model(model2).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Search for BERT
|
|
let query = SearchQuery::new().with_query("BERT");
|
|
let results = discovery.search(query).await.unwrap();
|
|
|
|
assert_eq!(results.len(), 1);
|
|
assert!(results[0].model_summary.title.contains("BERT"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_filter_by_task() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
let model1 = create_test_metadata(
|
|
ModelId::new("test", "classifier"),
|
|
"Classifier Model",
|
|
vec!["text-classification".to_string()],
|
|
);
|
|
|
|
let model2 = create_test_metadata(
|
|
ModelId::new("test", "generator"),
|
|
"Generator Model",
|
|
vec!["text-generation".to_string()],
|
|
);
|
|
|
|
registry.register_model(model1).await.unwrap();
|
|
registry.register_model(model2).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Search by task
|
|
let results = discovery
|
|
.get_by_task("text-classification", 10)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(results.len(), 1);
|
|
assert!(
|
|
results[0]
|
|
.model_summary
|
|
.tags
|
|
.contains(&"text-classification".to_string())
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_trending() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
let model = create_test_metadata(
|
|
ModelId::new("test", "model"),
|
|
"Test Model",
|
|
vec!["test".to_string()],
|
|
);
|
|
|
|
registry.register_model(model).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Record some activity
|
|
let model_id = ModelId::new("test", "model");
|
|
discovery.record_download(&model_id).await.unwrap();
|
|
discovery.record_like(&model_id).await.unwrap();
|
|
|
|
// Get trending
|
|
let results = discovery.get_trending(10).await.unwrap();
|
|
|
|
assert!(!results.is_empty());
|
|
assert!(results[0].trending_score > 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_popular() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
let model = create_test_metadata(
|
|
ModelId::new("test", "model"),
|
|
"Test Model",
|
|
vec!["test".to_string()],
|
|
);
|
|
|
|
registry.register_model(model).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Record downloads
|
|
let model_id = ModelId::new("test", "model");
|
|
for _ in 0..5 {
|
|
discovery.record_download(&model_id).await.unwrap();
|
|
}
|
|
|
|
// Get popular
|
|
let results = discovery.get_popular(10).await.unwrap();
|
|
|
|
assert!(!results.is_empty());
|
|
assert_eq!(results[0].downloads, 5);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_recent() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
let model = create_test_metadata(
|
|
ModelId::new("test", "model"),
|
|
"Test Model",
|
|
vec!["test".to_string()],
|
|
);
|
|
|
|
registry.register_model(model).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Get recent
|
|
let results = discovery.get_recent(10).await.unwrap();
|
|
|
|
assert!(!results.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_similar() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
|
|
let model1 = create_test_metadata(
|
|
ModelId::new("test", "bert-base"),
|
|
"BERT Base",
|
|
vec!["nlp".to_string(), "encoder".to_string()],
|
|
);
|
|
|
|
let model2 = create_test_metadata(
|
|
ModelId::new("test", "roberta"),
|
|
"RoBERTa",
|
|
vec!["nlp".to_string(), "encoder".to_string()],
|
|
);
|
|
|
|
let model3 = create_test_metadata(
|
|
ModelId::new("test", "gpt2"),
|
|
"GPT-2",
|
|
vec!["nlp".to_string(), "decoder".to_string()],
|
|
);
|
|
|
|
registry.register_model(model1).await.unwrap();
|
|
registry.register_model(model2).await.unwrap();
|
|
registry.register_model(model3).await.unwrap();
|
|
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
// Find similar to BERT
|
|
let similar = discovery
|
|
.get_similar(&ModelId::new("test", "bert-base"), 10)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(!similar.is_empty());
|
|
// RoBERTa should be more similar than GPT-2
|
|
assert_eq!(similar[0].model_id.name, "roberta");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_model_discovery_record_metrics() {
|
|
let (registry, _temp_dir) = create_test_registry().await;
|
|
let discovery = ModelDiscovery::new(registry);
|
|
|
|
let model_id = ModelId::new("test", "model");
|
|
|
|
// Record downloads and likes
|
|
discovery.record_download(&model_id).await.unwrap();
|
|
discovery.record_download(&model_id).await.unwrap();
|
|
discovery.record_like(&model_id).await.unwrap();
|
|
|
|
// Check metrics
|
|
let metrics = discovery.get_metrics(&model_id);
|
|
assert_eq!(metrics.downloads, 2);
|
|
assert_eq!(metrics.likes, 1);
|
|
assert!(metrics.last_downloaded.is_some());
|
|
assert!(metrics.last_liked.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_metrics_record_download() {
|
|
let model_id = ModelId::new("test", "model");
|
|
let mut metrics = ModelMetrics::new(model_id);
|
|
|
|
metrics.record_download();
|
|
|
|
assert_eq!(metrics.downloads, 1);
|
|
assert!(metrics.last_downloaded.is_some());
|
|
assert!(!metrics.daily_downloads.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_metrics_record_like() {
|
|
let model_id = ModelId::new("test", "model");
|
|
let mut metrics = ModelMetrics::new(model_id);
|
|
|
|
metrics.record_like();
|
|
|
|
assert_eq!(metrics.likes, 1);
|
|
assert!(metrics.last_liked.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_metrics_daily_downloads() {
|
|
let model_id = ModelId::new("test", "model");
|
|
let mut metrics = ModelMetrics::new(model_id);
|
|
|
|
// Record multiple downloads
|
|
for _ in 0..3 {
|
|
metrics.record_download();
|
|
}
|
|
|
|
let today = Utc::now().format("%Y-%m-%d").to_string();
|
|
assert_eq!(metrics.get_downloads_for_date(&today), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_metrics_recent_downloads() {
|
|
let model_id = ModelId::new("test", "model");
|
|
let mut metrics = ModelMetrics::new(model_id);
|
|
|
|
// Record downloads
|
|
metrics.record_download();
|
|
metrics.record_download();
|
|
|
|
let recent = metrics.get_recent_downloads(7);
|
|
assert_eq!(recent, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_search_query_pagination() {
|
|
let query = SearchQuery::new().with_limit(20).with_offset(10);
|
|
|
|
assert_eq!(query.limit, 20);
|
|
assert_eq!(query.offset, 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sort_by_variants() {
|
|
assert_eq!(SortBy::Downloads, SortBy::Downloads);
|
|
assert_ne!(SortBy::Downloads, SortBy::Likes);
|
|
}
|
|
}
|