chore: Refactor model_path to model_uri

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 2, 2024 · cb1a228 · cb1a228
1 parent c9ca168
commit cb1a228
Show file tree

Hide file tree

Showing 7 changed files with 32 additions and 30 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
         }}
@@ -68,20 +68,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}" }},
 }}"""
 
@@ -96,7 +96,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,
                 "max_samples": request.max_samples,
                 "storage_path": self.storage_path,

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/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"],