def initialize()

in community-content/vertex_model_garden/model_oss/open_clip/handler.py [0:0]
38 lines of code
8 McCabe index (conditional complexity)

  def initialize(self, context: Any):
    """Custom initialize."""

    properties = context.system_properties
    self.map_location = (
        "cuda"
        if torch.cuda.is_available() and properties.get("gpu_id") is not None
        else "cpu"
    )
    self.device = torch.device(
        self.map_location + ":" + str(properties.get("gpu_id"))
        if torch.cuda.is_available() and properties.get("gpu_id") is not None
        else self.map_location
    )
    self.manifest = context.manifest

    self.model_name = os.environ.get("MODEL", None)
    if not self.model_name:
      self.model_name = os.environ.get("MODEL_ID", _DEFAULT_MODEL)
    precision = os.environ.get("PRECISION", self._DEFAULT_PRECISION)
    checkpoint = os.environ.get("CHECKPOINT")
    self.task = os.environ.get("TASK", _FEATURE_EMBEDDING)
    if self.task not in _VALID_TASKS:
      raise ValueError(f"Invalid task: {self.task}.")
    logging.info(
        "Handler initializing task:%s, model:%s, precision:%s, checkpoint:%s",
        self.task,
        self.model_name,
        precision,
        checkpoint,
    )

    if fileutils.is_gcs_path(checkpoint):
      local_fname = os.path.join(constants.LOCAL_MODEL_DIR, "model.pt")
      fileutils.download_gcs_file_to_local(checkpoint, local_fname)
      checkpoint = local_fname

    self.model, self.preprocessor = open_clip.create_model_from_pretrained(
        self.model_name, pretrained=checkpoint, precision=precision
    )
    self.model.to(self.device)
    self.tokenizer = open_clip.get_tokenizer(self.model_name)

    self.initialized = True