use crate::error::{Result, TrustformersError};
use std::path::Path;
use super::manager::OfflineModelPackManager;
use super::pack_types::PackCreationConfig;
use super::resolution::{empty_model_info, ModelInfo};
pub struct HubIntegration {
pub hub_options: crate::hub::HubOptions,
}
impl HubIntegration {
pub fn new(options: Option<crate::hub::HubOptions>) -> Self {
Self {
hub_options: options.unwrap_or_default(),
}
}
pub async fn download_model_to_pack(
&self,
pack_manager: &mut OfflineModelPackManager,
model_id: &str,
pack_id: &str,
) -> Result<()> {
let _model_path = crate::hub::download_file_from_hub(
model_id,
"config.json",
Some(self.hub_options.clone()),
)
.map_err(|e| TrustformersError::io_error(format!("Hub download failed: {}", e)))?;
let _model_info = self.get_hub_model_info(model_id).await?;
let additional_models = vec![model_id.to_string()];
pack_manager.update_pack(pack_id, additional_models).await?;
Ok(())
}
pub async fn create_pack_from_hub_collection(
&self,
pack_manager: &mut OfflineModelPackManager,
collection_name: &str,
model_ids: Vec<String>,
config: PackCreationConfig,
) -> Result<String> {
for model_id in &model_ids {
let _ = self.get_hub_model_info(model_id).await?;
}
pack_manager
.create_pack(
format!("Hub Collection: {}", collection_name),
format!(
"Model pack created from Hub collection: {}",
collection_name
),
model_ids,
config,
)
.await
}
pub(super) async fn get_hub_model_info(&self, model_id: &str) -> Result<ModelInfo> {
if Path::new(model_id).is_dir() {
return Ok(empty_model_info(model_id));
}
match crate::hub::load_model_card_from_hub(model_id, Some(self.hub_options.clone())) {
Ok(model_card) => {
Ok(ModelInfo {
model_id: model_id.to_string(),
library_name: Some("transformers".to_string()),
pipeline_tag: model_card.pipeline_tag.clone(),
tags: model_card.tags.unwrap_or_default(),
config: model_card.extra.into_iter().collect(),
downloads: None, likes: None, created_at: None,
updated_at: None,
author: None,
description: None,
license: model_card.license,
task: model_card.pipeline_tag,
language: model_card.language.unwrap_or_default(),
dataset: model_card.datasets.unwrap_or_default(),
model_type: None,
architecture: None,
})
},
Err(_) => {
Ok(empty_model_info(model_id))
},
}
}
}