space_worker/model_store.py

from pathlib import Path


PROTEXA_INSTRUCTION_EXPECTED_FILES = (
    ".gitattributes",
    "README.md",
    "checkpoint_best.pt",
    "instruction_audit.json",
    "metrics_history.jsonl",
    "tokenizer_map.json",
    "train.txt",
)


class ModelStore:
    def __init__(self, settings):
        self.settings = settings
        self.root = Path(settings.model_bucket_mount_path)
        self.protexa_instruction_path = Path(settings.protexa_instruction_model_dir)

    def path_for_model(self, model_name):
        return self.root / model_name

    def describe(self):
        missing_instruction_files = [
            name
            for name in PROTEXA_INSTRUCTION_EXPECTED_FILES
            if not (self.protexa_instruction_path / name).exists()
        ]

        return {
            "bucket_id": self.settings.model_bucket_id,
            "mount_path": str(self.root),
            "mount_exists": self.root.exists(),
            "protexa_instruction_repo_id": self.settings.protexa_instruction_repo_id,
            "protexa_instruction_path": str(self.protexa_instruction_path),
            "protexa_instruction_exists": self.protexa_instruction_path.exists(),
            "protexa_instruction_ready": len(missing_instruction_files) == 0,
            "missing_instruction_files": missing_instruction_files,
        }