File size: 6,533 Bytes
60b21d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | # SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT
import torch
from typing import Optional, Union
from enum import Enum
from huggingface_hub import hf_hub_download
from .OpenTSLMSP import OpenTSLMSP
from .OpenTSLMFlamingo import OpenTSLMFlamingo
class ModelType(Enum):
"""Enumeration of supported model types."""
SP = "sp"
FLAMINGO = "flamingo"
class OpenTSLM:
"""
Factory class for loading EmbedHealth models from Hugging Face Hub.
Automatically detects model type based on repository ID suffix and returns
the appropriate model instance (EmbedHealthSP or EmbedHealthFlamingo) with
optimal parameters from curriculum learning training.
- Repository IDs ending with "-sp" load EmbedHealthSP models
- Repository IDs ending with "-flamingo" load EmbedHealthFlamingo models
The factory automatically applies the exact same parameters used in curriculum learning:
- EmbedHealthSP: Uses default constructor parameters
- EmbedHealthFlamingo: cross_attn_every_n_layers=1, gradient_checkpointing=False
These parameters are fixed and cannot be overridden since they were determined during training.
Example:
>>> model = OpenTSLM.load_pretrained("OpenTSLM/gemma-3-270m-pt-sleep-flamingo")
>>>
>>> from opentslm.prompt.full_prompt import FullPrompt
>>> prompt = FullPrompt(...)
>>> response = model.eval_prompt(prompt)
"""
@classmethod
def load_pretrained(
cls,
repo_id: str,
device: Optional[str] = None,
cache_dir: Optional[str] = None,
enable_lora: Optional[bool] = False,
) -> Union[OpenTSLMSP, OpenTSLMFlamingo]:
"""
Load a pretrained model from Hugging Face Hub.
Args:
repo_id: Hugging Face repository ID (e.g., "OpenTSLM/gemma-3-270m-pt-sleep-flamingo")
device: Device to load the model on (default: auto-detect)
cache_dir: Directory to cache downloaded models (optional)
enable_lora: Whether to enable LoRA (default: False)
Returns:
Union[OpenTSLMSP, OpenTSLMFlamingo]: The loaded model instance
Example:
>>> model = OpenTSLM.load_pretrained("OpenTSLM/gemma-3-270m-pt-sleep-flamingo")
>>> prompt = FullPrompt(...)
>>> response = model.eval_prompt(prompt)
"""
device = cls._get_device(device)
model_type = cls._detect_model_type(repo_id)
checkpoint_path = cls._download_model_files(repo_id, cache_dir)
base_llm_id = cls._get_base_llm_id(repo_id)
print(f"🚀 Loading {model_type.value.upper()} model...")
print(f" Repository: {repo_id}")
print(f" Base LLM: {base_llm_id}")
print(f" Device: {device}")
# Instantiate model with fixed training parameters
if model_type == ModelType.SP:
# OpenTSLMSP uses default parameters from curriculum learning
model = OpenTSLMSP(llm_id=base_llm_id, device=device)
if enable_lora:
model.enable_lora()
elif model_type == ModelType.FLAMINGO:
# OpenTSLMFlamingo with fixed parameters from curriculum learning
model = OpenTSLMFlamingo(
device=device,
llm_id=base_llm_id,
cross_attn_every_n_layers=1,
gradient_checkpointing=False,
)
else:
raise ValueError(f"Unknown model type: {model_type}")
# Load the checkpoint
model.load_from_file(checkpoint_path)
model.eval()
print(f"✅ {model_type.value.upper()} model loaded successfully!")
return model
@staticmethod
def _get_device(device: Optional[str]) -> str:
"""Auto-detect device if not specified."""
if device is not None:
return device
if torch.cuda.is_available():
return "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return "mps"
else:
return "cpu"
@staticmethod
def _detect_model_type(repo_id: str) -> ModelType:
"""Detect model type from repository ID suffix."""
if repo_id.endswith("-sp"):
return ModelType.SP
elif repo_id.endswith("-flamingo"):
return ModelType.FLAMINGO
else:
raise ValueError(
f"Repository ID '{repo_id}' must end with either '-sp' or '-flamingo' "
f"to indicate the model type."
)
@staticmethod
def _download_model_files(repo_id: str, cache_dir: Optional[str] = None) -> str:
"""Download model checkpoint from Hugging Face Hub."""
try:
# Download the main model checkpoint file
checkpoint_path = hf_hub_download(
repo_id=repo_id,
filename="model_checkpoint.pt",
cache_dir=cache_dir,
local_files_only=False,
)
print(f"✅ Downloaded model checkpoint from {repo_id}")
return checkpoint_path
except Exception as e:
raise RuntimeError(
f"Failed to download model from {repo_id}. "
f"Tried 'model_checkpoint.pt'. "
f"Original error: {e}"
)
@staticmethod
def _get_base_llm_id(repo_id: str) -> str:
"""Get the base LLM ID from static mapping based on repository ID pattern."""
repo_name = repo_id.split("/")[-1] if "/" in repo_id else repo_id
# Extract base model from repository name pattern
if repo_name.startswith("llama-3.2-3b"):
return "meta-llama/Llama-3.2-3B"
elif repo_name.startswith("llama-3.2-1b"):
return "meta-llama/Llama-3.2-1B"
elif repo_name.startswith("gemma-3-1b"):
return "google/gemma-3-1b"
elif repo_name.startswith("gemma-3-270m"):
return "google/gemma-3-270m"
else:
# Raise exception if pattern doesn't match
raise ValueError(
f"Unable to determine base LLM ID from repository name '{repo_name}'. "
f"Repository name must start with one of: 'llama-3.2-3b', 'llama-3.2-1b', "
f"'gemma-3-1b', or 'gemma-3-270m'."
)
|