osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
937 lines
28 KiB
Rust
937 lines
28 KiB
Rust
//! HuggingFace Hub compatibility layer.
|
|
//!
|
|
//! This module provides seamless integration with HuggingFace Hub, enabling:
|
|
//! - Loading models from HuggingFace Hub
|
|
//! - Converting HuggingFace configs to RustyTorch format
|
|
//! - Auto-downloading and caching models
|
|
//! - Support for private models with authentication
|
|
|
|
use crate::{HubError, HubResult, LocalCache, PretrainedConfig};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::path::{Path, PathBuf};
|
|
use tokio::fs;
|
|
use tracing::{debug, info, warn};
|
|
|
|
/// HuggingFace Hub base URL.
|
|
const HF_HUB_URL: &str = "https://huggingface.co";
|
|
|
|
/// HuggingFace model configuration (config.json).
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct HFModelConfig {
|
|
/// Model architecture (e.g., "LlamaForCausalLM", "GPT2LMHeadModel")
|
|
#[serde(default)]
|
|
pub architectures: Vec<String>,
|
|
/// Model type (e.g., "llama", "gpt2", "bert")
|
|
#[serde(default)]
|
|
pub model_type: Option<String>,
|
|
/// Hidden size
|
|
#[serde(default)]
|
|
pub hidden_size: Option<usize>,
|
|
/// Number of hidden layers
|
|
#[serde(default)]
|
|
pub num_hidden_layers: Option<usize>,
|
|
/// Number of attention heads
|
|
#[serde(default)]
|
|
pub num_attention_heads: Option<usize>,
|
|
/// Number of key-value heads (for GQA)
|
|
#[serde(default)]
|
|
pub num_key_value_heads: Option<usize>,
|
|
/// Intermediate size (FFN)
|
|
#[serde(default)]
|
|
pub intermediate_size: Option<usize>,
|
|
/// Vocabulary size
|
|
#[serde(default)]
|
|
pub vocab_size: Option<usize>,
|
|
/// Maximum position embeddings
|
|
#[serde(default)]
|
|
pub max_position_embeddings: Option<usize>,
|
|
/// RMS norm epsilon
|
|
#[serde(default)]
|
|
pub rms_norm_eps: Option<f64>,
|
|
/// Layer norm epsilon
|
|
#[serde(default)]
|
|
pub layer_norm_eps: Option<f64>,
|
|
/// Rope theta
|
|
#[serde(default)]
|
|
pub rope_theta: Option<f64>,
|
|
/// Rope scaling configuration
|
|
#[serde(default)]
|
|
pub rope_scaling: Option<RopeScalingConfig>,
|
|
/// Tie word embeddings
|
|
#[serde(default)]
|
|
pub tie_word_embeddings: Option<bool>,
|
|
/// Hidden activation function
|
|
#[serde(default)]
|
|
pub hidden_act: Option<String>,
|
|
/// Torch dtype (e.g., "float16", "bfloat16")
|
|
#[serde(default)]
|
|
pub torch_dtype: Option<String>,
|
|
/// Use cache for generation
|
|
#[serde(default)]
|
|
pub use_cache: Option<bool>,
|
|
/// Beginning of sequence token ID
|
|
#[serde(default)]
|
|
pub bos_token_id: Option<u32>,
|
|
/// End of sequence token ID
|
|
#[serde(default)]
|
|
pub eos_token_id: Option<EosTokenId>,
|
|
/// Padding token ID
|
|
#[serde(default)]
|
|
pub pad_token_id: Option<u32>,
|
|
/// Additional configuration fields
|
|
#[serde(flatten)]
|
|
pub extra: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// EOS token ID can be a single value or list.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(untagged)]
|
|
pub enum EosTokenId {
|
|
/// Single EOS token
|
|
Single(u32),
|
|
/// Multiple EOS tokens
|
|
Multiple(Vec<u32>),
|
|
}
|
|
|
|
impl EosTokenId {
|
|
/// Get the primary EOS token ID.
|
|
pub fn primary(&self) -> u32 {
|
|
match self {
|
|
EosTokenId::Single(id) => *id,
|
|
EosTokenId::Multiple(ids) => ids.first().copied().unwrap_or(0),
|
|
}
|
|
}
|
|
|
|
/// Get all EOS token IDs.
|
|
pub fn all(&self) -> Vec<u32> {
|
|
match self {
|
|
EosTokenId::Single(id) => vec![*id],
|
|
EosTokenId::Multiple(ids) => ids.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Rope scaling configuration.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RopeScalingConfig {
|
|
/// Scaling type (e.g., "linear", "dynamic")
|
|
#[serde(rename = "type")]
|
|
pub scaling_type: String,
|
|
/// Scaling factor
|
|
pub factor: f64,
|
|
/// Original max position embeddings (for dynamic scaling)
|
|
#[serde(default)]
|
|
pub original_max_position_embeddings: Option<usize>,
|
|
}
|
|
|
|
/// HuggingFace tokenizer configuration.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct HFTokenizerConfig {
|
|
/// Tokenizer class name
|
|
#[serde(default)]
|
|
pub tokenizer_class: Option<String>,
|
|
/// Model max length
|
|
#[serde(default)]
|
|
pub model_max_length: Option<usize>,
|
|
/// Padding side
|
|
#[serde(default)]
|
|
pub padding_side: Option<String>,
|
|
/// Truncation side
|
|
#[serde(default)]
|
|
pub truncation_side: Option<String>,
|
|
/// Chat template
|
|
#[serde(default)]
|
|
pub chat_template: Option<String>,
|
|
/// Beginning of sequence token
|
|
#[serde(default)]
|
|
pub bos_token: Option<TokenValue>,
|
|
/// End of sequence token
|
|
#[serde(default)]
|
|
pub eos_token: Option<TokenValue>,
|
|
/// Padding token
|
|
#[serde(default)]
|
|
pub pad_token: Option<TokenValue>,
|
|
/// Unknown token
|
|
#[serde(default)]
|
|
pub unk_token: Option<TokenValue>,
|
|
/// Additional configuration fields
|
|
#[serde(flatten)]
|
|
pub extra: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// Token value (can be string or object with content).
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(untagged)]
|
|
pub enum TokenValue {
|
|
/// Simple string token
|
|
Simple(String),
|
|
/// Token with metadata
|
|
WithMetadata {
|
|
content: String,
|
|
#[serde(default)]
|
|
lstrip: bool,
|
|
#[serde(default)]
|
|
rstrip: bool,
|
|
#[serde(default)]
|
|
single_word: bool,
|
|
#[serde(default)]
|
|
normalized: bool,
|
|
},
|
|
}
|
|
|
|
impl TokenValue {
|
|
/// Get the token content.
|
|
pub fn content(&self) -> &str {
|
|
match self {
|
|
TokenValue::Simple(s) => s,
|
|
TokenValue::WithMetadata { content, .. } => content,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// HuggingFace generation configuration.
|
|
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
|
pub struct HFGenerationConfig {
|
|
/// Maximum new tokens to generate
|
|
#[serde(default)]
|
|
pub max_new_tokens: Option<usize>,
|
|
/// Maximum total length
|
|
#[serde(default)]
|
|
pub max_length: Option<usize>,
|
|
/// Temperature for sampling
|
|
#[serde(default)]
|
|
pub temperature: Option<f64>,
|
|
/// Top-p (nucleus) sampling
|
|
#[serde(default)]
|
|
pub top_p: Option<f64>,
|
|
/// Top-k sampling
|
|
#[serde(default)]
|
|
pub top_k: Option<usize>,
|
|
/// Repetition penalty
|
|
#[serde(default)]
|
|
pub repetition_penalty: Option<f64>,
|
|
/// Do sample
|
|
#[serde(default)]
|
|
pub do_sample: Option<bool>,
|
|
/// Number of beams
|
|
#[serde(default)]
|
|
pub num_beams: Option<usize>,
|
|
/// Early stopping
|
|
#[serde(default)]
|
|
pub early_stopping: Option<bool>,
|
|
/// EOS token ID
|
|
#[serde(default)]
|
|
pub eos_token_id: Option<EosTokenId>,
|
|
/// PAD token ID
|
|
#[serde(default)]
|
|
pub pad_token_id: Option<u32>,
|
|
/// BOS token ID
|
|
#[serde(default)]
|
|
pub bos_token_id: Option<u32>,
|
|
/// Additional configuration fields
|
|
#[serde(flatten)]
|
|
pub extra: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// RustyTorch model configuration.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RTXModelConfig {
|
|
/// Model architecture type
|
|
pub architecture: RTXArchitecture,
|
|
/// Hidden dimension size
|
|
pub hidden_size: usize,
|
|
/// Number of transformer layers
|
|
pub num_layers: usize,
|
|
/// Number of attention heads
|
|
pub num_heads: usize,
|
|
/// Number of key-value heads (for GQA, defaults to num_heads)
|
|
pub num_kv_heads: usize,
|
|
/// FFN intermediate size
|
|
pub intermediate_size: usize,
|
|
/// Vocabulary size
|
|
pub vocab_size: usize,
|
|
/// Maximum sequence length
|
|
pub max_seq_len: usize,
|
|
/// RoPE theta parameter
|
|
pub rope_theta: f64,
|
|
/// Layer norm epsilon
|
|
pub norm_eps: f64,
|
|
/// Activation function
|
|
pub activation: RTXActivation,
|
|
/// Data type for inference
|
|
pub dtype: RTXDType,
|
|
/// Whether to tie input/output embeddings
|
|
pub tie_embeddings: bool,
|
|
/// RoPE scaling configuration
|
|
pub rope_scaling: Option<RTXRopeScaling>,
|
|
/// Beginning of sequence token ID
|
|
pub bos_token_id: u32,
|
|
/// End of sequence token IDs
|
|
pub eos_token_ids: Vec<u32>,
|
|
/// Padding token ID
|
|
pub pad_token_id: Option<u32>,
|
|
}
|
|
|
|
/// RustyTorch model architecture types.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum RTXArchitecture {
|
|
/// LLaMA family (LLaMA, LLaMA 2, LLaMA 3)
|
|
Llama,
|
|
/// Mistral family
|
|
Mistral,
|
|
/// Mixtral (MoE)
|
|
Mixtral,
|
|
/// GPT-2
|
|
Gpt2,
|
|
/// GPT-NeoX / Pythia
|
|
GptNeoX,
|
|
/// Falcon
|
|
Falcon,
|
|
/// Phi family
|
|
Phi,
|
|
/// Qwen family
|
|
Qwen,
|
|
/// BERT
|
|
Bert,
|
|
/// RoBERTa
|
|
Roberta,
|
|
/// T5
|
|
T5,
|
|
/// Mamba (state space model)
|
|
Mamba,
|
|
/// Gemma
|
|
Gemma,
|
|
/// Other/unknown
|
|
Other,
|
|
}
|
|
|
|
impl RTXArchitecture {
|
|
/// Detect architecture from HuggingFace model type.
|
|
pub fn from_hf_model_type(model_type: &str) -> Self {
|
|
match model_type.to_lowercase().as_str() {
|
|
"llama" => RTXArchitecture::Llama,
|
|
"mistral" => RTXArchitecture::Mistral,
|
|
"mixtral" => RTXArchitecture::Mixtral,
|
|
"gpt2" => RTXArchitecture::Gpt2,
|
|
"gpt_neox" | "gptneox" => RTXArchitecture::GptNeoX,
|
|
"falcon" => RTXArchitecture::Falcon,
|
|
"phi" | "phi3" => RTXArchitecture::Phi,
|
|
"qwen" | "qwen2" => RTXArchitecture::Qwen,
|
|
"bert" => RTXArchitecture::Bert,
|
|
"roberta" => RTXArchitecture::Roberta,
|
|
"t5" => RTXArchitecture::T5,
|
|
"mamba" => RTXArchitecture::Mamba,
|
|
"gemma" | "gemma2" => RTXArchitecture::Gemma,
|
|
_ => RTXArchitecture::Other,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Activation function types.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum RTXActivation {
|
|
/// SiLU / Swish
|
|
SiLU,
|
|
/// GELU
|
|
GELU,
|
|
/// GELUApprox (tanh approximation)
|
|
GELUApprox,
|
|
/// ReLU
|
|
ReLU,
|
|
/// Mish
|
|
Mish,
|
|
}
|
|
|
|
impl RTXActivation {
|
|
/// Convert from HuggingFace activation name.
|
|
pub fn from_hf(name: &str) -> Self {
|
|
match name.to_lowercase().as_str() {
|
|
"silu" | "swish" | "silu_and_mul" => RTXActivation::SiLU,
|
|
"gelu" | "gelu_new" => RTXActivation::GELU,
|
|
"gelu_pytorch_tanh" | "gelu_fast" => RTXActivation::GELUApprox,
|
|
"relu" => RTXActivation::ReLU,
|
|
"mish" => RTXActivation::Mish,
|
|
_ => RTXActivation::SiLU, // Default
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Data type for model weights.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum RTXDType {
|
|
/// 32-bit float
|
|
F32,
|
|
/// 16-bit float
|
|
F16,
|
|
/// Brain float 16
|
|
BF16,
|
|
/// 8-bit float (E4M3)
|
|
F8E4M3,
|
|
/// 8-bit float (E5M2)
|
|
F8E5M2,
|
|
}
|
|
|
|
impl RTXDType {
|
|
/// Convert from HuggingFace dtype string.
|
|
pub fn from_hf(dtype: &str) -> Self {
|
|
match dtype.to_lowercase().as_str() {
|
|
"float32" | "f32" => RTXDType::F32,
|
|
"float16" | "f16" => RTXDType::F16,
|
|
"bfloat16" | "bf16" => RTXDType::BF16,
|
|
"float8_e4m3fn" | "fp8_e4m3" => RTXDType::F8E4M3,
|
|
"float8_e5m2" | "fp8_e5m2" => RTXDType::F8E5M2,
|
|
_ => RTXDType::F16, // Default
|
|
}
|
|
}
|
|
}
|
|
|
|
/// RoPE scaling configuration.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RTXRopeScaling {
|
|
/// Scaling type
|
|
pub scaling_type: RopeScalingType,
|
|
/// Scaling factor
|
|
pub factor: f64,
|
|
/// Original max position (for dynamic)
|
|
pub original_max_position: Option<usize>,
|
|
}
|
|
|
|
/// RoPE scaling types.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum RopeScalingType {
|
|
/// Linear interpolation
|
|
Linear,
|
|
/// Dynamic NTK-aware scaling
|
|
Dynamic,
|
|
/// YaRN scaling
|
|
Yarn,
|
|
/// LongRoPE
|
|
LongRope,
|
|
}
|
|
|
|
impl RopeScalingType {
|
|
/// Convert from HuggingFace scaling type.
|
|
pub fn from_hf(s: &str) -> Self {
|
|
match s.to_lowercase().as_str() {
|
|
"linear" => RopeScalingType::Linear,
|
|
"dynamic" => RopeScalingType::Dynamic,
|
|
"yarn" => RopeScalingType::Yarn,
|
|
"longrope" => RopeScalingType::LongRope,
|
|
_ => RopeScalingType::Linear,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Convert HuggingFace config to RustyTorch config.
|
|
pub fn convert_hf_config(hf_config: &HFModelConfig) -> HubResult<RTXModelConfig> {
|
|
let model_type = hf_config
|
|
.model_type
|
|
.as_deref()
|
|
.or_else(|| {
|
|
hf_config
|
|
.architectures
|
|
.first()
|
|
.map(std::string::String::as_str)
|
|
})
|
|
.unwrap_or("unknown");
|
|
|
|
let architecture = RTXArchitecture::from_hf_model_type(model_type);
|
|
|
|
let hidden_size = hf_config
|
|
.hidden_size
|
|
.ok_or_else(|| HubError::InvalidPackage {
|
|
reason: "Missing hidden_size in config".to_string(),
|
|
})?;
|
|
|
|
let num_layers = hf_config
|
|
.num_hidden_layers
|
|
.ok_or_else(|| HubError::InvalidPackage {
|
|
reason: "Missing num_hidden_layers in config".to_string(),
|
|
})?;
|
|
|
|
let num_heads = hf_config
|
|
.num_attention_heads
|
|
.ok_or_else(|| HubError::InvalidPackage {
|
|
reason: "Missing num_attention_heads in config".to_string(),
|
|
})?;
|
|
|
|
let num_kv_heads = hf_config.num_key_value_heads.unwrap_or(num_heads);
|
|
|
|
let intermediate_size =
|
|
hf_config
|
|
.intermediate_size
|
|
.ok_or_else(|| HubError::InvalidPackage {
|
|
reason: "Missing intermediate_size in config".to_string(),
|
|
})?;
|
|
|
|
let vocab_size = hf_config
|
|
.vocab_size
|
|
.ok_or_else(|| HubError::InvalidPackage {
|
|
reason: "Missing vocab_size in config".to_string(),
|
|
})?;
|
|
|
|
let max_seq_len = hf_config.max_position_embeddings.unwrap_or(4096);
|
|
|
|
let rope_theta = hf_config.rope_theta.unwrap_or(10000.0);
|
|
|
|
let norm_eps = hf_config
|
|
.rms_norm_eps
|
|
.or(hf_config.layer_norm_eps)
|
|
.unwrap_or(1e-5);
|
|
|
|
let activation = hf_config
|
|
.hidden_act
|
|
.as_deref()
|
|
.map_or(RTXActivation::SiLU, RTXActivation::from_hf);
|
|
|
|
let dtype = hf_config
|
|
.torch_dtype
|
|
.as_deref()
|
|
.map_or(RTXDType::F16, RTXDType::from_hf);
|
|
|
|
let tie_embeddings = hf_config.tie_word_embeddings.unwrap_or(false);
|
|
|
|
let rope_scaling = hf_config.rope_scaling.as_ref().map(|rs| RTXRopeScaling {
|
|
scaling_type: RopeScalingType::from_hf(&rs.scaling_type),
|
|
factor: rs.factor,
|
|
original_max_position: rs.original_max_position_embeddings,
|
|
});
|
|
|
|
let bos_token_id = hf_config.bos_token_id.unwrap_or(1);
|
|
let eos_token_ids = hf_config
|
|
.eos_token_id
|
|
.as_ref()
|
|
.map_or_else(|| vec![2], EosTokenId::all);
|
|
let pad_token_id = hf_config.pad_token_id;
|
|
|
|
Ok(RTXModelConfig {
|
|
architecture,
|
|
hidden_size,
|
|
num_layers,
|
|
num_heads,
|
|
num_kv_heads,
|
|
intermediate_size,
|
|
vocab_size,
|
|
max_seq_len,
|
|
rope_theta,
|
|
norm_eps,
|
|
activation,
|
|
dtype,
|
|
tie_embeddings,
|
|
rope_scaling,
|
|
bos_token_id,
|
|
eos_token_ids,
|
|
pad_token_id,
|
|
})
|
|
}
|
|
|
|
/// HuggingFace Hub client for downloading models.
|
|
pub struct HFHubClient {
|
|
/// HTTP client
|
|
client: reqwest::Client,
|
|
/// Authentication token
|
|
token: Option<String>,
|
|
/// Local cache
|
|
cache: LocalCache,
|
|
/// Base URL (can be changed for mirrors)
|
|
base_url: String,
|
|
}
|
|
|
|
impl HFHubClient {
|
|
/// Create a new HuggingFace Hub client.
|
|
pub fn new(cache_dir: PathBuf) -> Self {
|
|
Self {
|
|
client: reqwest::Client::new(),
|
|
token: std::env::var("HF_TOKEN")
|
|
.or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
|
|
.ok(),
|
|
cache: LocalCache::new(cache_dir),
|
|
base_url: HF_HUB_URL.to_string(),
|
|
}
|
|
}
|
|
|
|
/// Set authentication token.
|
|
pub fn with_token(mut self, token: impl Into<String>) -> Self {
|
|
self.token = Some(token.into());
|
|
self
|
|
}
|
|
|
|
/// Set base URL (for mirrors).
|
|
pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
|
|
self.base_url = url.into();
|
|
self
|
|
}
|
|
|
|
/// Download a model from HuggingFace Hub.
|
|
pub async fn download_model(
|
|
&self,
|
|
model_id: &str,
|
|
config: &PretrainedConfig,
|
|
) -> HubResult<PathBuf> {
|
|
let revision = config.revision.as_deref().unwrap_or("main");
|
|
let cache_path = self.cache.cache_path(model_id, revision);
|
|
|
|
// Check cache
|
|
if !config.force_download && self.cache.is_cached(&cache_path).await {
|
|
debug!("Using cached model: {}", cache_path.display());
|
|
return Ok(cache_path);
|
|
}
|
|
|
|
if config.local_files_only {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: format!("{} (local files only)", model_id),
|
|
});
|
|
}
|
|
|
|
info!("Downloading model from HuggingFace: {}", model_id);
|
|
|
|
// Create cache directory
|
|
fs::create_dir_all(&cache_path).await?;
|
|
|
|
// Download essential files
|
|
let files_to_download = [
|
|
"config.json",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"generation_config.json",
|
|
"model.safetensors.index.json",
|
|
];
|
|
|
|
for filename in &files_to_download {
|
|
if let Err(e) = self
|
|
.download_file(model_id, filename, revision, &cache_path)
|
|
.await
|
|
{
|
|
debug!("Optional file {} not found: {}", filename, e);
|
|
}
|
|
}
|
|
|
|
// Download weight files
|
|
self.download_weights(model_id, revision, &cache_path)
|
|
.await?;
|
|
|
|
Ok(cache_path)
|
|
}
|
|
|
|
/// Download a single file.
|
|
async fn download_file(
|
|
&self,
|
|
model_id: &str,
|
|
filename: &str,
|
|
revision: &str,
|
|
output_dir: &Path,
|
|
) -> HubResult<PathBuf> {
|
|
let url = format!(
|
|
"{}/{}/resolve/{}/{}",
|
|
self.base_url, model_id, revision, filename
|
|
);
|
|
|
|
let mut request = self.client.get(&url);
|
|
|
|
if let Some(ref token) = self.token {
|
|
request = request.header("Authorization", format!("Bearer {}", token));
|
|
}
|
|
|
|
let response = request.send().await.map_err(|e| HubError::DownloadFailed {
|
|
url: url.clone(),
|
|
reason: e.to_string(),
|
|
})?;
|
|
|
|
if !response.status().is_success() {
|
|
return Err(HubError::DownloadFailed {
|
|
url,
|
|
reason: format!("HTTP {}", response.status()),
|
|
});
|
|
}
|
|
|
|
let output_path = output_dir.join(filename);
|
|
let content = response
|
|
.bytes()
|
|
.await
|
|
.map_err(|e| HubError::DownloadFailed {
|
|
url: url.clone(),
|
|
reason: e.to_string(),
|
|
})?;
|
|
|
|
fs::write(&output_path, &content).await?;
|
|
debug!("Downloaded: {}", filename);
|
|
|
|
Ok(output_path)
|
|
}
|
|
|
|
/// Download model weights.
|
|
async fn download_weights(
|
|
&self,
|
|
model_id: &str,
|
|
revision: &str,
|
|
output_dir: &Path,
|
|
) -> HubResult<()> {
|
|
// Check for sharded model first
|
|
let index_path = output_dir.join("model.safetensors.index.json");
|
|
|
|
if index_path.exists() {
|
|
// Sharded model - download all shards
|
|
let index_content = fs::read_to_string(&index_path).await?;
|
|
let index: crate::safetensors::SafeTensorsIndex = serde_json::from_str(&index_content)?;
|
|
|
|
let shard_files: std::collections::HashSet<_> = index.weight_map.values().collect();
|
|
|
|
for shard_file in shard_files {
|
|
self.download_file(model_id, shard_file, revision, output_dir)
|
|
.await?;
|
|
}
|
|
} else {
|
|
// Try single file model
|
|
if let Err(e) = self
|
|
.download_file(model_id, "model.safetensors", revision, output_dir)
|
|
.await
|
|
{
|
|
warn!("SafeTensors not found, trying PyTorch format: {}", e);
|
|
self.download_file(model_id, "pytorch_model.bin", revision, output_dir)
|
|
.await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Load model configuration from cache.
|
|
pub async fn load_config(&self, model_path: &Path) -> HubResult<HFModelConfig> {
|
|
let config_path = model_path.join("config.json");
|
|
let config_content = fs::read_to_string(&config_path).await?;
|
|
let config: HFModelConfig = serde_json::from_str(&config_content)?;
|
|
Ok(config)
|
|
}
|
|
|
|
/// Load tokenizer configuration from cache.
|
|
pub async fn load_tokenizer_config(&self, model_path: &Path) -> HubResult<HFTokenizerConfig> {
|
|
let config_path = model_path.join("tokenizer_config.json");
|
|
let config_content = fs::read_to_string(&config_path).await?;
|
|
let config: HFTokenizerConfig = serde_json::from_str(&config_content)?;
|
|
Ok(config)
|
|
}
|
|
|
|
/// Load generation configuration from cache.
|
|
pub async fn load_generation_config(&self, model_path: &Path) -> HubResult<HFGenerationConfig> {
|
|
let config_path = model_path.join("generation_config.json");
|
|
let config_content = fs::read_to_string(&config_path).await?;
|
|
let config: HFGenerationConfig = serde_json::from_str(&config_content)?;
|
|
Ok(config)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_architecture_detection() {
|
|
assert_eq!(
|
|
RTXArchitecture::from_hf_model_type("llama"),
|
|
RTXArchitecture::Llama
|
|
);
|
|
assert_eq!(
|
|
RTXArchitecture::from_hf_model_type("mistral"),
|
|
RTXArchitecture::Mistral
|
|
);
|
|
assert_eq!(
|
|
RTXArchitecture::from_hf_model_type("gpt2"),
|
|
RTXArchitecture::Gpt2
|
|
);
|
|
assert_eq!(
|
|
RTXArchitecture::from_hf_model_type("bert"),
|
|
RTXArchitecture::Bert
|
|
);
|
|
assert_eq!(
|
|
RTXArchitecture::from_hf_model_type("unknown"),
|
|
RTXArchitecture::Other
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_activation_conversion() {
|
|
assert_eq!(RTXActivation::from_hf("silu"), RTXActivation::SiLU);
|
|
assert_eq!(RTXActivation::from_hf("gelu"), RTXActivation::GELU);
|
|
assert_eq!(RTXActivation::from_hf("relu"), RTXActivation::ReLU);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_conversion() {
|
|
assert_eq!(RTXDType::from_hf("float32"), RTXDType::F32);
|
|
assert_eq!(RTXDType::from_hf("float16"), RTXDType::F16);
|
|
assert_eq!(RTXDType::from_hf("bfloat16"), RTXDType::BF16);
|
|
}
|
|
|
|
#[test]
|
|
fn test_hf_config_parsing() {
|
|
let config_json = r#"{
|
|
"architectures": ["LlamaForCausalLM"],
|
|
"model_type": "llama",
|
|
"hidden_size": 4096,
|
|
"num_hidden_layers": 32,
|
|
"num_attention_heads": 32,
|
|
"num_key_value_heads": 8,
|
|
"intermediate_size": 14336,
|
|
"vocab_size": 128256,
|
|
"max_position_embeddings": 8192,
|
|
"rms_norm_eps": 1e-5,
|
|
"rope_theta": 500000.0,
|
|
"hidden_act": "silu",
|
|
"torch_dtype": "bfloat16",
|
|
"tie_word_embeddings": false,
|
|
"bos_token_id": 128000,
|
|
"eos_token_id": [128001, 128008, 128009]
|
|
}"#;
|
|
|
|
let config: HFModelConfig = serde_json::from_str(config_json).unwrap();
|
|
|
|
assert_eq!(config.model_type, Some("llama".to_string()));
|
|
assert_eq!(config.hidden_size, Some(4096));
|
|
assert_eq!(config.num_hidden_layers, Some(32));
|
|
assert_eq!(config.num_attention_heads, Some(32));
|
|
assert_eq!(config.num_key_value_heads, Some(8));
|
|
assert_eq!(config.vocab_size, Some(128256));
|
|
assert_eq!(config.rope_theta, Some(500000.0));
|
|
|
|
// Test EOS token ID
|
|
match &config.eos_token_id {
|
|
Some(EosTokenId::Multiple(ids)) => {
|
|
assert_eq!(ids.len(), 3);
|
|
assert_eq!(ids[0], 128001);
|
|
}
|
|
_ => panic!("Expected multiple EOS tokens"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_convert_hf_config() {
|
|
let hf_config = HFModelConfig {
|
|
architectures: vec!["LlamaForCausalLM".to_string()],
|
|
model_type: Some("llama".to_string()),
|
|
hidden_size: Some(4096),
|
|
num_hidden_layers: Some(32),
|
|
num_attention_heads: Some(32),
|
|
num_key_value_heads: Some(8),
|
|
intermediate_size: Some(14336),
|
|
vocab_size: Some(128256),
|
|
max_position_embeddings: Some(8192),
|
|
rms_norm_eps: Some(1e-5),
|
|
layer_norm_eps: None,
|
|
rope_theta: Some(500000.0),
|
|
rope_scaling: None,
|
|
tie_word_embeddings: Some(false),
|
|
hidden_act: Some("silu".to_string()),
|
|
torch_dtype: Some("bfloat16".to_string()),
|
|
use_cache: Some(true),
|
|
bos_token_id: Some(128000),
|
|
eos_token_id: Some(EosTokenId::Single(128001)),
|
|
pad_token_id: None,
|
|
extra: HashMap::new(),
|
|
};
|
|
|
|
let rtx_config = convert_hf_config(&hf_config).unwrap();
|
|
|
|
assert_eq!(rtx_config.architecture, RTXArchitecture::Llama);
|
|
assert_eq!(rtx_config.hidden_size, 4096);
|
|
assert_eq!(rtx_config.num_layers, 32);
|
|
assert_eq!(rtx_config.num_heads, 32);
|
|
assert_eq!(rtx_config.num_kv_heads, 8);
|
|
assert_eq!(rtx_config.intermediate_size, 14336);
|
|
assert_eq!(rtx_config.vocab_size, 128256);
|
|
assert_eq!(rtx_config.max_seq_len, 8192);
|
|
assert_eq!(rtx_config.rope_theta, 500000.0);
|
|
assert_eq!(rtx_config.activation, RTXActivation::SiLU);
|
|
assert_eq!(rtx_config.dtype, RTXDType::BF16);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tokenizer_config_parsing() {
|
|
let config_json = r#"{
|
|
"tokenizer_class": "PreTrainedTokenizerFast",
|
|
"model_max_length": 8192,
|
|
"padding_side": "left",
|
|
"chat_template": "{% for message in messages %}..."
|
|
}"#;
|
|
|
|
let config: HFTokenizerConfig = serde_json::from_str(config_json).unwrap();
|
|
|
|
assert_eq!(
|
|
config.tokenizer_class,
|
|
Some("PreTrainedTokenizerFast".to_string())
|
|
);
|
|
assert_eq!(config.model_max_length, Some(8192));
|
|
assert_eq!(config.padding_side, Some("left".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_generation_config_parsing() {
|
|
let config_json = r#"{
|
|
"max_new_tokens": 512,
|
|
"temperature": 0.7,
|
|
"top_p": 0.9,
|
|
"do_sample": true
|
|
}"#;
|
|
|
|
let config: HFGenerationConfig = serde_json::from_str(config_json).unwrap();
|
|
|
|
assert_eq!(config.max_new_tokens, Some(512));
|
|
assert_eq!(config.temperature, Some(0.7));
|
|
assert_eq!(config.top_p, Some(0.9));
|
|
assert_eq!(config.do_sample, Some(true));
|
|
}
|
|
|
|
#[test]
|
|
fn test_eos_token_id_variants() {
|
|
// Single
|
|
let single: EosTokenId = serde_json::from_str("128001").unwrap();
|
|
assert_eq!(single.primary(), 128001);
|
|
assert_eq!(single.all(), vec![128001]);
|
|
|
|
// Multiple
|
|
let multiple: EosTokenId = serde_json::from_str("[128001, 128008]").unwrap();
|
|
assert_eq!(multiple.primary(), 128001);
|
|
assert_eq!(multiple.all(), vec![128001, 128008]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_value_variants() {
|
|
// Simple
|
|
let simple: TokenValue = serde_json::from_str(r#""<|begin_of_text|>""#).unwrap();
|
|
assert_eq!(simple.content(), "<|begin_of_text|>");
|
|
|
|
// With metadata
|
|
let with_meta: TokenValue = serde_json::from_str(
|
|
r#"{"content": "<|end_of_text|>", "lstrip": false, "rstrip": false}"#,
|
|
)
|
|
.unwrap();
|
|
assert_eq!(with_meta.content(), "<|end_of_text|>");
|
|
}
|
|
|
|
#[test]
|
|
fn test_rope_scaling() {
|
|
let config_json = r#"{
|
|
"model_type": "llama",
|
|
"hidden_size": 4096,
|
|
"num_hidden_layers": 32,
|
|
"num_attention_heads": 32,
|
|
"intermediate_size": 14336,
|
|
"vocab_size": 128256,
|
|
"rope_scaling": {
|
|
"type": "dynamic",
|
|
"factor": 2.0,
|
|
"original_max_position_embeddings": 8192
|
|
}
|
|
}"#;
|
|
|
|
let config: HFModelConfig = serde_json::from_str(config_json).unwrap();
|
|
let rtx_config = convert_hf_config(&config).unwrap();
|
|
|
|
assert!(rtx_config.rope_scaling.is_some());
|
|
let scaling = rtx_config.rope_scaling.unwrap();
|
|
assert_eq!(scaling.scaling_type, RopeScalingType::Dynamic);
|
|
assert_eq!(scaling.factor, 2.0);
|
|
assert_eq!(scaling.original_max_position, Some(8192));
|
|
}
|
|
}
|