Merge remote-tracking branch 'origin/main' into support_more_data_format

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -18,8 +18,8 @@ from llama_stack.providers.datatypes import EvalTasksProtocolPrivate
 
 from llama_stack.providers.utils.common.data_schema_validator import (
     ColumnName,
-    DataSchemaValidatorMixin,
     get_valid_schemas,
+    validate_dataset_schema,
 )
 from llama_stack.providers.utils.kvstore import kvstore_impl
 
@@ -31,7 +31,10 @@ from .config import MetaReferenceEvalConfig
 EVAL_TASKS_PREFIX = "eval_tasks:"
 
 
-class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate, DataSchemaValidatorMixin):
+class MetaReferenceEvalImpl(
+    Eval,
+    EvalTasksProtocolPrivate,
+):
     def __init__(
         self,
         config: MetaReferenceEvalConfig,
@@ -85,7 +88,7 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate, DataSchemaValidatorM
         candidate = task_config.eval_candidate
         scoring_functions = task_def.scoring_functions
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        self.validate_dataset_schema(
+        validate_dataset_schema(
             dataset_def.dataset_schema, get_valid_schemas(Api.eval.value)
         )
         all_rows = await self.datasetio_api.get_rows_paginated(

llama_stack/providers/inline/post_training/torchtune/common/checkpointer.py

@@ -90,18 +90,24 @@ class TorchtuneCheckpointer:
         model_file_path.mkdir(parents=True, exist_ok=True)
 
         # copy the related files for inference
-        shutil.copy(
-            Path.joinpath(self._checkpoint_dir, "params.json"),
-            Path.joinpath(model_file_path, "params.json"),
-        )
-        shutil.copy(
-            Path.joinpath(self._checkpoint_dir, "tokenizer.model"),
-            Path.joinpath(model_file_path, "tokenizer.model"),
-        )
-        shutil.copy(
-            Path.joinpath(self._checkpoint_dir, "orig_params.json"),
-            Path.joinpath(model_file_path, "orig_params.json"),
-        )
+        source_path = Path.joinpath(self._checkpoint_dir, "params.json")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "params.json"),
+            )
+        source_path = Path.joinpath(self._checkpoint_dir, "tokenizer.model")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "tokenizer.model"),
+            )
+        source_path = Path.joinpath(self._checkpoint_dir, "orig_params.json")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "orig_params.json"),
+            )
 
         if not adapter_only:
             model_state_dict = state_dict[training.MODEL_KEY]

llama_stack/providers/inline/post_training/torchtune/common/utils.py

@@ -29,8 +29,9 @@ from torchtune.data._messages import (
     ShareGPTToMessages,
 )
 
-from torchtune.models.llama3 import llama3_tokenizer, lora_llama3_8b
+from torchtune.models.llama3 import llama3_tokenizer
 from torchtune.models.llama3._tokenizer import Llama3Tokenizer
+from torchtune.models.llama3_1 import lora_llama3_1_8b
 from torchtune.models.llama3_2 import lora_llama3_2_3b
 from torchtune.modules.transforms import Transform
 
@@ -63,8 +64,8 @@ MODEL_CONFIGS: Dict[str, ModelConfig] = {
         tokenizer_type=llama3_tokenizer,
         checkpoint_type="LLAMA3_2",
     ),
-    "Llama-3-8B-Instruct": ModelConfig(
-        model_definition=lora_llama3_8b,
+    "Llama3.1-8B-Instruct": ModelConfig(
+        model_definition=lora_llama3_1_8b,
         tokenizer_type=llama3_tokenizer,
         checkpoint_type="LLAMA3",
     ),

llama_stack/providers/inline/scoring/basic/scoring.py

@@ -18,8 +18,8 @@ from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
 from llama_stack.distribution.datatypes import Api
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
-    DataSchemaValidatorMixin,
     get_valid_schemas,
+    validate_dataset_schema,
 )
 from .config import BasicScoringConfig
 from .scoring_fn.equality_scoring_fn import EqualityScoringFn
@@ -30,7 +30,8 @@ FIXED_FNS = [EqualityScoringFn, SubsetOfScoringFn, RegexParserScoringFn]
 
 
 class BasicScoringImpl(
-    Scoring, ScoringFunctionsProtocolPrivate, DataSchemaValidatorMixin
+    Scoring,
+    ScoringFunctionsProtocolPrivate,
 ):
     def __init__(
         self,
@@ -75,7 +76,7 @@ class BasicScoringImpl(
         save_results_dataset: bool = False,
     ) -> ScoreBatchResponse:
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        self.validate_dataset_schema(
+        validate_dataset_schema(
             dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value)
         )
 

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -35,8 +35,9 @@ from llama_stack.distribution.datatypes import Api
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
-    DataSchemaValidatorMixin,
     get_valid_schemas,
+    validate_dataset_schema,
+    validate_row_schema,
 )
 
 from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_metrics
@@ -111,7 +112,6 @@ class BraintrustScoringImpl(
     Scoring,
     ScoringFunctionsProtocolPrivate,
     NeedsRequestProviderData,
-    DataSchemaValidatorMixin,
 ):
     def __init__(
         self,
@@ -171,7 +171,7 @@ class BraintrustScoringImpl(
         await self.set_api_key()
 
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        self.validate_dataset_schema(
+        validate_dataset_schema(
             dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value)
         )
 
@@ -194,7 +194,7 @@ class BraintrustScoringImpl(
     async def score_row(
         self, input_row: Dict[str, Any], scoring_fn_identifier: Optional[str] = None
     ) -> ScoringResultRow:
-        self.validate_row_schema(input_row, get_valid_schemas(Api.scoring.value))
+        validate_row_schema(input_row, get_valid_schemas(Api.scoring.value))
         await self.set_api_key()
         assert scoring_fn_identifier is not None, "scoring_fn_identifier cannot be None"
         expected_answer = input_row["expected_answer"]

llama_stack/providers/inline/scoring/llm_as_judge/scoring.py

@@ -19,8 +19,8 @@ from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
 from llama_stack.distribution.datatypes import Api
 from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
 from llama_stack.providers.utils.common.data_schema_validator import (
-    DataSchemaValidatorMixin,
     get_valid_schemas,
+    validate_dataset_schema,
 )
 
 from .config import LlmAsJudgeScoringConfig
@@ -31,7 +31,8 @@ LLM_JUDGE_FNS = [LlmAsJudgeScoringFn]
 
 
 class LlmAsJudgeScoringImpl(
-    Scoring, ScoringFunctionsProtocolPrivate, DataSchemaValidatorMixin
+    Scoring,
+    ScoringFunctionsProtocolPrivate,
 ):
     def __init__(
         self,
@@ -79,7 +80,7 @@ class LlmAsJudgeScoringImpl(
         save_results_dataset: bool = False,
     ) -> ScoreBatchResponse:
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        self.validate_dataset_schema(
+        validate_dataset_schema(
             dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value)
         )