clarifai 11.3.0rc2__py3-none-any.whl → 11.4.0__py3-none-any.whl

clarifai/runners/utils/loader.py

@@ -11,236 +11,256 @@ from clarifai.utils.logging import logger
 
 
 class HuggingFaceLoader:
-
-  HF_DOWNLOAD_TEXT = "The 'huggingface_hub' package is not installed. Please install it using 'pip install huggingface_hub'."
-
-  def __init__(self, repo_id=None, token=None, model_type_id=None):
-    self.repo_id = repo_id
-    self.token = token
-    self.clarifai_model_type_id = model_type_id
-    if token:
-      if self.validate_hftoken(token):
+    HF_DOWNLOAD_TEXT = "The 'huggingface_hub' package is not installed. Please install it using 'pip install huggingface_hub'."
+
+    def __init__(self, repo_id=None, token=None, model_type_id=None):
+        self.repo_id = repo_id
+        self.token = token
+        self.clarifai_model_type_id = model_type_id
+        if token:
+            if self.validate_hftoken(token):
+                try:
+                    from huggingface_hub import login
+                except ImportError:
+                    raise ImportError(self.HF_DOWNLOAD_TEXT)
+                login(token=token)
+                logger.info("Hugging Face token validated")
+            else:
+                self.token = None
+                logger.info("Continuing without Hugging Face token")
+
+    @classmethod
+    def validate_hftoken(cls, hf_token: str):
+        try:
+            if importlib.util.find_spec("huggingface_hub") is None:
+                raise ImportError(cls.HF_DOWNLOAD_TEXT)
+            from huggingface_hub import HfApi
+
+            api = HfApi()
+            api.whoami(token=hf_token)
+            return True
+        except Exception as e:
+            logger.error(
+                f"Error setting up Hugging Face token, please make sure you have the correct token: {e}"
+            )
+            return False
+
+    def download_checkpoints(
+        self, checkpoint_path: str, allowed_file_patterns=None, ignore_file_patterns=None
+    ):
+        # throw error if huggingface_hub wasn't installed
         try:
-          from huggingface_hub import login
+            from huggingface_hub import snapshot_download
         except ImportError:
-          raise ImportError(self.HF_DOWNLOAD_TEXT)
-        login(token=token)
-        logger.info("Hugging Face token validated")
-      else:
-        self.token = None
-        logger.info("Continuing without Hugging Face token")
-
-  @classmethod
-  def validate_hftoken(cls, hf_token: str):
-    try:
-      if importlib.util.find_spec("huggingface_hub") is None:
-        raise ImportError(cls.HF_DOWNLOAD_TEXT)
-      from huggingface_hub import HfApi
-
-      api = HfApi()
-      api.whoami(token=hf_token)
-      return True
-    except Exception as e:
-      logger.error(
-          f"Error setting up Hugging Face token, please make sure you have the correct token: {e}")
-      return False
-
-  def download_checkpoints(self,
-                           checkpoint_path: str,
-                           allowed_file_patterns=None,
-                           ignore_file_patterns=None):
-    # throw error if huggingface_hub wasn't installed
-    try:
-      from huggingface_hub import snapshot_download
-    except ImportError:
-      raise ImportError(self.HF_DOWNLOAD_TEXT)
-    if os.path.exists(checkpoint_path) and self.validate_download(
-        checkpoint_path, allowed_file_patterns, ignore_file_patterns):
-      logger.info("Checkpoints already exist")
-      return True
-    else:
-      os.makedirs(checkpoint_path, exist_ok=True)
-      try:
-        is_hf_model_exists = self.validate_hf_model()
-        if not is_hf_model_exists:
-          logger.error("Model %s not found on Hugging Face" % (self.repo_id))
-          return False
+            raise ImportError(self.HF_DOWNLOAD_TEXT)
+        if os.path.exists(checkpoint_path) and self.validate_download(
+            checkpoint_path, allowed_file_patterns, ignore_file_patterns
+        ):
+            logger.info("Checkpoints already exist")
+            return True
+        else:
+            os.makedirs(checkpoint_path, exist_ok=True)
+            try:
+                is_hf_model_exists = self.validate_hf_model()
+                if not is_hf_model_exists:
+                    logger.error("Model %s not found on Hugging Face" % (self.repo_id))
+                    return False
+
+                self.ignore_patterns = self._get_ignore_patterns()
+                if ignore_file_patterns:
+                    if self.ignore_patterns:
+                        self.ignore_patterns.extend(ignore_file_patterns)
+                    else:
+                        self.ignore_patterns = ignore_file_patterns
+                snapshot_download(
+                    repo_id=self.repo_id,
+                    local_dir=checkpoint_path,
+                    local_dir_use_symlinks=False,
+                    allow_patterns=allowed_file_patterns,
+                    ignore_patterns=self.ignore_patterns,
+                )
+                # Remove the `.cache` folder if it exists
+                cache_path = os.path.join(checkpoint_path, ".cache")
+                if os.path.exists(cache_path) and os.path.isdir(cache_path):
+                    shutil.rmtree(cache_path)
+
+            except Exception as e:
+                logger.error(f"Error downloading model checkpoints {e}")
+                return False
+            finally:
+                is_downloaded = self.validate_download(
+                    checkpoint_path, allowed_file_patterns, ignore_file_patterns
+                )
+                if not is_downloaded:
+                    logger.error("Error validating downloaded model checkpoints")
+                    return False
+            return True
+
+    def download_config(self, checkpoint_path: str):
+        # throw error if huggingface_hub wasn't installed
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError:
+            raise ImportError(self.HF_DOWNLOAD_TEXT)
+        if os.path.exists(checkpoint_path) and os.path.exists(
+            os.path.join(checkpoint_path, 'config.json')
+        ):
+            logger.info("HF model's config.json already exists")
+            return True
+        os.makedirs(checkpoint_path, exist_ok=True)
+        try:
+            is_hf_model_exists = self.validate_hf_model()
+            if not is_hf_model_exists:
+                logger.error("Model %s not found on Hugging Face" % (self.repo_id))
+                return False
+            hf_hub_download(
+                repo_id=self.repo_id, filename='config.json', local_dir=checkpoint_path
+            )
+        except Exception as e:
+            logger.error(f"Error downloading model's config.json {e}")
+            return False
+        return True
+
+    def validate_hf_model(
+        self,
+    ):
+        # check if model exists on HF
+        try:
+            from huggingface_hub import file_exists, repo_exists
+        except ImportError:
+            raise ImportError(self.HF_DOWNLOAD_TEXT)
+        if self.clarifai_model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
+            return repo_exists(self.repo_id) and file_exists(self.repo_id, 'config.json')
+        else:
+            return repo_exists(self.repo_id)
+
+    def validate_download(
+        self, checkpoint_path: str, allowed_file_patterns: list, ignore_file_patterns: list
+    ):
+        # check if model exists on HF
+        try:
+            from huggingface_hub import list_repo_files
+        except ImportError:
+            raise ImportError(self.HF_DOWNLOAD_TEXT)
+        # Get the list of files on the repo
+        repo_files = list_repo_files(self.repo_id, token=self.token)
+
+        # Get the list of files on the repo that are allowed
+        if allowed_file_patterns:
+
+            def should_allow(file_path):
+                return any(
+                    fnmatch.fnmatch(file_path, pattern) for pattern in allowed_file_patterns
+                )
+
+            repo_files = [f for f in repo_files if should_allow(f)]
 
         self.ignore_patterns = self._get_ignore_patterns()
         if ignore_file_patterns:
-          if self.ignore_patterns:
-            self.ignore_patterns.extend(ignore_file_patterns)
-          else:
-            self.ignore_patterns = ignore_file_patterns
-        snapshot_download(
-            repo_id=self.repo_id,
-            local_dir=checkpoint_path,
-            local_dir_use_symlinks=False,
-            allow_patterns=allowed_file_patterns,
-            ignore_patterns=self.ignore_patterns)
-        # Remove the `.cache` folder if it exists
-        cache_path = os.path.join(checkpoint_path, ".cache")
-        if os.path.exists(cache_path) and os.path.isdir(cache_path):
-          shutil.rmtree(cache_path)
-
-      except Exception as e:
-        logger.error(f"Error downloading model checkpoints {e}")
-        return False
-      finally:
-        is_downloaded = self.validate_download(checkpoint_path, allowed_file_patterns,
-                                               ignore_file_patterns)
-        if not is_downloaded:
-          logger.error("Error validating downloaded model checkpoints")
-          return False
-      return True
-
-  def download_config(self, checkpoint_path: str):
-    # throw error if huggingface_hub wasn't installed
-    try:
-      from huggingface_hub import hf_hub_download
-    except ImportError:
-      raise ImportError(self.HF_DOWNLOAD_TEXT)
-    if os.path.exists(checkpoint_path) and os.path.exists(
-        os.path.join(checkpoint_path, 'config.json')):
-      logger.info("HF model's config.json already exists")
-      return True
-    os.makedirs(checkpoint_path, exist_ok=True)
-    try:
-      is_hf_model_exists = self.validate_hf_model()
-      if not is_hf_model_exists:
-        logger.error("Model %s not found on Hugging Face" % (self.repo_id))
+            if self.ignore_patterns:
+                self.ignore_patterns.extend(ignore_file_patterns)
+            else:
+                self.ignore_patterns = ignore_file_patterns
+        # Get the list of files on the repo that are not ignored
+        if getattr(self, "ignore_patterns", None):
+            patterns = self.ignore_patterns
+
+            def should_ignore(file_path):
+                return any(fnmatch.fnmatch(file_path, pattern) for pattern in patterns)
+
+            repo_files = [f for f in repo_files if not should_ignore(f)]
+
+        # Check if downloaded files match the files we expect (ignoring ignored patterns)
+        checkpoint_dir_files = []
+        for dp, dn, fn in os.walk(os.path.expanduser(checkpoint_path)):
+            checkpoint_dir_files.extend(
+                [os.path.relpath(os.path.join(dp, f), checkpoint_path) for f in fn]
+            )
+
+        # Validate by comparing file lists
+        return (
+            len(checkpoint_dir_files) >= len(repo_files)
+            and not (len(set(repo_files) - set(checkpoint_dir_files)) > 0)
+            and len(repo_files) > 0
+        )
+
+    def _get_ignore_patterns(self):
+        # check if model exists on HF
+        try:
+            from huggingface_hub import list_repo_files
+        except ImportError:
+            raise ImportError(self.HF_DOWNLOAD_TEXT)
+
+        # Get the list of files on the repo that are not ignored
+        repo_files = list_repo_files(self.repo_id, token=self.token)
+        self.ignore_patterns = None
+        if any(f.endswith(".safetensors") for f in repo_files):
+            self.ignore_patterns = [
+                "**/original/*",
+                "**/*.pth",
+                "**/*.bin",
+                "*.pth",
+                "*.bin",
+                "**/.cache/*",
+            ]
+        return self.ignore_patterns
+
+    @staticmethod
+    def validate_config(checkpoint_path: str):
+        # check if downloaded config.json exists
+        return os.path.exists(checkpoint_path) and os.path.exists(
+            os.path.join(checkpoint_path, 'config.json')
+        )
+
+    @staticmethod
+    def validate_concept(checkpoint_path: str):
+        # check if downloaded concept exists in hf model
+        config_path = os.path.join(checkpoint_path, 'config.json')
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+
+        labels = config.get('id2label', None)
+        if labels:
+            return True
         return False
-      hf_hub_download(repo_id=self.repo_id, filename='config.json', local_dir=checkpoint_path)
-    except Exception as e:
-      logger.error(f"Error downloading model's config.json {e}")
-      return False
-    return True
-
-  def validate_hf_model(self,):
-    # check if model exists on HF
-    try:
-      from huggingface_hub import file_exists, repo_exists
-    except ImportError:
-      raise ImportError(self.HF_DOWNLOAD_TEXT)
-    if self.clarifai_model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
-      return repo_exists(self.repo_id) and file_exists(self.repo_id, 'config.json')
-    else:
-      return repo_exists(self.repo_id)
-
-  def validate_download(self, checkpoint_path: str, allowed_file_patterns: list,
-                        ignore_file_patterns: list):
-    # check if model exists on HF
-    try:
-      from huggingface_hub import list_repo_files
-    except ImportError:
-      raise ImportError(self.HF_DOWNLOAD_TEXT)
-    # Get the list of files on the repo
-    repo_files = list_repo_files(self.repo_id, token=self.token)
-
-    # Get the list of files on the repo that are allowed
-    if allowed_file_patterns:
-
-      def should_allow(file_path):
-        return any(fnmatch.fnmatch(file_path, pattern) for pattern in allowed_file_patterns)
-
-      repo_files = [f for f in repo_files if should_allow(f)]
-
-    self.ignore_patterns = self._get_ignore_patterns()
-    if ignore_file_patterns:
-      if self.ignore_patterns:
-        self.ignore_patterns.extend(ignore_file_patterns)
-      else:
-        self.ignore_patterns = ignore_file_patterns
-    # Get the list of files on the repo that are not ignored
-    if getattr(self, "ignore_patterns", None):
-      patterns = self.ignore_patterns
-
-      def should_ignore(file_path):
-        return any(fnmatch.fnmatch(file_path, pattern) for pattern in patterns)
-
-      repo_files = [f for f in repo_files if not should_ignore(f)]
-
-    # Check if downloaded files match the files we expect (ignoring ignored patterns)
-    checkpoint_dir_files = []
-    for dp, dn, fn in os.walk(os.path.expanduser(checkpoint_path)):
-      checkpoint_dir_files.extend(
-          [os.path.relpath(os.path.join(dp, f), checkpoint_path) for f in fn])
-
-    # Validate by comparing file lists
-    return len(checkpoint_dir_files) >= len(repo_files) and not (
-        len(set(repo_files) - set(checkpoint_dir_files)) > 0) and len(repo_files) > 0
-
-  def _get_ignore_patterns(self):
-    # check if model exists on HF
-    try:
-      from huggingface_hub import list_repo_files
-    except ImportError:
-      raise ImportError(self.HF_DOWNLOAD_TEXT)
-
-    # Get the list of files on the repo that are not ignored
-    repo_files = list_repo_files(self.repo_id, token=self.token)
-    self.ignore_patterns = None
-    if any(f.endswith(".safetensors") for f in repo_files):
-      self.ignore_patterns = [
-          "**/original/*", "**/*.pth", "**/*.bin", "*.pth", "*.bin", "**/.cache/*"
-      ]
-    return self.ignore_patterns
-
-  @staticmethod
-  def validate_config(checkpoint_path: str):
-    # check if downloaded config.json exists
-    return os.path.exists(checkpoint_path) and os.path.exists(
-        os.path.join(checkpoint_path, 'config.json'))
-
-  @staticmethod
-  def validate_concept(checkpoint_path: str):
-    # check if downloaded concept exists in hf model
-    config_path = os.path.join(checkpoint_path, 'config.json')
-    with open(config_path, 'r') as f:
-      config = json.load(f)
-
-    labels = config.get('id2label', None)
-    if labels:
-      return True
-    return False
-
-  @staticmethod
-  def fetch_labels(checkpoint_path: str):
-    # Fetch labels for classification, detection and segmentation models
-    config_path = os.path.join(checkpoint_path, 'config.json')
-    with open(config_path, 'r') as f:
-      config = json.load(f)
-
-    labels = config['id2label']
-    return labels
-
-  @staticmethod
-  def get_huggingface_checkpoint_total_size(repo_name):
-    """
-    Fetches the JSON data for a Hugging Face model using the API with `?blobs=true`.
-    Calculates the total size from the JSON output.
-
-    Args:
-        repo_name (str): The name of the model on Hugging Face Hub. e.g. "casperhansen/llama-3-8b-instruct-awq"
-
-    Returns:
-        int: The total size in bytes.
-    """
-    try:
-      url = f"https://huggingface.co/api/models/{repo_name}?blobs=true"
-      response = requests.get(url)
-      response.raise_for_status()  # Raise an exception for bad status codes
-      json_data = response.json()
-
-      if isinstance(json_data, str):
-        data = json.loads(json_data)
-      else:
-        data = json_data
-
-      total_size = 0
-      for file in data['siblings']:
-        total_size += file['size']
-      return total_size
-    except Exception as e:
-      logger.error(f"Error fetching checkpoint size from huggingface.co: {e}")
-      return 0
+
+    @staticmethod
+    def fetch_labels(checkpoint_path: str):
+        # Fetch labels for classification, detection and segmentation models
+        config_path = os.path.join(checkpoint_path, 'config.json')
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+
+        labels = config['id2label']
+        return labels
+
+    @staticmethod
+    def get_huggingface_checkpoint_total_size(repo_name):
+        """
+        Fetches the JSON data for a Hugging Face model using the API with `?blobs=true`.
+        Calculates the total size from the JSON output.
+
+        Args:
+            repo_name (str): The name of the model on Hugging Face Hub. e.g. "casperhansen/llama-3-8b-instruct-awq"
+
+        Returns:
+            int: The total size in bytes.
+        """
+        try:
+            url = f"https://huggingface.co/api/models/{repo_name}?blobs=true"
+            response = requests.get(url)
+            response.raise_for_status()  # Raise an exception for bad status codes
+            json_data = response.json()
+
+            if isinstance(json_data, str):
+                data = json.loads(json_data)
+            else:
+                data = json_data
+
+            total_size = 0
+            for file in data['siblings']:
+                total_size += file['size']
+            return total_size
+        except Exception as e:
+            logger.error(f"Error fetching checkpoint size from huggingface.co: {e}")
+            return 0