chore: Refactor model_path to model_uri (#464)

In the codebase, we currently name terms like hf://... as "model paths" and "URIs." This inconsistency can cause confusion and make the code harder to read and maintain. Refactor `model_path` to `model_uri`. Closes #428 Signed-off-by: Dimitris Poulopoulos <[email protected]>
mozilla-ai · Dec 11, 2024 · d665e68 · d665e68
1 parent 30c5ba6
commit d665e68
Show file tree

Hide file tree

Showing 11 changed files with 42 additions and 40 deletions.
diff --git a/lumigator/python/mzai/backend/backend/config_templates.py b/lumigator/python/mzai/backend/backend/config_templates.py
@@ -3,7 +3,7 @@
 
 seq2seq_eval_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
+    "model": {{ "path": "{model_uri}" }},
     "dataset": {{ "path": "{dataset_path}" }},
     "evaluation": {{
         "metrics": ["rouge", "meteor", "bertscore"],
@@ -17,8 +17,8 @@
 
 bart_eval_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
-    "tokenizer": {{ "path": "{model_path}", "mod_max_length": 1024 }},
+    "model": {{ "path": "{model_uri}" }},
+    "tokenizer": {{ "path": "{model_uri}", "mod_max_length": 1024 }},
     "dataset": {{ "path": "{dataset_path}" }},
     "evaluation": {{
         "metrics": ["rouge", "meteor", "bertscore"],
@@ -32,7 +32,7 @@
 
 causal_eval_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
+    "model": {{ "path": "{model_uri}" }},
     "dataset": {{ "path": "{dataset_path}" }},
     "evaluation": {{
         "metrics": ["rouge", "meteor", "bertscore"],
@@ -49,7 +49,7 @@
     "model": {{
         "inference": {{
             "base_url": "{model_url}",
-            "engine": "{model_path}",
+            "engine": "{model_uri}",
             "system_prompt": "{system_prompt}",
             "max_retries": 3
         }}
@@ -70,7 +70,7 @@
     "name": "{job_name}/{job_id}",
     "dataset": {{ "path": "{dataset_path}" }},
     "hf_pipeline": {{
-        "model_path": "{model_path}",
+        "model_uri": "{model_uri}",
         "task": "{task}",
         "accelerator": "{accelerator}",
         "revision": "{revision}",
@@ -86,20 +86,20 @@
 
 seq2seq_infer_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
+    "model": {{ "path": "{model_uri}" }},
     "dataset": {{ "path": "{dataset_path}" }},
 }}"""
 
 bart_infer_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
-    "tokenizer": {{ "path": "{model_path}", "mod_max_length": 1024 }},
+    "model": {{ "path": "{model_uri}" }},
+    "tokenizer": {{ "path": "{model_uri}", "mod_max_length": 1024 }},
     "dataset": {{ "path": "{dataset_path}" }},
 }}"""
 
 causal_infer_template = """{{
     "name": "{job_name}/{job_id}",
-    "model": {{ "path": "{model_path}" }},
+    "model": {{ "path": "{model_uri}" }},
     "dataset": {{ "path": "{dataset_path}" }},
 }}"""
 
@@ -114,7 +114,7 @@
     }},
     "inference_server": {{
         "base_url": "{model_url}",
-        "engine": "{model_path}",
+        "engine": "{model_uri}",
         "system_prompt": "{system_prompt}",
         "max_retries": 3
     }},

diff --git a/lumigator/python/mzai/backend/backend/services/jobs.py b/lumigator/python/mzai/backend/backend/services/jobs.py
@@ -130,7 +130,7 @@ def _get_job_params(self, job_type: str, record, request: BaseModel) -> dict:
             job_params = {
                 "job_id": record.id,
                 "job_name": request.name,
-                "model_path": request.model,
+                "model_uri": request.model,
                 "dataset_path": dataset_s3_path,
                 "max_samples": request.max_samples,
                 "storage_path": self.storage_path,
@@ -141,7 +141,7 @@ def _get_job_params(self, job_type: str, record, request: BaseModel) -> dict:
             job_params = {
                 "job_id": record.id,
                 "job_name": request.name,
-                "model_path": request.model,
+                "model_uri": request.model,
                 "dataset_path": dataset_s3_path,
                 "task": request.task,
                 "accelerator": request.accelerator,

diff --git a/lumigator/python/mzai/jobs/evaluator/evaluator/configs/jobs/hf_evaluate.py b/lumigator/python/mzai/jobs/evaluator/evaluator/configs/jobs/hf_evaluate.py
@@ -18,7 +18,7 @@ class HuggingFaceEvaluationConfig(EvaluatorConfig):
     metrics: conlist(str, min_length=1)
     use_pipeline: bool = False
     enable_tqdm: bool = False
-    max_samples: int = -1 # set to all samples by default
+    max_samples: int = -1  # set to all samples by default
     storage_path: str | None = None
     return_input_data: bool = False
     return_predictions: bool = False
@@ -41,8 +41,8 @@ def ensure_tokenizer_config(cls, values):
         if values.get("tokenizer") is None:
             values["tokenizer"] = {}
             match values["model"]:
-                case str() as model_path:
-                    values["tokenizer"]["path"] = model_path
+                case str() as model_uri:
+                    values["tokenizer"]["path"] = model_uri
                 case dict() as model_data:
                     # if dict we might have model.path specified
                     # if we don't it is VLLMCompletion and we are ok

diff --git a/lumigator/python/mzai/jobs/evaluator/evaluator/jobs/asset_loader.py b/lumigator/python/mzai/jobs/evaluator/evaluator/jobs/asset_loader.py
@@ -49,12 +49,14 @@ def resolve_asset_path(self, path: AssetPath) -> str:
         """
         raw_path = strip_path_prefix(path)
         if path.startswith(
-    (
-        PathPrefix.FILE, PathPrefix.HUGGINGFACE, PathPrefix.OPENAI, 
-        PathPrefix.MISTRAL, PathPrefix.LLAMAFILE
-    )
-):
-
+            (
+                PathPrefix.FILE,
+                PathPrefix.HUGGINGFACE,
+                PathPrefix.OPENAI,
+                PathPrefix.MISTRAL,
+                PathPrefix.LLAMAFILE,
+            )
+        ):
             return raw_path
         elif path.startswith(PathPrefix.WANDB):
             artifact = get_artifact_from_api(raw_path)
@@ -136,7 +138,7 @@ def load_pretrained_model(
 
         # TODO: HuggingFace has many AutoModel classes with different "language model heads"
         #   Can we abstract this to load with any type of AutoModel class?
-        model_path = self.resolve_asset_path(config.path)
+        model_uri = self.resolve_asset_path(config.path)
 
         # load config first to get the model type
         model_config = self.load_pretrained_config(config)
@@ -150,7 +152,7 @@ def load_pretrained_model(
             automodel_class = AutoModelForCausalLM
 
         return automodel_class.from_pretrained(
-            pretrained_model_name_or_path=model_path,
+            pretrained_model_name_or_path=model_uri,
             trust_remote_code=config.trust_remote_code,
             torch_dtype=config.torch_dtype,
             quantization_config=bnb_config,

diff --git a/lumigator/python/mzai/jobs/evaluator/evaluator/jobs/evaluation/lm_harness.py b/lumigator/python/mzai/jobs/evaluator/evaluator/jobs/evaluation/lm_harness.py
@@ -43,14 +43,14 @@ def load_harness_model(config: LMHarnessJobConfig) -> HFLM | OpenaiCompletionsLM
     hf_model_loader = HuggingFaceModelLoader()
     match config.model:
         case AutoModelConfig() as model_config:
-            model_path, peft_path = hf_model_loader.resolve_peft_and_pretrained(model_config.path)
+            model_uri, peft_path = hf_model_loader.resolve_peft_and_pretrained(model_config.path)
             quantization_kwargs: dict[str, Any] = (
                 config.quantization.model_dump() if config.quantization else {}
             )
             # TODO: Fix this up by passing in the instantiated model directly
             return HFLM(
-                pretrained=model_path,
-                tokenizer=model_path,
+                pretrained=model_uri,
+                tokenizer=model_uri,
                 peft=peft_path,
                 device="cuda" if torch.cuda.device_count() > 0 else "cpu",
                 trust_remote_code=model_config.trust_remote_code,

diff --git a/lumigator/python/mzai/jobs/evaluator/evaluator/paths.py b/lumigator/python/mzai/jobs/evaluator/evaluator/paths.py
@@ -116,13 +116,13 @@ def format_s3_path(bucket: str, key: str) -> AssetPath:
     return f"{PathPrefix.S3.value}{bucket}/{key}"
 
 
-def format_openai_model_path(model_name: str) -> AssetPath:
+def format_openai_model_uri(model_name: str) -> AssetPath:
     return f"{PathPrefix.OPENAI.value}{model_name}"
 
 
-def format_mistral_model_path(model_name: str) -> AssetPath:
+def format_mistral_model_uri(model_name: str) -> AssetPath:
     return f"{PathPrefix.MISTRAL.value}{model_name}"
 
 
-def format_llamafile_model_path(model_name: str) -> AssetPath:
+def format_llamafile_model_uri(model_name: str) -> AssetPath:
     return f"{PathPrefix.LLAMAFILE.value}{model_name}"
diff --git a/lumigator/python/mzai/jobs/inference/inference.py b/lumigator/python/mzai/jobs/inference/inference.py
@@ -104,7 +104,7 @@ def run_inference(config: InferenceJobConfig) -> Path:
             logger.info(f"Using OAI client. Endpoint: {base_url}")
             model_client = OpenAIModelClient(base_url, config)
     elif config.hf_pipeline:
-        if config.hf_pipeline.model_path.startswith(PathPrefix.HUGGINGFACE):
+        if config.hf_pipeline.model_uri.startswith(PathPrefix.HUGGINGFACE):
             logger.info("Using HuggingFace client.")
             model_client = HuggingFaceModelClient(config)
             output_model_name = config.hf_pipeline.model

diff --git a/lumigator/python/mzai/jobs/inference/inference_config.py b/lumigator/python/mzai/jobs/inference/inference_config.py
@@ -35,13 +35,13 @@ def _validate_torch_dtype(x: str | torch.dtype) -> str | torch.dtype:
 TorchDtype = Annotated[torch.dtype | str, BeforeValidator(lambda x: _validate_torch_dtype(x))]
 
 
-def _validate_model_path(path: str) -> str:
-    """Validate the given model path.
+def _validate_model_uri(path: str) -> str:
+    """Validate the given model RI.
 
     Resorve the model repo ID and check if it is a valid HuggingFace repo ID.
 
     Args:
-        path (str): The model path to validate.
+        path (str): The model URI to validate.
 
     Raises:
         ValueError: If the path is not a valid HuggingFace repo ID.
@@ -54,7 +54,7 @@ def _validate_model_path(path: str) -> str:
     return path
 
 
-AssetPath = Annotated[str, AfterValidator(lambda x: _validate_model_path(x))]
+AssetPath = Annotated[str, AfterValidator(lambda x: _validate_model_uri(x))]
 
 
 def _validate_task(task: str) -> None:
@@ -119,7 +119,7 @@ class SamplingParameters(BaseModel):
 
 
 class HfPipelineConfig(BaseModel, arbitrary_types_allowed=True):
-    model_path: AssetPath = Field(title="The Model HF Hub repo ID", exclude=True)
+    model_uri: AssetPath = Field(title="The Model HF Hub repo ID", exclude=True)
     task: SupportedTask
     revision: str = "main"  # Model version: branch, tag, or commit ID
     use_fast: bool = True  # Whether or not to use a Fast tokenizer if possible
@@ -136,7 +136,7 @@ def model(self) -> str:
         Returns:
             str: The model name.
         """
-        return resolve_model_repo(self.model_path)
+        return resolve_model_repo(self.model_uri)
 
     @computed_field
     @property

diff --git a/lumigator/python/mzai/jobs/inference/tests/data/config_full_hf.json b/lumigator/python/mzai/jobs/inference/tests/data/config_full_hf.json
@@ -1,7 +1,7 @@
 {
     "name": "full_inference_config",
     "dataset": { "path": "s3://lumigator-storage/datasets/deaddead-dead-dead-dead-deaddeaddead/dataset_name.csv" },
-    "pipeline": { "model_path": "hf://facebook/bart-large-cnn", "task": "summarization" },
+    "pipeline": { "model_uri": "hf://facebook/bart-large-cnn", "task": "summarization" },
     "job": {
         "max_samples": 10,
         "storage_path": "s3://lumigator-storage/jobs/results/",

diff --git a/lumigator/python/mzai/jobs/inference/tests/test_configs.py b/lumigator/python/mzai/jobs/inference/tests/test_configs.py
@@ -77,7 +77,7 @@ def test_invalid_hf_config(json_config_minimal):
     with pytest.raises(ValidationError):
         InferenceJobConfig.model_validate(invalid_task)
 
-    invalid_model_repo_id = json_config_minimal["pipeline"]["model_path"] = "hf://foo"
+    invalid_model_repo_id = json_config_minimal["pipeline"]["model_uri"] = "hf://foo"
     with pytest.raises(ValidationError):
         InferenceJobConfig.model_validate(invalid_model_repo_id)
 

diff --git a/lumigator/python/mzai/sdk/tests/conftest.py b/lumigator/python/mzai/sdk/tests/conftest.py
@@ -165,7 +165,7 @@ def dialog_data(common_resources_dir):
 def simple_eval_template():
     return """{{
         "name": "{job_name}/{job_id}",
-        "model": {{ "path": "{model_path}" }},
+        "model": {{ "path": "{model_uri}" }},
         "dataset": {{ "path": "{dataset_path}" }},
         "evaluation": {{
             "metrics": ["meteor", "rouge"],