clarifai 9.8.1__py3-none-any.whl → 9.9.0__py3-none-any.whl

clarifai/client/runner.py

@@ -0,0 +1,161 @@
+# Copyright 2023 Clarifai, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Interface to Clarifai Runners API."""
+
+import os
+
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
+from google.protobuf import json_format
+
+from clarifai.client.base import BaseClient
+from clarifai.errors import UserError
+from clarifai.utils.logging import get_logger
+
+
+class Runner(BaseClient):
+  """Base class for remote inference runners. This should be subclassed with the run_input method
+  implemented to process each input in the request.
+
+  Then on the subclass call start() to start the run loop.
+  """
+
+  def __init__(self,
+               runner_id: str,
+               user_id: str = "",
+               check_runner_exists: bool = True,
+               base_url: str = "https://api.clarifai.com",
+               **kwargs) -> None:
+    """
+    Args:
+      runner_id (str): the id of the runner to use. Create the runner in the Clarifai API first
+      user_id (str): Clarifai User ID
+      base_url (str): Base API url. Default "https://api.clarifai.com"
+    """
+    user_id = user_id or os.environ.get("CLARIFAI_USER_ID", "")
+
+    if user_id == "":
+      raise UserError(
+          "Set CLARIFAI_USER_ID as environment variables or pass user_id as input arguments")
+
+    self.runner_id = runner_id
+    self.logger = get_logger("INFO", __name__)
+    self.kwargs = {**kwargs, 'id': runner_id, 'user_id': user_id}
+    self.runner_info = resources_pb2.Runner(**self.kwargs)
+    BaseClient.__init__(self, user_id=self.user_id, app_id="", base=base_url)
+
+    # Check that the runner exists.
+    if check_runner_exists:
+      request = service_pb2.GetRunnerRequest(user_app_id=self.user_app_id, runner_id=runner_id)
+      response = self._grpc_request(self.STUB.GetRunner, request)
+      if response.status.code != status_code_pb2.SUCCESS:
+        raise Exception(
+            f"""Error getting runner, are you use this is a valid runner id {runner_id} at the user_id
+            {self.user_app_id.user_id}.
+            Error: {response.status.description}""")
+
+  def start(self):
+    """Start the run loop. This will ask the Clarifai API for work, and when it gets work, it will run
+    the model on the inputs and post the results back to the Clarifai API. It will then ask for more
+    work again.
+    """
+    self._long_poll_loop()
+
+  def _run(self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
+    """Run the model on the given request. You shouldn't need to override this method, see run_input
+    for the implementation to process each input in the request.
+
+    Args:
+      request: service_pb2.PostModelOutputsRequest - the request to run the model on
+
+    Returns:
+      service_pb2.MultiOutputResponse - the response from the model's run_input implementation.
+    """
+    outputs = []
+    # TODO: parallelize this
+    for inp in request.inputs:
+      # TODO: handle errors
+      outputs.append(self.run_input(inp))
+
+    return service_pb2.MultiOutputResponse(
+        status=status_pb2.Status(
+            code=status_code_pb2.SUCCESS,
+            description="Success",
+        ),
+        outputs=outputs,
+    )
+
+  def run_input(self, input: resources_pb2.Input) -> resources_pb2.Output:
+    """Run the model on the given input in the request. This is the method you should override to
+    process each input in the request.
+
+    Args:
+      input: resources_pb2.Input - the input to run the model on
+
+    Returns:
+      resources_pb2.Output - the response from the model's run_input implementation.
+    """
+    raise NotImplementedError("run_input() not implemented")
+
+  def _long_poll_loop(self):
+    """This method will long poll for work, and when it gets work, it will run the model on the inputs
+    and post the results back to the Clarifai API. It will then long poll again for more work.
+    """
+    c = 0
+    # TODO: handle more errors within this loop so it never stops.
+    # TODO: perhaps have multiple processes running this loop to handle more work.
+    while True:
+      # Long poll waiting for work.
+      self.logger.info("Loop iteration: {}".format(c))
+      request = service_pb2.ListRunnerItemsRequest(
+          user_app_id=self.user_app_id, runner_id=self.runner_id)
+      work_response = self._grpc_request(self.STUB.ListRunnerItems, request)
+      if work_response.status.code == status_code_pb2.RUNNER_NEEDS_RETRY:
+        c += 1
+        continue  # immediate restart the long poll
+      if work_response.status.code != status_code_pb2.SUCCESS:
+        raise Exception("Error getting work: {}".format(work_response.status.description))
+      if len(work_response.items) == 0:
+        self.logger.info("No work to do. Waiting...")
+        continue
+
+      # We have work to do. Run the model on the inputs.
+      for item in work_response.items:
+        if not item.HasField('post_model_outputs_request'):
+          raise Exception("Unexpected work item type: {}".format(item))
+        self.logger.info(
+            f"Working on item: {item.id} with inputs {len(item.post_model_outputs_request.inputs)}"
+        )
+        result = self._run(item.post_model_outputs_request)
+
+        request = service_pb2.PostRunnerItemOutputsRequest(
+            user_app_id=self.user_app_id,
+            item_id=item.id,
+            runner_id=self.runner_id,
+            runner_item_outputs=[service_pb2.RunnerItemOutput(multi_output_response=result)])
+        result_response = self._grpc_request(self.STUB.PostRunnerItemOutputs, request)
+        if result_response.status.code != status_code_pb2.SUCCESS:
+          raise Exception(
+              json_format.MessageToJson(result_response, preserving_proto_field_name=True))
+          # raise Exception("Error posting result: {}".format(result_response.status.description))
+
+  def __getattr__(self, name):
+    return getattr(self.runner_info, name)
+
+  def __str__(self):
+    init_params = [param for param in self.kwargs.keys()]
+    attribute_strings = [
+        f"{param}={getattr(self.runner_info, param)}" for param in init_params
+        if hasattr(self.runner_info, param)
+    ]
+    return f"Runner Details: \n{', '.join(attribute_strings)}\n"

clarifai/client/search.py

@@ -0,0 +1,173 @@
+from typing import Any, Callable, Dict, Generator
+
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2
+from google.protobuf.json_format import MessageToDict
+from google.protobuf.struct_pb2 import Struct
+from schema import SchemaError
+
+from clarifai.client.base import BaseClient
+from clarifai.client.input import Inputs
+from clarifai.client.lister import Lister
+from clarifai.constants.search import DEFAULT_SEARCH_METRIC, DEFAULT_TOP_K
+from clarifai.errors import UserError
+from clarifai.schema.search import get_schema
+
+
+class Search(Lister, BaseClient):
+
+  def __init__(self,
+               user_id,
+               app_id,
+               top_k: int = DEFAULT_TOP_K,
+               metric: str = DEFAULT_SEARCH_METRIC):
+    """Initialize the Search object.
+
+        Args:
+            user_id (str): User ID.
+            app_id (str): App ID.
+            top_k (int, optional): Top K results to retrieve. Defaults to 10.
+            metric (str, optional): Similarity metric (either 'cosine' or 'euclidean'). Defaults to 'cosine'.
+
+        Raises:
+            UserError: If the metric is not 'cosine' or 'euclidean'.
+        """
+    if metric not in ["cosine", "euclidean"]:
+      raise UserError("Metric should be either cosine or euclidean")
+
+    self.user_id = user_id
+    self.app_id = app_id
+    self.metric_distance = dict(cosine="COSINE_DISTANCE", euclidean="EUCLIDEAN_DISTANCE")[metric]
+    self.data_proto = resources_pb2.Data()
+
+    self.inputs = Inputs(user_id=self.user_id, app_id=self.app_id)
+    self.rank_filter_schema = get_schema()
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id)
+    Lister.__init__(self, page_size=top_k)
+
+  def _get_annot_proto(self, **kwargs):
+    """Get an Annotation proto message based on keyword arguments.
+
+        Args:
+            **kwargs: Keyword arguments specifying the resource.
+
+        Returns:
+            resources_pb2.Annotation: An Annotation proto message.
+        """
+    if not kwargs:
+      return resources_pb2.Annotation()
+
+    self.data_proto = resources_pb2.Data()
+    for key, value in kwargs.items():
+      if key == "image_bytes":
+        image_proto = self.inputs.get_input_from_bytes("", image_bytes=value).data.image
+        self.data_proto.image.CopyFrom(image_proto)
+
+      elif key == "image_url":
+        image_proto = self.inputs.get_input_from_url("", image_url=value).data.image
+        self.data_proto.image.CopyFrom(image_proto)
+
+      elif key == "concepts":
+        for concept in value:
+          concept_proto = resources_pb2.Concept(**concept)
+          self.data_proto.concepts.add().CopyFrom(concept_proto)
+
+      elif key == "text_raw":
+        text_proto = self.inputs.get_input_from_bytes(
+            "", text_bytes=bytes(value, 'utf-8')).data.text
+        self.data_proto.text.CopyFrom(text_proto)
+
+      elif key == "metadata":
+        metadata_struct = Struct()
+        metadata_struct.update(value)
+        self.data_proto.metadata.CopyFrom(metadata_struct)
+
+      elif key == "geo_point":
+        geo_point_proto = self._get_geo_point_proto(value["longitude"], value["latitude"],
+                                                    value["geo_limit"])
+        self.data_proto.geo.CopyFrom(geo_point_proto)
+
+      else:
+        raise UserError(f"kwargs contain key that is not supported: {key}")
+    return resources_pb2.Annotation(data=self.data_proto)
+
+  def _get_geo_point_proto(self, longitude: float, latitude: float,
+                           geo_limit: float) -> resources_pb2.Geo:
+    """Get a GeoPoint proto message based on geographical data.
+
+        Args:
+            longitude (float): Longitude coordinate.
+            latitude (float): Latitude coordinate.
+            geo_limit (float): Geographical limit.
+
+        Returns:
+            resources_pb2.Geo: A Geo proto message.
+        """
+    return resources_pb2.Geo(
+        geo_point=resources_pb2.GeoPoint(longitude=longitude, latitude=latitude),
+        geo_limit=resources_pb2.GeoLimit(type="withinKilometers", value=geo_limit))
+
+  def list_all_pages_generator(
+      self, endpoint: Callable[..., Any], proto_message: Any,
+      request_data: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
+    """Lists all pages of a resource.
+
+        Args:
+            endpoint (Callable): The endpoint to call.
+            proto_message (Any): The proto message to use.
+            request_data (dict): The request data to use.
+
+        Yields:
+            response_dict: The next item in the listing.
+        """
+    page = 1
+    request_data['pagination'] = service_pb2.Pagination(page=page, per_page=self.default_page_size)
+    while True:
+      request_data['pagination'].page = page
+      response = self._grpc_request(endpoint, proto_message(**request_data))
+      dict_response = MessageToDict(response, preserving_proto_field_name=True)
+      if response.status.code != status_code_pb2.SUCCESS:
+        raise Exception(f"Listing failed with response {response!r}")
+
+      if 'hits' not in list(dict_response.keys()):
+        break
+      page += 1
+      yield response
+
+  def query(self, ranks=[{}], filters=[{}]):
+    """Perform a query with rank and filters.
+
+        Args:
+            ranks (List[Dict], optional): List of rank parameters. Defaults to [{}].
+            filters (List[Dict], optional): List of filter parameters. Defaults to [{}].
+
+        Returns:
+            Generator[Dict[str, Any], None, None]: A generator of query results.
+        """
+    try:
+      self.rank_filter_schema.validate(ranks)
+      self.rank_filter_schema.validate(filters)
+    except SchemaError as err:
+      raise UserError(f"Invalid rank or filter input: {err}")
+
+    rank_annot_proto, filters_annot_proto = [], []
+    for rank_dict in ranks:
+      rank_annot_proto.append(self._get_annot_proto(**rank_dict))
+    for filter_dict in filters:
+      filters_annot_proto.append(self._get_annot_proto(**filter_dict))
+
+    all_ranks = [resources_pb2.Rank(annotation=rank_annot) for rank_annot in rank_annot_proto]
+    all_filters = [
+        resources_pb2.Filter(annotation=filter_annot) for filter_annot in filters_annot_proto
+    ]
+
+    request_data = dict(
+        user_app_id=self.user_app_id,
+        searches=[
+            resources_pb2.Search(
+                query=resources_pb2.Query(ranks=all_ranks, filters=all_filters),
+                metric=self.metric_distance)
+        ])
+
+    return self.list_all_pages_generator(self.STUB.PostAnnotationsSearches,
+                                         service_pb2.PostAnnotationsSearchesRequest, request_data)

clarifai/client/user.py

@@ -1,28 +1,32 @@
 from typing import Any, Dict, List
 
-from clarifai_grpc.grpc.api import resources_pb2, service_pb2  # noqa: F401
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2
+from google.protobuf.json_format import MessageToDict
 
 from clarifai.client.app import App
 from clarifai.client.base import BaseClient
 from clarifai.client.lister import Lister
+from clarifai.client.runner import Runner
+from clarifai.errors import UserError
 from clarifai.utils.logging import get_logger
 
 
 class User(Lister, BaseClient):
   """User is a class that provides access to Clarifai API endpoints related to user information."""
 
-  def __init__(self, user_id: str = "", **kwargs):
+  def __init__(self, user_id: str = "", base_url: str = "https://api.clarifai.com", **kwargs):
     """Initializes an User object.
 
     Args:
         user_id (str): The user ID for the user to interact with.
-        **kwargs: Additional keyword arguments to be passed to the ClarifaiAuthHelper.
+        base_url (str): Base API url. Default "https://api.clarifai.com"
+        **kwargs: Additional keyword arguments to be passed to the User.
     """
     self.kwargs = {**kwargs, 'id': user_id}
     self.user_info = resources_pb2.User(**self.kwargs)
     self.logger = get_logger(logger_level="INFO", name=__name__)
-    BaseClient.__init__(self, user_id=self.id, app_id="")
+    BaseClient.__init__(self, user_id=self.id, app_id="", base=base_url)
     Lister.__init__(self)
 
   def list_apps(self, filter_by: Dict[str, Any] = {}) -> List[App]:
@@ -45,6 +49,27 @@ class User(Lister, BaseClient):
 
     return [App(**app_info) for app_info in all_apps_info]
 
+  def list_runners(self, filter_by: Dict[str, Any] = {}) -> List[Runner]:
+    """List all runners for the user
+
+    Args:
+        filter_by (dict): A dictionary of filters to apply to the list of runners.
+
+    Returns:
+        List[Runner]: A list of Runner objects for the runners.
+
+    Example:
+        >>> from clarifai.client.user import User
+        >>> client = User(user_id="user_id")
+        >>> all_runners= client.list_runners()
+    """
+    request_data = dict(user_app_id=self.user_app_id, per_page=self.default_page_size, **filter_by)
+    all_runners_info = list(
+        self.list_all_pages_generator(self.STUB.ListRunners, service_pb2.ListRunnersRequest,
+                                      request_data))
+
+    return [Runner(check_runner_exists=False, **runner_info) for runner_info in all_runners_info]
+
   def create_app(self, app_id: str, base_workflow: str = 'Language-Understanding',
                  **kwargs) -> App:
     """Creates an app for the user.
@@ -73,6 +98,41 @@ class User(Lister, BaseClient):
     kwargs.update({'user_id': self.id})
     return App(app_id=app_id, **kwargs)
 
+  def create_runner(self, runner_id: str, labels: List[str], description: str) -> Runner:
+    """Create a runner
+
+    Args:
+      runner_id (str): The Id of runner to create
+      labels (List[str]): Labels to match runner
+      description (str): Description of Runner
+
+    Returns:
+      Runner: A runner object for the specified Runner ID
+
+    Example:
+        >>> from clarifai.client.user import User
+        >>> client = User(user_id="user_id")
+        >>> runner = client.create_runner(runner_id="runner_id", labels=["label to link runner"], description="laptop runner")
+    """
+    if not isinstance(labels, List):
+      raise UserError("Labels must be a List of strings")
+
+    request = service_pb2.PostRunnersRequest(
+        user_app_id=self.user_app_id,
+        runners=[resources_pb2.Runner(id=runner_id, labels=labels, description=description)])
+    response = self._grpc_request(self.STUB.PostRunners, request)
+
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(response.status)
+    self.logger.info("\nRunner created\n%s", response.status)
+
+    return Runner(
+        runner_id=runner_id,
+        user_id=self.id,
+        labels=labels,
+        description=description,
+        check_runner_exists=False)
+
   def app(self, app_id: str, **kwargs) -> App:
     """Returns an App object for the specified app ID.
 
@@ -96,6 +156,34 @@ class User(Lister, BaseClient):
     kwargs['user_id'] = self.id
     return App(app_id=app_id, **kwargs)
 
+  def runner(self, runner_id: str) -> Runner:
+    """Returns a Runner object if exists.
+
+    Args:
+        runner_id (str): The runner ID to interact with
+
+    Returns:
+        Runner: A Runner object for the existing runner ID.
+
+    Example:
+        >>> from clarifai.client.user import User
+        >>> client = User(user_id="user_id")
+        >>> runner = client.runner(runner_id="runner_id")
+    """
+    request = service_pb2.GetRunnerRequest(user_app_id=self.user_app_id, runner_id=runner_id)
+    response = self._grpc_request(self.STUB.GetRunner, request)
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(
+          f"""Error getting runner, are you use this is a valid runner id {runner_id} at the user_id/app_id
+          {self.user_app_id.user_id}/{self.user_app_id.app_id}.
+          Error: {response.status.description}""")
+
+    dict_response = MessageToDict(response, preserving_proto_field_name=True)
+    kwargs = self.process_response_keys(dict_response[list(dict_response.keys())[1]],
+                                        list(dict_response.keys())[1])
+
+    return Runner(check_runner_exists=False, **kwargs)
+
   def delete_app(self, app_id: str) -> None:
     """Deletes an app for the user.
 
@@ -113,6 +201,24 @@ class User(Lister, BaseClient):
       raise Exception(response.status)
     self.logger.info("\nApp Deleted\n%s", response.status)
 
+  def delete_runner(self, runner_id: str) -> None:
+    """Deletes all spectified runner ids
+
+    Args:
+        runner_ids (str): List of runners to delete
+
+    Example:
+        >>> from clarifai.client.user import User
+        >>> client = User(user_id="user_id")
+        >>> client.delete_runner(runner_id="runner_id")
+    """
+    request = service_pb2.DeleteRunnersRequest(user_app_id=self.user_app_id, ids=[runner_id])
+    response = self._grpc_request(self.STUB.DeleteRunners, request)
+
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(response.status)
+    self.logger.info("\nRunner Deleted\n%s", response.status)
+
   def __getattr__(self, name):
     return getattr(self.user_info, name)
 

clarifai/client/workflow.py

@@ -10,6 +10,7 @@ from clarifai.client.lister import Lister
 from clarifai.errors import UserError
 from clarifai.urls.helper import ClarifaiUrlHelper
 from clarifai.utils.logging import get_logger
+from clarifai.workflows.export import Exporter
 
 
 class Workflow(Lister, BaseClient):
@@ -20,6 +21,7 @@ class Workflow(Lister, BaseClient):
                workflow_id: str = "",
                workflow_version: Dict = {'id': ""},
                output_config: Dict = {'min_value': 0},
+               base_url: str = "https://api.clarifai.com",
                **kwargs):
     """Initializes a Workflow object.
 
@@ -32,7 +34,8 @@ class Workflow(Lister, BaseClient):
           max_concepts (int): The maximum number of concepts to return.
           select_concepts (list[Concept]): The concepts to select.
           sample_ms (int): The number of milliseconds to sample.
-        **kwargs: Additional keyword arguments to be passed to the ClarifaiAuthHelper.
+        base_url (str): Base API url. Default "https://api.clarifai.com"
+        **kwargs: Additional keyword arguments to be passed to the Workflow.
     """
     if url_init != "" and workflow_id != "":
       raise UserError("You can only specify one of url_init or workflow_id.")
@@ -47,7 +50,7 @@ class Workflow(Lister, BaseClient):
     self.output_config = output_config
     self.workflow_info = resources_pb2.Workflow(**self.kwargs)
     self.logger = get_logger(logger_level="INFO")
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url)
     Lister.__init__(self)
 
   def predict(self, inputs: List[Input]):
@@ -182,6 +185,28 @@ class Workflow(Lister, BaseClient):
         for workflow_version_info in all_workflow_versions_info
     ]
 
+  def export(self, out_path: str):
+    """Exports the workflow to a yaml file.
+
+    Args:
+        out_path (str): The path to save the yaml file to.
+
+    Example:
+        >>> from clarifai.client.workflow import Workflow
+        >>> workflow = Workflow("https://clarifai.com/clarifai/main/workflows/Demographics")
+        >>> workflow.export('out_path.yml')
+    """
+    request = service_pb2.GetWorkflowRequest(user_app_id=self.user_app_id, workflow_id=self.id)
+    response = self._grpc_request(self.STUB.GetWorkflow, request)
+    if response.status.code != status_code_pb2.SUCCESS:
+      raise Exception(f"Workflow Export failed with response {response.status!r}")
+
+    with Exporter(response) as e:
+      e.parse()
+      e.export(out_path)
+
+    self.logger.info(f"Exported workflow to {out_path}")
+
   def __getattr__(self, name):
     return getattr(self.workflow_info, name)
 

clarifai/constants/search.py

@@ -0,0 +1,2 @@
+DEFAULT_TOP_K = 10
+DEFAULT_SEARCH_METRIC = "cosine"

clarifai/datasets/upload/loaders/xview_detection.py

@@ -42,7 +42,7 @@ class xviewDetectionDataLoader(ClarifaiDataLoader):
     self.img_comp_dir = os.path.join(self.data_dir, "train_images_comp")
     self.label_file = os.path.join(self.data_dir, "xview_train.geojson")
 
-    # self.preprocess()
+    self.preprocess()
     self.all_data = self.xview_data_parser()
 
     self.load_data()

clarifai/models/model_serving/README.md

@@ -16,10 +16,10 @@ $ clarifai-model-upload-init --model_name <Your model name> \
 3. Add your model loading and inference code inside `inference.py` script of the generated model repository under the `setup()` and `predict()` functions respectively. Refer to  The [Inference Script section]() for a description of this file.
 4. Testing your implementation locally by running `<your_triton_folder>/1/test.py` with basic predefined tests.
 To avoid missing dependencies when deploying, recommend to use conda to create clean environment from [Clarifai base envs](./envs/). Then install everything in `requirements.txt`. Follow instruction inside [test.py](./models/test.py) for implementing custom tests.
-  * Create conda env: The base envs are attached in [./envs](./envs/), these are yaml file named as `triton_conda-python_version-torch_version-xxxx.yaml` format. Make sure base env that you're about to create is matched the one in your_triton_folder/triton_conda.yaml. To create conda env and install requirements run:
+  * Create conda env and install requirements:
 ```bash
-# create env
-conda env create -n <your_env> -f <base env name>.yaml
+# create env (note: only python version 3.8 is supported currently)
+conda create -n <your_env> python=3.8
 # activate it
 conda activate <your_env>
 # install dependencies

clarifai/models/model_serving/cli/deploy_cli.py

@@ -15,8 +15,7 @@ import argparse
 
 from clarifai.auth.helper import ClarifaiAuthHelper
 from clarifai.models.api import Models
-from clarifai.models.model_serving.constants import MODEL_TYPES
-from clarifai.models.model_serving.model_config.deploy import ClarifaiFieldsMap
+from clarifai.models.model_serving.model_config import MODEL_TYPES, get_model_config
 
 
 def deploy(model_url,
@@ -50,7 +49,7 @@ def deploy(model_url,
   assert model_type, "Can not parse model_type from url, please input it directly"
   # key map
   assert model_type in MODEL_TYPES, f"model_type should be one of {MODEL_TYPES}"
-  clarifai_key_map = ClarifaiFieldsMap(model_type=model_type)
+  clarifai_key_map = get_model_config(model_type=model_type).field_maps
   # if updating new version of existing model
   if update_version:
     resp = model_api.post_model_version(

clarifai/models/model_serving/cli/repository.py

@@ -14,8 +14,8 @@
 
 import argparse
 
-from ..constants import MODEL_TYPES
-from ..model_config.triton_config import TritonModelConfig
+from ..constants import MAX_HW_DIM
+from ..model_config import MODEL_TYPES, get_model_config
 from ..pb_model_repository import TritonModelRepository
 
 
@@ -63,7 +63,6 @@ def model_upload_init():
       help="Directory to create triton repository.")
 
   args = parser.parse_args()
-  MAX_HW_DIM = 1024
 
   if len(args.image_shape) != 2:
     raise ValueError(
@@ -75,10 +74,9 @@ def model_upload_init():
         f"H and W each have a maximum value of 1024. Got H: {args.image_shape[0]}, W: {args.image_shape[1]}"
     )
 
-  model_config = TritonModelConfig(
+  model_config = get_model_config(args.model_type).make_triton_model_config(
       model_name=args.model_name,
       model_version="1",
-      model_type=args.model_type,
       image_shape=args.image_shape,
   )
 

clarifai/models/model_serving/constants.py

@@ -1,5 +1 @@
-# Clarifai model types
-MODEL_TYPES = [
-    "visual-detector", "visual-classifier", "text-classifier", "text-to-text", "text-embedder",
-    "text-to-image", "visual-embedder", "visual-segmenter"
-]
+MAX_HW_DIM = 1024

clarifai/models/model_serving/docs/custom_config.md

@@ -15,17 +15,16 @@ $ clarifai-model-upload-init --model_name <Your model name> \
 
 ## Generating the triton model repository without the commandline
 
-The triton model repository can be generated via a python script specifying the same values as required in the commandline. Below is a sample of how the code would be structured.
+The triton model repository can be generated via a python script specifying the same values as required in the commandline. Below is a sample of how the code would be structured with `visual_classifier`.
 
 ```python
-from clarifai.models.model_serving.model_config.triton_config import TritonModelConfig
-from clarifai.models.model_serving pb_model_repository import TritonModelRepository
+from clarifai.models.model_serving.model_config import get_model_config, ModelTypes, TritonModelConfig
+from clarifai.models.model_serving.pb_model_repository import TritonModelRepository
 
-
-model_config = TritonModelConfig(
+model_type = ModelTypes.visual_classifier
+model_config: TritonModelConfig = get_model_config(model_type).make_triton_model_config(
 	model_name="<model_name>",
 	model_version="1",
-	model_type="<model_type>",
 	image_shape=<[H,W]>, # 0 < [H,W] <= 1024
 )