count prompt tokens

charts/model-engine/values_circleci.yaml

@@ -151,7 +151,7 @@ config:
       user_inference_pytorch_repository: "hosted-model-inference/async-pytorch"
       user_inference_tensorflow_repository: "hosted-model-inference/async-tensorflow-cpu"
       docker_image_layer_cache_repository: "kaniko-cache"
-      hf_user_fine_tuned_weights_prefix: "s3://$CIRCLECI_AWS_S3_BUCKET"
+      hf_user_fine_tuned_weights_prefix: "s3://$CIRCLECI_AWS_S3_BUCKET/model-weights"
 
 # Service Account
 serviceAccount:

integration_tests/test_endpoints.py

@@ -2,6 +2,7 @@
 import time
 
 import pytest
+from model_engine_server.common.env_vars import CIRCLECI
 from tenacity import RetryError, retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
 from .rest_api_utils import (
@@ -234,3 +235,17 @@ def test_sync_streaming_model_endpoint(capsys):
                 )
         finally:
             delete_model_endpoint(create_endpoint_request["name"], user)
+
+
+@pytest.mark.skipif(CIRCLECI, reason="skip on circleci since need to figure out s3 access")
+def test_models_tokenizers() -> None:
+    from model_engine_server.infra.gateways.s3_llm_artifact_gateway import S3LLMArtifactGateway
+    from model_engine_server.infra.repositories import LiveTokenizerRepository
+    from model_engine_server.infra.repositories.live_tokenizer_repository import (
+        SUPPORTED_MODELS_INFO,
+    )
+
+    llm_artifact_gateway = S3LLMArtifactGateway()
+    tokenizer_repository = LiveTokenizerRepository(llm_artifact_gateway=llm_artifact_gateway)
+    for model_name in SUPPORTED_MODELS_INFO:
+        tokenizer_repository.load_tokenizer(model_name)

model-engine/model_engine_server/api/dependencies.py

@@ -34,6 +34,7 @@
     DockerRepository,
     LLMFineTuneEventsRepository,
     ModelBundleRepository,
+    TokenizerRepository,
     TriggerRepository,
 )
 from model_engine_server.domain.services import (
@@ -87,6 +88,7 @@
     DbTriggerRepository,
     ECRDockerRepository,
     FakeDockerRepository,
+    LiveTokenizerRepository,
     RedisModelEndpointCacheRepository,
     S3FileLLMFineTuneEventsRepository,
     S3FileLLMFineTuneRepository,
@@ -134,6 +136,7 @@ class ExternalInterfaces:
     llm_artifact_gateway: LLMArtifactGateway
     cron_job_gateway: CronJobGateway
     monitoring_metrics_gateway: MonitoringMetricsGateway
+    tokenizer_repository: TokenizerRepository
 
 
 def get_default_monitoring_metrics_gateway() -> MonitoringMetricsGateway:
@@ -260,6 +263,8 @@ def _get_external_interfaces(
 
     docker_repository = ECRDockerRepository() if not CIRCLECI else FakeDockerRepository()
 
+    tokenizer_repository = LiveTokenizerRepository(llm_artifact_gateway=llm_artifact_gateway)
+
     external_interfaces = ExternalInterfaces(
         docker_repository=docker_repository,
         model_bundle_repository=model_bundle_repository,
@@ -281,6 +286,7 @@ def _get_external_interfaces(
         trigger_repository=trigger_repository,
         cron_job_gateway=cron_job_gateway,
         monitoring_metrics_gateway=monitoring_metrics_gateway,
+        tokenizer_repository=tokenizer_repository,
     )
     return external_interfaces
 

model-engine/model_engine_server/api/llms_v1.py

@@ -247,6 +247,7 @@ async def create_completion_sync_task(
         use_case = CompletionSyncV1UseCase(
             model_endpoint_service=external_interfaces.model_endpoint_service,
             llm_model_endpoint_service=external_interfaces.llm_model_endpoint_service,
+            tokenizer_repository=external_interfaces.tokenizer_repository,
         )
         return await use_case.execute(
             user=auth, model_endpoint_name=model_endpoint_name, request=request
@@ -290,6 +291,7 @@ async def create_completion_stream_task(
     use_case = CompletionStreamV1UseCase(
         model_endpoint_service=external_interfaces.model_endpoint_service,
         llm_model_endpoint_service=external_interfaces.llm_model_endpoint_service,
+        tokenizer_repository=external_interfaces.tokenizer_repository,
     )
     response = use_case.execute(user=auth, model_endpoint_name=model_endpoint_name, request=request)
 

model-engine/model_engine_server/common/dtos/llms.py

@@ -80,7 +80,7 @@ class GetLLMModelEndpointV1Response(BaseModel):
     """
 
     name: str
-    model_name: Optional[str] = None
+    model_name: str
     source: LLMSource
     status: ModelEndpointStatus
     inference_framework: LLMInferenceFramework
@@ -143,6 +143,7 @@ class TokenOutput(BaseModel):
 
 class CompletionOutput(BaseModel):
     text: str
+    num_prompt_tokens: int
     num_completion_tokens: int
     tokens: Optional[List[TokenOutput]] = None
 
@@ -198,6 +199,7 @@ class CompletionStreamV1Request(BaseModel):
 class CompletionStreamOutput(BaseModel):
     text: str
     finished: bool
+    num_prompt_tokens: Optional[int] = None
     num_completion_tokens: Optional[int] = None
     token: Optional[TokenOutput] = None
 

model-engine/model_engine_server/domain/gateways/llm_artifact_gateway.py

@@ -11,12 +11,31 @@ class LLMArtifactGateway(ABC):
     def list_files(self, path: str, **kwargs) -> List[str]:
         """
         Gets a list of files from a given path.
+
+        Args:
+            path (str): path to list files
+        """
+        pass
+
+    @abstractmethod
+    def download_files(self, path: str, target_path: str, overwrite=False, **kwargs) -> List[str]:
+        """
+        Download files from a given path to a target path.
+
+        Args:
+            path (str): path to list files
+            target_path (str): local path to download files
+            overwrite (bool): whether to overwrite existing local files
         """
         pass
 
     @abstractmethod
     def get_model_weights_urls(self, owner: str, model_name: str, **kwargs) -> List[str]:
         """
         Gets a list of URLs for all files associated with a given model.
+
+        Args:
+            owner (str): owner of the model
+            model_name (str): name of the model
         """
         pass

model-engine/model_engine_server/domain/repositories/__init__.py

@@ -4,12 +4,14 @@
 from .docker_repository import DockerRepository
 from .llm_fine_tune_events_repository import LLMFineTuneEventsRepository
 from .model_bundle_repository import ModelBundleRepository
+from .tokenizer_repository import TokenizerRepository
 from .trigger_repository import TriggerRepository
 
 __all__: Sequence[str] = [
     "DockerRepository",
     "DockerImageBatchJobBundleRepository",
     "LLMFineTuneEventsRepository",
     "ModelBundleRepository",
+    "TokenizerRepository",
     "TriggerRepository",
 ]

model-engine/model_engine_server/domain/repositories/tokenizer_repository.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+from transformers import AutoTokenizer
+
+
+class TokenizerRepository(ABC):
+    @abstractmethod
+    def load_tokenizer(self, model_name: str) -> AutoTokenizer:
+        """
+        Loads a tokenizer from a model name.
+
+        Args:
+            model_name: The model name to load the tokenizer for.
+
+        Returns:
+            A tokenizer.
+        """
+        pass