diff --git a/src/heretic/config.py b/src/heretic/config.py index 6189a31..55bdc3c 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -88,6 +88,11 @@ class BenchmarkSpecification(BaseModel): class Settings(BaseSettings): model: str = Field(description="Hugging Face model ID, or path to model on disk.") + model_commit: str | None = Field( + default=None, + description="Hugging Face commit hash of the model.", + ) + evaluate_model: str | None = Field( default=None, description=( diff --git a/src/heretic/main.py b/src/heretic/main.py index 427ce83..a2d207f 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -66,7 +66,7 @@ from .utils import ( ) -def obtain_merge_strategy(settings: Settings) -> str | None: +def obtain_merge_strategy(settings: Settings, model: Model) -> str | None: """ Prompts the user for how to proceed with saving the model. Provides info to the user if the model is quantized on memory use. @@ -95,7 +95,8 @@ def obtain_merge_strategy(settings: Settings) -> str | None: settings.model, device_map="meta", torch_dtype=torch.bfloat16, - trust_remote_code=True, + trust_remote_code=model.trusted_models.get(settings.model), + **model.revision_kwargs, ) footprint_bytes = meta_model.get_memory_footprint() footprint_gb = footprint_bytes / (1024**3) @@ -752,7 +753,7 @@ def run(): if not save_directory: continue - strategy = obtain_merge_strategy(settings) + strategy = obtain_merge_strategy(settings, model) if strategy is None: continue @@ -803,7 +804,7 @@ def run(): continue private = visibility == "Private" - strategy = obtain_merge_strategy(settings) + strategy = obtain_merge_strategy(settings, model) if strategy is None: continue diff --git a/src/heretic/model.py b/src/heretic/model.py index 52e6add..753a8bd 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -62,12 +62,17 @@ class Model: self.settings = settings self.needs_reload = False + self.revision_kwargs = {} + if settings.model_commit is not None: + self.revision_kwargs["revision"] = settings.model_commit + print() print(f"Loading model [bold]{settings.model}[/]...") self.tokenizer = AutoTokenizer.from_pretrained( settings.model, trust_remote_code=settings.trust_remote_code, + **self.revision_kwargs, ) # Fallback for tokenizers that don't declare a special pad token. @@ -108,6 +113,7 @@ class Model: device_map=settings.device_map, max_memory=self.max_memory, trust_remote_code=self.trusted_models.get(settings.model), + **self.revision_kwargs, **extra_kwargs, ) @@ -257,6 +263,7 @@ class Model: torch_dtype=self.model.dtype, device_map="cpu", trust_remote_code=self.trusted_models.get(self.settings.model), + **self.revision_kwargs, ) # Apply LoRA adapters to the CPU model @@ -318,6 +325,7 @@ class Model: device_map=self.settings.device_map, max_memory=self.max_memory, trust_remote_code=self.trusted_models.get(self.settings.model), + **self.revision_kwargs, **extra_kwargs, ) diff --git a/src/heretic/utils.py b/src/heretic/utils.py index b2f2619..c2970a5 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -357,7 +357,6 @@ def generate_reproduce_readme( checkpoint_filename: str, trial: Trial, timestamp: str | None = None, - base_model_commit: str | None = None, ) -> str: """Generates the contents of a README.md for the reproduce/ folder.""" @@ -404,7 +403,7 @@ def generate_reproduce_readme( > This system installed `heretic-llm` from an unknown non-standard source. **Reproducibility ***cannot*** be guaranteed in this environment.** """ - model_link = format_hf_link(settings.model, base_model_commit) + model_link = format_hf_link(settings.model, settings.model_commit) dataset_info = f"""## Dataset Information - **Good Prompts:** {format_hf_link(settings.good_prompts.dataset, settings.good_prompts.commit, is_dataset=True)} @@ -518,7 +517,6 @@ def generate_reproduce_json( settings: Settings, trial: Trial, timestamp: str | None = None, - base_model_commit: str | None = None, uploaded_model_hashes: dict[str, str] | None = None, ) -> str: """Generates the contents of a reproduce.json file for the reproduce/ folder.""" @@ -528,11 +526,6 @@ def generate_reproduce_json( data = { "version": "1", # Version number of the reproduce.json file format, to allow for future changes. "timestamp": timestamp, - # TODO: Remove this, it's redundant with settings! - "base_model": { - "id": settings.model, - "commit": base_model_commit, - }, "system": { "python": get_python_env_info_dict(), "os": { @@ -592,6 +585,9 @@ def create_reproduce_folder( checkpoint_filename = Path(checkpoint_path).name + # Fetch commit hash for the base model. + settings.model_commit = huggingface_hub.model_info(settings.model).sha + # Fetch commit hashes for all HF datasets to ensure reproducibility. for spec in [ settings.good_prompts, @@ -601,9 +597,6 @@ def create_reproduce_folder( ]: spec.commit = huggingface_hub.dataset_info(spec.dataset).sha - # Fetch commit hash for the base model. - base_model_commit = huggingface_hub.model_info(settings.model).sha - # Strip microseconds and timezone for a clean format. timestamp = ( datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None).isoformat() @@ -623,7 +616,6 @@ def create_reproduce_folder( checkpoint_filename, trial, timestamp=timestamp, - base_model_commit=base_model_commit, ), encoding="utf-8", ) @@ -637,7 +629,6 @@ def create_reproduce_folder( settings, trial, timestamp=timestamp, - base_model_commit=base_model_commit, uploaded_model_hashes=uploaded_model_hashes, ), encoding="utf-8",