llama-stack 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

llama_stack/providers/inline/inference/meta_reference/parallel_utils.py

@@ -33,10 +33,6 @@ from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 
 from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import GenerationResult
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    ChatCompletionRequestWithRawContent,
-    CompletionRequestWithRawContent,
-)
 
 log = get_logger(name=__name__, category="inference")
 
@@ -69,10 +65,7 @@ class CancelSentinel(BaseModel):
 
 class TaskRequest(BaseModel):
     type: Literal[ProcessingMessageName.task_request] = ProcessingMessageName.task_request
-    task: tuple[
-        str,
-        list[CompletionRequestWithRawContent] | list[ChatCompletionRequestWithRawContent],
-    ]
+    task: tuple[str, list]
 
 
 class TaskResponse(BaseModel):
@@ -328,10 +321,7 @@ class ModelParallelProcessGroup:
 
     def run_inference(
         self,
-        req: tuple[
-            str,
-            list[CompletionRequestWithRawContent] | list[ChatCompletionRequestWithRawContent],
-        ],
+        req: tuple[str, list],
     ) -> Generator:
         assert not self.running, "inference already running"
 

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -6,24 +6,20 @@
 
 from collections.abc import AsyncIterator
 
-from llama_stack.apis.inference import (
-    InferenceProvider,
-    OpenAIChatCompletionRequestWithExtraBody,
-    OpenAICompletionRequestWithExtraBody,
-)
-from llama_stack.apis.inference.inference import (
-    OpenAIChatCompletion,
-    OpenAIChatCompletionChunk,
-    OpenAICompletion,
-)
-from llama_stack.apis.models import ModelType
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.embedding_mixin import (
     SentenceTransformerEmbeddingMixin,
 )
-from llama_stack.providers.utils.inference.openai_compat import (
-    OpenAIChatCompletionToLlamaStackMixin,
+from llama_stack_api import (
+    InferenceProvider,
+    Model,
+    ModelsProtocolPrivate,
+    ModelType,
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAICompletion,
+    OpenAICompletionRequestWithExtraBody,
 )
 
 from .config import SentenceTransformersInferenceConfig
@@ -32,7 +28,6 @@ log = get_logger(name=__name__, category="inference")
 
 
 class SentenceTransformersInferenceImpl(
-    OpenAIChatCompletionToLlamaStackMixin,
     SentenceTransformerEmbeddingMixin,
     InferenceProvider,
     ModelsProtocolPrivate,

llama_stack/providers/inline/post_training/common/validator.py

@@ -12,14 +12,10 @@
 
 from typing import Any
 
-from llama_stack.apis.common.type_system import (
-    ChatCompletionInputType,
-    DialogType,
-    StringType,
-)
 from llama_stack.providers.utils.common.data_schema_validator import (
     ColumnName,
 )
+from llama_stack_api import ChatCompletionInputType, DialogType, StringType
 
 EXPECTED_DATASET_SCHEMA: dict[str, list[dict[str, Any]]] = {
     "instruct": [

llama_stack/providers/inline/post_training/huggingface/post_training.py

@@ -6,11 +6,16 @@
 from enum import Enum
 from typing import Any
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.post_training import (
+from llama_stack.providers.inline.post_training.huggingface.config import (
+    HuggingFacePostTrainingConfig,
+)
+from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
+from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
+from llama_stack_api import (
     AlgorithmConfig,
     Checkpoint,
+    DatasetIO,
+    Datasets,
     DPOAlignmentConfig,
     JobStatus,
     ListPostTrainingJobsResponse,
@@ -19,11 +24,6 @@ from llama_stack.apis.post_training import (
     PostTrainingJobStatusResponse,
     TrainingConfig,
 )
-from llama_stack.providers.inline.post_training.huggingface.config import (
-    HuggingFacePostTrainingConfig,
-)
-from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
-from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 
 
 class TrainingArtifactType(Enum):

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device.py

@@ -14,24 +14,24 @@ import torch
 from datasets import Dataset
 from peft import LoraConfig
 from transformers import (
-    AutoModelForCausalLM,
     AutoTokenizer,
 )
 from trl import SFTConfig, SFTTrainer
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.post_training import (
+from llama_stack.log import get_logger
+from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
+from llama_stack_api import (
     Checkpoint,
     DataConfig,
+    DatasetIO,
+    Datasets,
     LoraFinetuningConfig,
     TrainingConfig,
 )
-from llama_stack.log import get_logger
-from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 
 from ..config import HuggingFacePostTrainingConfig
 from ..utils import (
+    HFAutoModel,
     calculate_training_steps,
     create_checkpoints,
     get_memory_stats,
@@ -338,7 +338,7 @@ class HFFinetuningSingleDevice:
 
     def save_model(
         self,
-        model_obj: AutoModelForCausalLM,
+        model_obj: HFAutoModel,
         trainer: SFTTrainer,
         peft_config: LoraConfig | None,
         output_dir_path: Path,
@@ -350,14 +350,22 @@ class HFFinetuningSingleDevice:
             peft_config: Optional LoRA configuration
             output_dir_path: Path to save the model
         """
+        from typing import cast
+
         logger.info("Saving final model")
         model_obj.config.use_cache = True
 
         if peft_config:
             logger.info("Merging LoRA weights with base model")
-            model_obj = trainer.model.merge_and_unload()
+            # TRL's merge_and_unload returns a HuggingFace model
+            # Both cast() and type: ignore are needed here:
+            # - cast() tells mypy the return type is HFAutoModel for downstream code
+            # - type: ignore suppresses errors on the merge_and_unload() call itself,
+            #   which mypy can't type-check due to TRL library's incomplete type stubs
+            model_obj = cast(HFAutoModel, trainer.model.merge_and_unload())  # type: ignore[union-attr,operator]
         else:
-            model_obj = trainer.model
+            # trainer.model is the trained HuggingFace model
+            model_obj = cast(HFAutoModel, trainer.model)
 
         save_path = output_dir_path / "merged_model"
         logger.info(f"Saving model to {save_path}")
@@ -411,7 +419,7 @@ class HFFinetuningSingleDevice:
         # Initialize trainer
         logger.info("Initializing SFTTrainer")
         trainer = SFTTrainer(
-            model=model_obj,
+            model=model_obj,  # type: ignore[arg-type]
             train_dataset=train_dataset,
             eval_dataset=eval_dataset,
             peft_config=peft_config,

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device_dpo.py

@@ -16,15 +16,15 @@ from transformers import (
 )
 from trl import DPOConfig, DPOTrainer
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.post_training import (
+from llama_stack.log import get_logger
+from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
+from llama_stack_api import (
     Checkpoint,
+    DatasetIO,
+    Datasets,
     DPOAlignmentConfig,
     TrainingConfig,
 )
-from llama_stack.log import get_logger
-from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 
 from ..config import HuggingFacePostTrainingConfig
 from ..utils import (
@@ -309,7 +309,7 @@ class HFDPOAlignmentSingleDevice:
             save_total_limit=provider_config.save_total_limit,
             # DPO specific parameters
             beta=dpo_config.beta,
-            loss_type=provider_config.dpo_loss_type,
+            loss_type=provider_config.dpo_loss_type,  # type: ignore[arg-type]
         )
 
     def save_model(
@@ -381,13 +381,16 @@ class HFDPOAlignmentSingleDevice:
 
         # Initialize DPO trainer
         logger.info("Initializing DPOTrainer")
+        # TRL library has incomplete type stubs - use Any to bypass
+        from typing import Any, cast
+
         trainer = DPOTrainer(
-            model=model_obj,
-            ref_model=ref_model,
+            model=cast(Any, model_obj),  # HFAutoModel satisfies PreTrainedModel protocol
+            ref_model=cast(Any, ref_model),
             args=training_args,
             train_dataset=train_dataset,
             eval_dataset=eval_dataset,
-            processing_class=tokenizer,
+            processing_class=cast(Any, tokenizer),  # AutoTokenizer satisfies interface
         )
 
         try:

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -9,15 +9,33 @@ import signal
 import sys
 from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, Protocol
 
 import psutil
 import torch
 from datasets import Dataset
 from transformers import AutoConfig, AutoModelForCausalLM
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.post_training import Checkpoint, TrainingConfig
+from llama_stack_api import Checkpoint, DatasetIO, TrainingConfig
+
+if TYPE_CHECKING:
+    from transformers import PretrainedConfig
+
+
+class HFAutoModel(Protocol):
+    """Protocol describing HuggingFace AutoModel interface.
+
+    This protocol defines the common interface for HuggingFace AutoModelForCausalLM
+    and similar models, providing type safety without requiring type stubs.
+    """
+
+    config: PretrainedConfig
+    device: torch.device
+
+    def to(self, device: torch.device) -> "HFAutoModel": ...
+    def save_pretrained(self, save_directory: str | Path) -> None: ...
+
+
 from llama_stack.log import get_logger
 
 from .config import HuggingFacePostTrainingConfig
@@ -132,7 +150,7 @@ def load_model(
     model: str,
     device: torch.device,
     provider_config: HuggingFacePostTrainingConfig,
-) -> AutoModelForCausalLM:
+) -> HFAutoModel:
     """Load and initialize the model for training.
     Args:
         model: The model identifier to load
@@ -143,6 +161,8 @@ def load_model(
     Raises:
         RuntimeError: If model loading fails
     """
+    from typing import cast
+
     logger.info("Loading the base model")
     try:
         model_config = AutoConfig.from_pretrained(model, **provider_config.model_specific_config)
@@ -154,9 +174,10 @@ def load_model(
             **provider_config.model_specific_config,
         )
         # Always move model to specified device
-        model_obj = model_obj.to(device)
+        model_obj = model_obj.to(device)  # type: ignore[arg-type]
         logger.info(f"Model loaded and moved to device: {model_obj.device}")
-        return model_obj
+        # Cast to HFAutoModel protocol - transformers models satisfy this interface
+        return cast(HFAutoModel, model_obj)
     except Exception as e:
         raise RuntimeError(f"Failed to load model: {str(e)}") from e
 

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

@@ -91,7 +91,7 @@ class TorchtuneCheckpointer:
         if checkpoint_format == "meta" or checkpoint_format is None:
             self._save_meta_format_checkpoint(model_file_path, state_dict, adapter_only)
         elif checkpoint_format == "huggingface":
-            # Note: for saving hugging face format checkpoints, we only suppport saving adapter weights now
+            # Note: for saving hugging face format checkpoints, we only support saving adapter weights now
             self._save_hf_format_checkpoint(model_file_path, state_dict)
         else:
             raise ValueError(f"Unsupported checkpoint format: {format}")

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

@@ -21,9 +21,9 @@ 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
 
-from llama_stack.apis.post_training import DatasetFormat
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.models.llama.sku_types import Model
+from llama_stack_api import DatasetFormat
 
 BuildLoraModelCallable = Callable[..., torch.nn.Module]
 BuildTokenizerCallable = Callable[..., Llama3Tokenizer]

llama_stack/providers/inline/post_training/torchtune/datasets/format_adapter.py

@@ -25,7 +25,7 @@ def llama_stack_instruct_to_torchtune_instruct(
     )
     input_messages = json.loads(sample[ColumnName.chat_completion_input.value])
 
-    assert len(input_messages) == 1, "llama stack intruct dataset format only supports 1 user message"
+    assert len(input_messages) == 1, "llama stack instruct dataset format only supports 1 user message"
     input_message = input_messages[0]
 
     assert "content" in input_message, "content not found in input message"

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

@@ -6,11 +6,16 @@
 from enum import Enum
 from typing import Any
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.post_training import (
+from llama_stack.providers.inline.post_training.torchtune.config import (
+    TorchtunePostTrainingConfig,
+)
+from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
+from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
+from llama_stack_api import (
     AlgorithmConfig,
     Checkpoint,
+    DatasetIO,
+    Datasets,
     DPOAlignmentConfig,
     JobStatus,
     ListPostTrainingJobsResponse,
@@ -20,11 +25,6 @@ from llama_stack.apis.post_training import (
     PostTrainingJobStatusResponse,
     TrainingConfig,
 )
-from llama_stack.providers.inline.post_training.torchtune.config import (
-    TorchtunePostTrainingConfig,
-)
-from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
-from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 
 
 class TrainingArtifactType(Enum):

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -32,17 +32,6 @@ from torchtune.training.lr_schedulers import get_cosine_schedule_with_warmup
 from torchtune.training.metric_logging import DiskLogger
 from tqdm import tqdm
 
-from llama_stack.apis.common.training_types import PostTrainingMetric
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.post_training import (
-    Checkpoint,
-    DataConfig,
-    LoraFinetuningConfig,
-    OptimizerConfig,
-    QATFinetuningConfig,
-    TrainingConfig,
-)
 from llama_stack.core.utils.config_dirs import DEFAULT_CHECKPOINT_DIR
 from llama_stack.core.utils.model_utils import model_local_dir
 from llama_stack.log import get_logger
@@ -56,6 +45,17 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
     TorchtunePostTrainingConfig,
 )
 from llama_stack.providers.inline.post_training.torchtune.datasets.sft import SFTDataset
+from llama_stack_api import (
+    Checkpoint,
+    DataConfig,
+    DatasetIO,
+    Datasets,
+    LoraFinetuningConfig,
+    OptimizerConfig,
+    PostTrainingMetric,
+    QATFinetuningConfig,
+    TrainingConfig,
+)
 
 log = get_logger(name=__name__, category="post_training")
 
@@ -193,7 +193,7 @@ class LoraFinetuningSingleDevice:
         log.info("Optimizer is initialized.")
 
         self._loss_fn = CEWithChunkedOutputLoss()
-        self._model.set_num_output_chunks(self._loss_fn.num_output_chunks)
+        self._model.set_num_output_chunks(self._loss_fn.num_output_chunks)  # type: ignore[operator]
         log.info("Loss is initialized.")
 
         assert isinstance(self.training_config.data_config, DataConfig), "DataConfig must be initialized"
@@ -284,7 +284,7 @@ class LoraFinetuningSingleDevice:
         if self._is_dora:
             for m in model.modules():
                 if hasattr(m, "initialize_dora_magnitude"):
-                    m.initialize_dora_magnitude()
+                    m.initialize_dora_magnitude()  # type: ignore[operator]
         if lora_weights_state_dict:
             lora_missing, lora_unexpected = model.load_state_dict(lora_weights_state_dict, strict=False)
         else:
@@ -353,7 +353,7 @@ class LoraFinetuningSingleDevice:
             dataset_type=self._data_format.value,
         )
 
-        sampler = DistributedSampler(
+        sampler: DistributedSampler = DistributedSampler(
             ds,
             num_replicas=1,
             rank=0,
@@ -389,7 +389,7 @@ class LoraFinetuningSingleDevice:
             num_training_steps=num_training_steps,
             last_epoch=last_epoch,
         )
-        return lr_scheduler
+        return lr_scheduler  # type: ignore[no-any-return]
 
     async def save_checkpoint(self, epoch: int) -> str:
         ckpt_dict = {}
@@ -447,7 +447,7 @@ class LoraFinetuningSingleDevice:
         # free logits otherwise it peaks backward memory
         del logits
 
-        return loss
+        return loss  # type: ignore[no-any-return]
 
     async def train(self) -> tuple[dict[str, Any], list[Checkpoint]]:
         """

llama_stack/providers/inline/safety/code_scanner/code_scanner.py

@@ -10,19 +10,20 @@ from typing import TYPE_CHECKING, Any
 if TYPE_CHECKING:
     from codeshield.cs import CodeShieldScanResult
 
-from llama_stack.apis.inference import OpenAIMessageParam
-from llama_stack.apis.safety import (
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    interleaved_content_as_str,
+)
+from llama_stack_api import (
+    ModerationObject,
+    ModerationObjectResults,
+    OpenAIMessageParam,
     RunShieldResponse,
     Safety,
     SafetyViolation,
+    Shield,
     ViolationLevel,
 )
-from llama_stack.apis.safety.safety import ModerationObject, ModerationObjectResults
-from llama_stack.apis.shields import Shield
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    interleaved_content_as_str,
-)
 
 from .config import CodeScannerConfig
 
@@ -101,7 +102,10 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
             metadata=metadata,
         )
 
-    async def run_moderation(self, input: str | list[str], model: str) -> ModerationObject:
+    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
+        if model is None:
+            raise ValueError("Code scanner moderation requires a model identifier.")
+
         inputs = input if isinstance(input, list) else [input]
         results = []
 

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -9,29 +9,29 @@ import uuid
 from string import Template
 from typing import Any
 
-from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem
-from llama_stack.apis.inference import (
+from llama_stack.core.datatypes import Api
+from llama_stack.log import get_logger
+from llama_stack.models.llama.datatypes import Role
+from llama_stack.models.llama.sku_types import CoreModelId
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    interleaved_content_as_str,
+)
+from llama_stack_api import (
+    ImageContentItem,
     Inference,
+    ModerationObject,
+    ModerationObjectResults,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIMessageParam,
     OpenAIUserMessageParam,
-)
-from llama_stack.apis.safety import (
     RunShieldResponse,
     Safety,
     SafetyViolation,
+    Shield,
+    ShieldsProtocolPrivate,
+    TextContentItem,
     ViolationLevel,
 )
-from llama_stack.apis.safety.safety import ModerationObject, ModerationObjectResults
-from llama_stack.apis.shields import Shield
-from llama_stack.core.datatypes import Api
-from llama_stack.log import get_logger
-from llama_stack.models.llama.datatypes import Role
-from llama_stack.models.llama.sku_types import CoreModelId
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    interleaved_content_as_str,
-)
 
 from .config import LlamaGuardConfig
 
@@ -200,7 +200,10 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
 
         return await impl.run(messages)
 
-    async def run_moderation(self, input: str | list[str], model: str) -> ModerationObject:
+    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
+        if model is None:
+            raise ValueError("Llama Guard moderation requires a model identifier.")
+
         if isinstance(input, list):
             messages = input.copy()
         else:

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -9,20 +9,20 @@ from typing import Any
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-from llama_stack.apis.inference import OpenAIMessageParam
-from llama_stack.apis.safety import (
+from llama_stack.core.utils.model_utils import model_local_dir
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
+from llama_stack_api import (
+    ModerationObject,
+    OpenAIMessageParam,
     RunShieldResponse,
     Safety,
     SafetyViolation,
+    Shield,
+    ShieldsProtocolPrivate,
     ShieldStore,
     ViolationLevel,
 )
-from llama_stack.apis.safety.safety import ModerationObject
-from llama_stack.apis.shields import Shield
-from llama_stack.core.utils.model_utils import model_local_dir
-from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
-from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
 
 from .config import PromptGuardConfig, PromptGuardType
 
@@ -63,7 +63,7 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
 
         return await self.shield.run(messages)
 
-    async def run_moderation(self, input: str | list[str], model: str) -> ModerationObject:
+    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
         raise NotImplementedError("run_moderation is not implemented for Prompt Guard")
 
 

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

@@ -5,21 +5,17 @@
 # the root directory of this source tree.
 from typing import Any
 
-from llama_stack.apis.datasetio import DatasetIO
-from llama_stack.apis.datasets import Datasets
-from llama_stack.apis.scoring import (
+from llama_stack_api import (
+    DatasetIO,
+    Datasets,
     ScoreBatchResponse,
     ScoreResponse,
     Scoring,
+    ScoringFn,
+    ScoringFnParams,
+    ScoringFunctionsProtocolPrivate,
     ScoringResult,
 )
-from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
-from llama_stack.core.datatypes import Api
-from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate
-from llama_stack.providers.utils.common.data_schema_validator import (
-    get_valid_schemas,
-    validate_dataset_schema,
-)
 
 from .config import BasicScoringConfig
 from .scoring_fn.docvqa_scoring_fn import DocVQAScoringFn
@@ -83,9 +79,6 @@ class BasicScoringImpl(
         scoring_functions: dict[str, ScoringFnParams | None] = None,
         save_results_dataset: bool = False,
     ) -> ScoreBatchResponse:
-        dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        validate_dataset_schema(dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value))
-
         all_rows = await self.datasetio_api.iterrows(
             dataset_id=dataset_id,
             limit=-1,

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

@@ -8,9 +8,8 @@ import json
 import re
 from typing import Any
 
-from llama_stack.apis.scoring import ScoringResultRow
-from llama_stack.apis.scoring_functions import ScoringFnParams
 from llama_stack.providers.utils.scoring.base_scoring_fn import RegisteredBaseScoringFn
+from llama_stack_api import ScoringFnParams, ScoringResultRow
 
 from .fn_defs.docvqa import docvqa
 

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

@@ -6,9 +6,8 @@
 
 from typing import Any
 
-from llama_stack.apis.scoring import ScoringResultRow
-from llama_stack.apis.scoring_functions import ScoringFnParams
 from llama_stack.providers.utils.scoring.base_scoring_fn import RegisteredBaseScoringFn
+from llama_stack_api import ScoringFnParams, ScoringResultRow
 
 from .fn_defs.equality import equality
 

llama_stack/providers/inline/scoring/basic/scoring_fn/fn_defs/docvqa.py

@@ -4,10 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.apis.common.type_system import NumberType
-from llama_stack.apis.scoring_functions import (
+from llama_stack_api import (
     AggregationFunctionType,
     BasicScoringFnParams,
+    NumberType,
     ScoringFn,
 )