use std::path::Path;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LocalFileStatus {
Ready,
Missing,
Partial,
Invalid,
CannotCheck,
}
impl LocalFileStatus {
pub fn user_label(&self) -> &'static str {
match self {
LocalFileStatus::Ready => "Ready",
LocalFileStatus::Missing => "Needed",
LocalFileStatus::Partial => "Needs to finish",
LocalFileStatus::Invalid => "Needs to be replaced",
LocalFileStatus::CannotCheck => "Could not check",
}
}
pub fn needs_work(&self) -> bool {
!matches!(self, LocalFileStatus::Ready)
}
}
#[derive(Debug, Clone)]
pub struct FileReadiness {
pub logical_name: &'static str,
pub relative_path: &'static str,
pub status: LocalFileStatus,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelReadiness {
Ready,
NeedsDownload,
NeedsRepair,
CannotCheck,
}
#[derive(Debug, Clone)]
pub struct ModelReadinessReport {
pub overall: ModelReadiness,
pub files: Vec<FileReadiness>,
}
impl ModelReadinessReport {
pub fn files_needing_work(&self) -> Vec<&FileReadiness> {
self.files
.iter()
.filter(|f| f.status.needs_work())
.collect()
}
pub fn ready_count(&self) -> usize {
self.files
.iter()
.filter(|f| f.status == LocalFileStatus::Ready)
.count()
}
pub fn total_count(&self) -> usize {
self.files.len()
}
}
const REQUIRED_FILES: &[(&str, &str)] = &[
("tokenizer", "tokenizer.json"),
("model", "onnx/model.onnx"),
];
pub fn check_model_readiness(model_dir: &Path) -> ModelReadinessReport {
let mut files = Vec::new();
for (logical_name, relative_path) in REQUIRED_FILES {
let full_path = model_dir.join(relative_path);
let part_path = model_dir.join(format!("{relative_path}.part"));
let status = check_single_file(&full_path, &part_path);
files.push(FileReadiness {
logical_name,
relative_path,
status,
});
}
let overall = derive_overall_readiness(&files);
ModelReadinessReport { overall, files }
}
fn check_single_file(path: &Path, part_path: &Path) -> LocalFileStatus {
match std::fs::metadata(path) {
Ok(meta) => {
if !meta.is_file() {
return LocalFileStatus::Invalid;
}
if meta.len() == 0 {
return LocalFileStatus::Invalid;
}
LocalFileStatus::Ready
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
if part_path.exists() {
LocalFileStatus::Partial
} else {
LocalFileStatus::Missing
}
}
Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => LocalFileStatus::CannotCheck,
Err(_) => LocalFileStatus::CannotCheck,
}
}
fn derive_overall_readiness(files: &[FileReadiness]) -> ModelReadiness {
if files
.iter()
.any(|f| f.status == LocalFileStatus::CannotCheck)
{
return ModelReadiness::CannotCheck;
}
if files.iter().all(|f| f.status == LocalFileStatus::Ready) {
return ModelReadiness::Ready;
}
if files.iter().any(|f| f.status == LocalFileStatus::Missing) {
return ModelReadiness::NeedsDownload;
}
ModelReadiness::NeedsRepair
}