ImageProcessorPil` and warn once.
+ if "torchvision" in missing_backends and name.endswith("ImageProcessor"):
+ pil_name = f"{name}Pil"
+ if pil_name in self._class_to_module and pil_name not in self._object_missing_backend:
+ try:
+ pil_module = self._get_module(self._class_to_module[pil_name])
+ pil_value = getattr(pil_module, pil_name)
+ logger.warning_once(
+ f"`{name}` requires torchvision (not installed); falling back to `{pil_name}` "
+ f"for backward compatibility. Install torchvision to use the default backend, "
+ f"or import `{pil_name}` directly to silence this warning."
+ )
+ setattr(self, name, pil_value)
+ return pil_value
+ except Exception as e:
+ logger.debug(f"Could not load PIL fallback {pil_name}: {e}")
+
+ class Placeholder(metaclass=DummyObject):
+ _backends = missing_backends
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, missing_backends)
+
+ def call(self, *args, **kwargs):
+ pass
+
+ Placeholder.__name__ = name
+
+ if name not in self._class_to_module:
+ module_name = f"transformers.{name}"
+ else:
+ module_name = self._class_to_module[name]
+ if not module_name.startswith("transformers."):
+ module_name = f"transformers.{module_name}"
+
+ Placeholder.__module__ = module_name
+
+ value = Placeholder
+ elif name in self._class_to_module:
+ try:
+ module = self._get_module(self._class_to_module[name])
+ value = getattr(module, name)
+ except (ModuleNotFoundError, RuntimeError, AttributeError) as e:
+ # V5: If trying to import a *TokenizerFast symbol, transparently fall back to the
+ # non-Fast symbol from the same module when available. This lets us keep only one
+ # backend tokenizer class while preserving legacy public names.
+ if name.endswith("TokenizerFast"):
+ fallback_name = name[:-4]
+ # Prefer importing the module that declares the fallback symbol if known
+ try:
+ if fallback_name in self._class_to_module:
+ fb_module = self._get_module(self._class_to_module[fallback_name])
+ fallback_value = getattr(fb_module, fallback_name)
+ else:
+ module = self._get_module(self._class_to_module[name])
+ fallback_value = getattr(module, fallback_name)
+ setattr(self, fallback_name, fallback_value)
+ value = fallback_value
+ except Exception:
+ # If we can't find the fallback here, try converter logic as a last resort
+ # before giving up
+ value = None
+ # Try converter mapping for Fast tokenizers that don't exist
+ if value is None and name.endswith("TokenizerFast"):
+ lookup_name = name[:-4]
+ try:
+ from ..convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
+
+ if lookup_name in SLOW_TO_FAST_CONVERTERS:
+ converter_class = SLOW_TO_FAST_CONVERTERS[lookup_name]
+ converter_base_name = converter_class.__name__.replace("Converter", "")
+ preferred_tokenizer_name = f"{converter_base_name}Tokenizer"
+
+ candidate_names = [preferred_tokenizer_name]
+ for tokenizer_name, tokenizer_converter in SLOW_TO_FAST_CONVERTERS.items():
+ if tokenizer_converter is converter_class and tokenizer_name != lookup_name:
+ if tokenizer_name not in candidate_names:
+ candidate_names.append(tokenizer_name)
+
+ # Try to import the preferred candidate directly
+ import importlib
+
+ for candidate_name in candidate_names:
+ base_tokenizer_class = None
+
+ # Try to derive module path from tokenizer name (e.g., "AlbertTokenizer" -> "albert")
+ # Remove "Tokenizer" suffix and convert to lowercase
+ if candidate_name.endswith("Tokenizer"):
+ model_name = candidate_name[:-10].lower() # Remove "Tokenizer"
+ module_path = f"transformers.models.{model_name}.tokenization_{model_name}"
+ try:
+ module = importlib.import_module(module_path)
+ base_tokenizer_class = getattr(module, candidate_name)
+ except Exception:
+ logger.debug(f"{module_path} does not have {candidate_name} defined.")
+
+ # Fallback: try via _class_to_module
+ if base_tokenizer_class is None and candidate_name in self._class_to_module:
+ try:
+ alias_module_name = self._class_to_module[candidate_name]
+ alias_module = self._get_module(alias_module_name)
+ base_tokenizer_class = getattr(alias_module, candidate_name)
+ except Exception:
+ logger.debug(
+ f"{alias_module_name} does not have {candidate_name} defined"
+ )
+
+ # If we still don't have base_tokenizer_class, skip this candidate
+ if base_tokenizer_class is None:
+ logger.debug(f"skipping candidate {candidate_name}")
+ continue
+
+ # If we got here, we have base_tokenizer_class
+ value = base_tokenizer_class
+
+ setattr(self, candidate_name, base_tokenizer_class)
+ if lookup_name != candidate_name:
+ setattr(self, lookup_name, value)
+ setattr(self, name, value)
+ break
+ except Exception as e:
+ logger.debug(f"Could not create tokenizer alias: {e}")
+
+ if value is None:
+ raise ModuleNotFoundError(
+ f"Could not import module '{name}'. Are this object's requirements defined correctly?"
+ ) from e
+ else:
+ raise ModuleNotFoundError(
+ f"Could not import module '{name}'. Are this object's requirements defined correctly?"
+ ) from e
+
+ elif name in self._modules:
+ try:
+ value = self._get_module(name)
+ except (ModuleNotFoundError, RuntimeError) as e:
+ raise ModuleNotFoundError(
+ f"Could not import module '{name}'. Are this object's requirements defined correctly?"
+ ) from e
+ else:
+ # V5: If a *TokenizerFast symbol is requested but not present in the import structure,
+ # try to resolve to the corresponding non-Fast symbol's module if available.
+ if name.endswith("TokenizerFast"):
+ fallback_name = name[:-4]
+ if fallback_name in self._class_to_module:
+ try:
+ fb_module = self._get_module(self._class_to_module[fallback_name])
+ value = getattr(fb_module, fallback_name)
+ setattr(self, fallback_name, value)
+ setattr(self, name, value)
+ return value
+ except Exception as e:
+ logger.debug(f"Could not load fallback {fallback_name}: {e}")
+ # V5: Handle *ImageProcessorFast backward compatibility
+ # Similar to TokenizerFast, but for image processors
+ if name.endswith("ImageProcessorFast"):
+ fallback_name = name[:-4] # Remove "Fast"
+ if fallback_name in self._class_to_module:
+ logger.warning_once(
+ f"`{name}` is deprecated. The `Fast` suffix for image processors has been removed; "
+ f"use `{fallback_name}` instead."
+ )
+ if fallback_name in self._object_missing_backend:
+ # The Fast alias has no entry in the import structure, so `requires_backends` on
+ # the real class never runs. Handle the missing backend explicitly here, otherwise
+ # `_get_module` swallows the ImportError and the caller gets an AttributeError.
+ # Do not fall through to the PIL fallback since a legacy "Fast" image processor was explicitly requested.
+ missing_backends = self._object_missing_backend[fallback_name]
+
+ class Placeholder(metaclass=DummyObject):
+ _backends = missing_backends
+
+ def __init__(self, *args, **kwargs):
+ requires_backends(self, missing_backends)
+
+ def call(self, *args, **kwargs):
+ pass
+
+ Placeholder.__name__ = fallback_name
+ module_name = self._class_to_module[fallback_name]
+ Placeholder.__module__ = (
+ module_name if module_name.startswith("transformers.") else f"transformers.{module_name}"
+ )
+ setattr(self, name, Placeholder)
+ return Placeholder
+ try:
+ fb_module = self._get_module(self._class_to_module[fallback_name])
+ value = getattr(fb_module, fallback_name)
+ setattr(self, fallback_name, value)
+ setattr(self, name, value)
+ return value
+ except Exception as e:
+ logger.debug(f"Could not load fallback {fallback_name}: {e}")
+ # V5: If a tokenizer class doesn't exist, check if it should alias to another tokenizer
+ # via the converter mapping (e.g., FNetTokenizer -> AlbertTokenizer via AlbertConverter)
+ value = None
+ if name.endswith("Tokenizer") or name.endswith("TokenizerFast"):
+ # Strip "Fast" suffix for converter lookup if present
+ lookup_name = name[:-4] if name.endswith("TokenizerFast") else name
+
+ try:
+ # Lazy import to avoid circular dependencies
+ from ..convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
+
+ # Check if this tokenizer has a converter mapping
+ if lookup_name in SLOW_TO_FAST_CONVERTERS:
+ converter_class = SLOW_TO_FAST_CONVERTERS[lookup_name]
+
+ # Find which tokenizer class uses the same converter (reverse lookup)
+ # Prefer the tokenizer that matches the converter name pattern
+ # (e.g., AlbertConverter -> AlbertTokenizer)
+ converter_base_name = converter_class.__name__.replace("Converter", "")
+ preferred_tokenizer_name = f"{converter_base_name}Tokenizer"
+
+ # Try preferred tokenizer first
+ candidate_names = [preferred_tokenizer_name]
+ # Then try all other tokenizers with the same converter
+ for tokenizer_name, tokenizer_converter in SLOW_TO_FAST_CONVERTERS.items():
+ if tokenizer_converter is converter_class and tokenizer_name != lookup_name:
+ if tokenizer_name not in candidate_names:
+ candidate_names.append(tokenizer_name)
+
+ # Try to import one of the candidate tokenizers
+ for candidate_name in candidate_names:
+ if candidate_name in self._class_to_module:
+ try:
+ alias_module = self._get_module(self._class_to_module[candidate_name])
+ base_tokenizer_class = getattr(alias_module, candidate_name)
+ value = base_tokenizer_class
+
+ # Cache both names for future imports
+ setattr(self, candidate_name, base_tokenizer_class)
+ if lookup_name != candidate_name:
+ setattr(self, lookup_name, value)
+ setattr(self, name, value)
+ break
+ except Exception:
+ # If this candidate fails, try the next one
+ continue
+ else:
+ # Candidate not in _class_to_module - might need recursive resolution
+ # Try importing it directly to trigger lazy loading
+ try:
+ # Try to get it from transformers module to trigger lazy loading
+ transformers_module = sys.modules.get("transformers")
+ if transformers_module and hasattr(transformers_module, candidate_name):
+ base_tokenizer_class = getattr(transformers_module, candidate_name)
+ value = base_tokenizer_class
+
+ if lookup_name != candidate_name:
+ setattr(self, lookup_name, value)
+ setattr(self, name, value)
+ break
+ except Exception:
+ continue
+ except (ImportError, AttributeError):
+ pass
+
+ if value is None:
+ for key, values in self._explicit_import_shortcut.items():
+ if name in values:
+ value = self._get_module(key)
+ break
+
+ if value is None:
+ raise AttributeError(f"module {self.__name__} has no attribute {name}")
+
+ setattr(self, name, value)
+ return value
+
+ def _get_module(self, module_name: str):
+ try:
+ return importlib.import_module("." + module_name, self.__name__)
+ except Exception as e:
+ raise e
+
+ def __reduce__(self):
+ return (self.__class__, (self._name, self.__file__, self._import_structure))
+
+
+class OptionalDependencyNotAvailable(BaseException):
+ """Internally used error class for signalling an optional dependency was not found."""
+
+
+def direct_transformers_import(path: str, file="__init__.py") -> ModuleType:
+ """Imports transformers directly
+
+ Args:
+ path (`str`): The path to the source file
+ file (`str`, *optional*): The file to join with the path. Defaults to "__init__.py".
+
+ Returns:
+ `ModuleType`: The resulting imported module
+ """
+ name = "transformers"
+ location = os.path.join(path, file)
+ spec = importlib.util.spec_from_file_location(name, location, submodule_search_locations=[path])
+ if spec is not None and spec.loader is not None:
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ module = sys.modules[name]
+ return module
+ raise ImportError(f"Could not load module {name} from {location}")
+
+
+class VersionComparison(Enum):
+ EQUAL = operator.eq
+ NOT_EQUAL = operator.ne
+ GREATER_THAN = operator.gt
+ LESS_THAN = operator.lt
+ GREATER_THAN_OR_EQUAL = operator.ge
+ LESS_THAN_OR_EQUAL = operator.le
+
+ @staticmethod
+ def from_string(version_string: str) -> "VersionComparison":
+ string_to_operator = {
+ "=": VersionComparison.EQUAL,
+ "==": VersionComparison.EQUAL,
+ "!=": VersionComparison.NOT_EQUAL,
+ ">": VersionComparison.GREATER_THAN,
+ "<": VersionComparison.LESS_THAN,
+ ">=": VersionComparison.GREATER_THAN_OR_EQUAL,
+ "<=": VersionComparison.LESS_THAN_OR_EQUAL,
+ }
+
+ return string_to_operator[version_string]
+
+
+@lru_cache
+def split_package_version(package_version_str) -> tuple[str, str, str]:
+ pattern = r"([a-zA-Z0-9_-]+)([!<>=~]+)([0-9.]+)"
+ match = re.match(pattern, package_version_str)
+ if match:
+ return (match.group(1), match.group(2), match.group(3))
+ else:
+ raise ValueError(f"Invalid package version string: {package_version_str}")
+
+
+class Backend:
+ def __init__(self, backend_requirement: str):
+ self.package_name, self.version_comparison, self.version = split_package_version(backend_requirement)
+
+ if self.package_name not in BACKENDS_MAPPING:
+ raise ValueError(
+ f"Backends should be defined in the BACKENDS_MAPPING. Offending backend: {self.package_name}"
+ )
+
+ def get_installed_version(self) -> str:
+ """Return the currently installed version of the backend"""
+ is_available, current_version = _is_package_available(self.package_name, return_version=True)
+ if not is_available:
+ raise RuntimeError(f"Backend {self.package_name} is not available.")
+ return current_version
+
+ def is_satisfied(self) -> bool:
+ return VersionComparison.from_string(self.version_comparison).value(
+ version.parse(self.get_installed_version()), version.parse(self.version)
+ )
+
+ def __repr__(self) -> str:
+ return f'Backend("{self.package_name}", {VersionComparison[self.version_comparison]}, "{self.version}")'
+
+ @property
+ def error_message(self):
+ return (
+ f"{{0}} requires the {self.package_name} library version {self.version_comparison}{self.version}. That"
+ f" library was not found with this version in your environment."
+ )
+
+
+def requires(*, backends=()):
+ """
+ This decorator enables two things:
+ - Attaching a `__backends` tuple to an object to see what are the necessary backends for it
+ to execute correctly without instantiating it
+ - The '@requires' string is used to dynamically import objects
+ """
+
+ if not isinstance(backends, (tuple, list)):
+ raise TypeError("Backends should be a tuple or list.")
+ backends = tuple(backends)
+
+ applied_backends = []
+ for backend in backends:
+ if backend in BACKENDS_MAPPING:
+ applied_backends.append(backend)
+ else:
+ if any(key in backend for key in ["=", "<", ">"]):
+ applied_backends.append(Backend(backend))
+ else:
+ raise ValueError(f"Backend should be defined in the BACKENDS_MAPPING. Offending backend: {backend}")
+
+ def inner_fn(fun):
+ if isinstance(fun, type):
+ # For classes, just attach the metadata โ don't wrap, as that would
+ # turn the class into a plain function and break isinstance checks.
+ fun.__backends = applied_backends
+ return fun
+
+ @functools.wraps(fun)
+ def wrapper(*args, **kwargs):
+ requires_backends(fun, applied_backends)
+ return fun(*args, **kwargs)
+
+ wrapper.__backends = applied_backends # type: ignore [unresolved-attribute]
+ return wrapper
+
+ return inner_fn
+
+
+BASE_FILE_REQUIREMENTS = {
+ lambda name, content: "modeling_" in name: ("torch",),
+ lambda name, content: "tokenization_" in name and name.endswith("_fast"): ("tokenizers",),
+ lambda name, content: "image_processing_" in name and "TorchvisionBackend" in content: (
+ "vision",
+ "torch",
+ "torchvision",
+ ),
+ lambda name, content: "image_processing_" in name: ("vision",),
+ lambda name, content: "video_processing_" in name: ("vision", "torch", "torchvision"),
+}
+
+
+def fetch__all__(file_content) -> list[str]:
+ """
+ Returns the content of the __all__ variable in the file content.
+ Returns None if not defined, otherwise returns a list of strings.
+ """
+
+ if "__all__" not in file_content:
+ return []
+
+ start_index = None
+ lines = file_content.splitlines()
+ for index, line in enumerate(lines):
+ if line.startswith("__all__"):
+ start_index = index
+
+ # There is no line starting with `__all__`
+ if start_index is None:
+ return []
+
+ lines = lines[start_index:]
+
+ if not lines[0].startswith("__all__"):
+ raise ValueError(
+ "fetch__all__ accepts a list of lines, with the first line being the __all__ variable declaration"
+ )
+
+ # __all__ is defined on a single line
+ if lines[0].endswith("]"):
+ return [obj.strip("\"' ") for obj in lines[0].split("=")[1].strip(" []").split(",")]
+
+ # __all__ is defined on multiple lines
+ else:
+ _all: list[str] = []
+ for __all__line_index in range(1, len(lines)):
+ if lines[__all__line_index].strip() == "]":
+ return _all
+ else:
+ _all.append(lines[__all__line_index].strip("\"', "))
+
+ return _all
+
+
+@lru_cache
+def create_import_structure_from_path(module_path):
+ """
+ This method takes the path to a file/a folder and returns the import structure.
+ If a file is given, it will return the import structure of the parent folder.
+
+ Import structures are designed to be digestible by `_LazyModule` objects. They are
+ created from the __all__ definitions in each files as well as the `@require` decorators
+ above methods and objects.
+
+ The import structure allows explicit display of the required backends for a given object.
+ These backends are specified in two ways:
+
+ 1. Through their `@require`, if they are exported with that decorator. This `@require` decorator
+ accepts a `backend` tuple kwarg mentioning which backends are required to run this object.
+
+ 2. If an object is defined in a file with "default" backends, it will have, at a minimum, this
+ backend specified. The default backends are defined according to the filename:
+
+ - If a file is named like `modeling_*.py`, it will have a `torch` backend
+ - If a file is named like `tokenization_*_fast.py`, it will have a `tokenizers` backend
+ - If a file is named like `image_processing*_fast.py`, it will have a `torchvision` + `torch` backend
+
+ Backends serve the purpose of displaying a clear error message to the user in case the backends are not installed.
+ Should an object be imported without its required backends being in the environment, any attempt to use the
+ object will raise an error mentioning which backend(s) should be added to the environment in order to use
+ that object.
+
+ Here's an example of an input import structure at the src.transformers.models level:
+
+ {
+ 'albert': {
+ frozenset(): {
+ 'configuration_albert': {'AlbertConfig'}
+ },
+ frozenset({'tokenizers'}): {
+ 'tokenization_albert_fast': {'AlbertTokenizer'}
+ },
+ },
+ 'align': {
+ frozenset(): {
+ 'configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'},
+ 'processing_align': {'AlignProcessor'}
+ },
+ },
+ 'altclip': {
+ frozenset(): {
+ 'configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'},
+ 'processing_altclip': {'AltCLIPProcessor'},
+ }
+ }
+ }
+ """
+ import_structure = {}
+
+ if os.path.isfile(module_path):
+ module_path = os.path.dirname(module_path)
+
+ adjacent_modules = []
+
+ with os.scandir(module_path) as entries:
+ for entry in entries:
+ if entry.name == "__pycache__":
+ continue
+ if entry.is_dir():
+ import_structure[entry.name] = create_import_structure_from_path(entry.path)
+ elif not entry.name.startswith(("convert_", "modular_")):
+ adjacent_modules.append(entry.name)
+
+ # We're only taking a look at files different from __init__.py
+ # We could theoretically require things directly from the __init__.py
+ # files, but this is not supported at this time.
+ if "__init__.py" in adjacent_modules:
+ adjacent_modules.remove("__init__.py")
+
+ module_requirements = {}
+ for module_name in adjacent_modules:
+ # Only modules ending in `.py` are accepted here.
+ if not module_name.endswith(".py"):
+ continue
+
+ with open(os.path.join(module_path, module_name), encoding="utf-8") as f:
+ file_content = f.read()
+
+ # Remove the .py suffix
+ module_name = module_name[:-3]
+
+ previous_line = ""
+ previous_index = 0
+
+ # Some files have some requirements by default.
+ # For example, any file named `modeling_xxx.py`
+ # should have torch as a required backend.
+ base_requirements = ()
+ for check, requirements in BASE_FILE_REQUIREMENTS.items():
+ if check(module_name, file_content):
+ base_requirements = requirements
+ break
+
+ # Objects that have a `@require` assigned to them will get exported
+ # with the backends specified in the decorator as well as the file backends.
+ exported_objects = set()
+ if "@requires" in file_content:
+ lines = file_content.split("\n")
+ for index, line in enumerate(lines):
+ # This allows exporting items with other decorators. We'll take a look
+ # at the line that follows at the same indentation level.
+ if line.startswith((" ", "\t", "@", ")")) and not line.startswith("@requires"):
+ continue
+
+ # Skipping line enables putting whatever we want between the
+ # requires() call and the actual class/method definition.
+ # This is what enables having # Copied from statements, docs, etc.
+ skip_line = False
+
+ if "@requires" in previous_line:
+ skip_line = False
+
+ # Backends are defined on the same line as requires
+ if "backends" in previous_line:
+ try:
+ backends_string = previous_line.split("backends=")[1].split("(")[1].split(")")[0]
+ except IndexError:
+ raise ValueError(
+ f"Couldn't parse backends for @requires decorator in file {module_name}:{previous_line}"
+ )
+ backends = tuple(sorted([b.strip("'\",") for b in backends_string.split(", ") if b]))
+
+ # Backends are defined in the lines following requires, for example such as:
+ # @requires(
+ # backends=(
+ # "sentencepiece",
+ # "torch",
+ # )
+ # )
+ #
+ # or
+ #
+ # @requires(
+ # backends=(
+ # "sentencepiece",
+ # )
+ # )
+ elif "backends" in lines[previous_index + 1]:
+ backends = []
+ for backend_line in lines[previous_index:index]:
+ if "backends" in backend_line:
+ backend_line = backend_line.split("=")[1]
+ if '"' in backend_line or "'" in backend_line:
+ if ", " in backend_line:
+ backends.extend(backend.strip("()\"', ") for backend in backend_line.split(", "))
+ else:
+ backends.append(backend_line.strip("()\"', "))
+
+ # If the line is only a ')', then we reached the end of the backends and we break.
+ if backend_line.strip() == ")":
+ break
+ backends = tuple(backends)
+
+ # No backends are registered for requires
+ else:
+ backends = ()
+
+ backends = frozenset(backends + base_requirements)
+ if backends not in module_requirements:
+ module_requirements[backends] = {}
+ if module_name not in module_requirements[backends]:
+ module_requirements[backends][module_name] = set()
+
+ if not line.startswith("class") and not line.startswith("def"):
+ skip_line = True
+ else:
+ start_index = 6 if line.startswith("class") else 4
+ object_name = line[start_index:].split("(")[0].strip(":")
+ module_requirements[backends][module_name].add(object_name)
+ exported_objects.add(object_name)
+
+ if not skip_line:
+ previous_line = line
+ previous_index = index
+
+ # All objects that are in __all__ should be exported by default.
+ # These objects are exported with the file backends.
+ if "__all__" in file_content:
+ for _all_object in fetch__all__(file_content):
+ if _all_object not in exported_objects:
+ backends = frozenset(base_requirements)
+ if backends not in module_requirements:
+ module_requirements[backends] = {}
+ if module_name not in module_requirements[backends]:
+ module_requirements[backends][module_name] = set()
+
+ module_requirements[backends][module_name].add(_all_object)
+
+ import_structure = {**module_requirements, **import_structure}
+ return import_structure
+
+
+def spread_import_structure(nested_import_structure):
+ """
+ This method takes as input an unordered import structure and brings the required backends at the top-level,
+ aggregating modules and objects under their required backends.
+
+ Here's an example of an input import structure at the src.transformers.models level:
+
+ {
+ 'albert': {
+ frozenset(): {
+ 'configuration_albert': {'AlbertConfig'}
+ },
+ frozenset({'tokenizers'}): {
+ 'tokenization_albert_fast': {'AlbertTokenizer'}
+ },
+ },
+ 'align': {
+ frozenset(): {
+ 'configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'},
+ 'processing_align': {'AlignProcessor'}
+ },
+ },
+ 'altclip': {
+ frozenset(): {
+ 'configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'},
+ 'processing_altclip': {'AltCLIPProcessor'},
+ }
+ }
+ }
+
+ Here's an example of an output import structure at the src.transformers.models level:
+
+ {
+ frozenset({'tokenizers'}): {
+ 'albert.tokenization_albert_fast': {'AlbertTokenizer'}
+ },
+ frozenset(): {
+ 'albert.configuration_albert': {'AlbertConfig'},
+ 'align.processing_align': {'AlignProcessor'},
+ 'align.configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'},
+ 'altclip.configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'},
+ 'altclip.processing_altclip': {'AltCLIPProcessor'}
+ }
+ }
+
+ """
+
+ def propagate_frozenset(unordered_import_structure):
+ frozenset_first_import_structure = {}
+ for _key, _value in unordered_import_structure.items():
+ # If the value is not a dict but a string, no need for custom manipulation
+ if not isinstance(_value, dict):
+ frozenset_first_import_structure[_key] = _value
+
+ elif any(isinstance(v, frozenset) for v in _value):
+ for k, v in _value.items():
+ if isinstance(k, frozenset):
+ # Here we want to switch around _key and k to propagate k upstream if it is a frozenset
+ if k not in frozenset_first_import_structure:
+ frozenset_first_import_structure[k] = {}
+ if _key not in frozenset_first_import_structure[k]:
+ frozenset_first_import_structure[k][_key] = {}
+
+ frozenset_first_import_structure[k][_key].update(v)
+
+ else:
+ # If k is not a frozenset, it means that the dictionary is not "level": some keys (top-level)
+ # are frozensets, whereas some are not -> frozenset keys are at an unknown depth-level of the
+ # dictionary.
+ #
+ # We recursively propagate the frozenset for this specific dictionary so that the frozensets
+ # are at the top-level when we handle them.
+ propagated_frozenset = propagate_frozenset({k: v})
+ for r_k, r_v in propagated_frozenset.items():
+ if isinstance(_key, frozenset):
+ if r_k not in frozenset_first_import_structure:
+ frozenset_first_import_structure[r_k] = {}
+ if _key not in frozenset_first_import_structure[r_k]:
+ frozenset_first_import_structure[r_k][_key] = {}
+
+ # _key is a frozenset -> we switch around the r_k and _key
+ frozenset_first_import_structure[r_k][_key].update(r_v)
+ else:
+ if _key not in frozenset_first_import_structure:
+ frozenset_first_import_structure[_key] = {}
+ if r_k not in frozenset_first_import_structure[_key]:
+ frozenset_first_import_structure[_key][r_k] = {}
+
+ # _key is not a frozenset -> we keep the order of r_k and _key
+ frozenset_first_import_structure[_key][r_k].update(r_v)
+
+ else:
+ frozenset_first_import_structure[_key] = propagate_frozenset(_value)
+
+ return frozenset_first_import_structure
+
+ def flatten_dict(_dict, previous_key=None):
+ items = []
+ for _key, _value in _dict.items():
+ _key = f"{previous_key}.{_key}" if previous_key is not None else _key
+ if isinstance(_value, dict):
+ items.extend(flatten_dict(_value, _key).items())
+ else:
+ items.append((_key, _value))
+ return dict(items)
+
+ # The tuples contain the necessary backends. We want these first, so we propagate them up the
+ # import structure.
+ ordered_import_structure = nested_import_structure
+
+ # 6 is a number that gives us sufficient depth to go through all files and foreseeable folder depths
+ # while not taking too long to parse.
+ for i in range(6):
+ ordered_import_structure = propagate_frozenset(ordered_import_structure)
+
+ # We then flatten the dict so that it references a module path.
+ flattened_import_structure = {}
+ for key, value in ordered_import_structure.copy().items():
+ if isinstance(key, str):
+ del ordered_import_structure[key]
+ else:
+ flattened_import_structure[key] = flatten_dict(value)
+
+ return flattened_import_structure
+
+
+@lru_cache
+def define_import_structure(module_path: str, prefix: str | None = None) -> IMPORT_STRUCTURE_T:
+ """
+ This method takes a module_path as input and creates an import structure digestible by a _LazyModule.
+
+ Here's an example of an output import structure at the src.transformers.models level:
+
+ {
+ frozenset({'tokenizers'}): {
+ 'albert.tokenization_albert_fast': {'AlbertTokenizer'}
+ },
+ frozenset(): {
+ 'albert.configuration_albert': {'AlbertConfig'},
+ 'align.processing_align': {'AlignProcessor'},
+ 'align.configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'},
+ 'altclip.configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'},
+ 'altclip.processing_altclip': {'AltCLIPProcessor'}
+ }
+ }
+
+ The import structure is a dict defined with frozensets as keys, and dicts of strings to sets of objects.
+
+ If `prefix` is not None, it will add that prefix to all keys in the returned dict.
+ """
+ import_structure = create_import_structure_from_path(module_path)
+ spread_dict = spread_import_structure(import_structure)
+
+ if prefix is None:
+ return spread_dict
+ else:
+ spread_dict = {k: {f"{prefix}.{kk}": vv for kk, vv in v.items()} for k, v in spread_dict.items()}
+ return spread_dict
+
+
+def clear_import_cache() -> None:
+ """
+ Clear cached Transformers modules to allow reloading modified code.
+
+ This is useful when actively developing/modifying Transformers code.
+ """
+ # Get all transformers modules
+ transformers_modules = [mod_name for mod_name in sys.modules if mod_name.startswith("transformers.")]
+
+ # Remove them from sys.modules
+ for mod_name in transformers_modules:
+ module = sys.modules[mod_name]
+ # Clear _LazyModule caches if applicable
+ if isinstance(module, _LazyModule):
+ module._objects = {} # Clear cached objects
+ del sys.modules[mod_name]
+
+ # Force reload main transformers module
+ if "transformers" in sys.modules:
+ main_module = sys.modules["transformers"]
+ if isinstance(main_module, _LazyModule):
+ main_module._objects = {} # Clear cached objects
+ importlib.reload(main_module)
diff --git a/third_party/transformers/src/transformers/utils/kernel_config.py b/third_party/transformers/src/transformers/utils/kernel_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb4f965ddbf4b22cb61dc465832af338478eb2f5
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/kernel_config.py
@@ -0,0 +1,281 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+from ..utils import PushToHubMixin
+
+
+def infer_device(model):
+ """
+ Infers the device type from the model parameters.
+ Args:
+ model: The model instance.
+
+ Returns:
+ The device type.
+ """
+ EXAMPLE_MAPPING = """
+ {
+ "RMSNorm": {
+ "cuda":
+ "kernels-community/layer_norm:LlamaRMSNorm",
+ ...
+ },
+ ...
+ }
+ """
+ try:
+ param = next(model.parameters())
+ except StopIteration:
+ raise ValueError(
+ f"Cannot determine model device, please provide a device to the mapping. Example: {EXAMPLE_MAPPING}"
+ )
+
+ dev_type = param.device.type
+ if dev_type == "cuda":
+ # Refine based on actual platform
+ from ..utils import is_torch_available
+
+ if is_torch_available():
+ import torch
+
+ if getattr(torch, "version").hip is not None:
+ return "rocm"
+
+ return dev_type
+
+
+def add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping):
+ from kernels import LayerRepository
+
+ if device not in ["cuda", "rocm", "xpu", "npu", "neuron"]:
+ raise ValueError(f"Only cuda, rocm, xpu, npu and neuron devices supported, got: {device}")
+ repo_layer_name = repo_name.split(":")[1]
+ repo_id = repo_name.split(":")[0]
+ compatible_mapping[layer_name] = {
+ device: {
+ mode: LayerRepository(
+ repo_id=repo_id,
+ layer_name=repo_layer_name,
+ )
+ }
+ }
+
+
+def add_to_mapping_local(layer_name, device, repo_name, mode, compatible_mapping):
+ from pathlib import Path
+
+ from kernels import LocalLayerRepository
+
+ if device not in ["cuda", "rocm", "xpu", "npu", "neuron"]:
+ raise ValueError(f"Only cuda, rocm, xpu, npu and neuron devices supported, got: {device}")
+ repo_layer_name = repo_name.split(":")[1]
+ repo_path = repo_name.split(":")[0]
+ repo_package_name = repo_path.split("/")[-1]
+ compatible_mapping[layer_name] = {
+ device: {
+ mode: LocalLayerRepository(
+ repo_path=Path(repo_path),
+ package_name=repo_package_name,
+ layer_name=repo_layer_name,
+ )
+ }
+ }
+
+
+class KernelConfig(PushToHubMixin):
+ """
+ Kernel configuration class. This class is used to configure the kernel mapping for a model.
+ """
+
+ def __init__(self, kernel_mapping=None, use_local_kernel=False):
+ self.kernel_mapping = kernel_mapping if kernel_mapping is not None else {}
+ self.registered_layer_names = {}
+ self.use_local_kernel = use_local_kernel
+
+ def update_kernel(self, repo_id, registered_name, layer_name, device, mode, revision=None):
+ from kernels import LayerRepository
+
+ self.kernel_mapping[registered_name] = {
+ device: {
+ mode: LayerRepository(
+ repo_id=repo_id,
+ layer_name=layer_name,
+ revision=revision,
+ )
+ }
+ }
+
+ def store_registered_layer_names(self, model):
+ for name, module in model.named_modules():
+ if hasattr(module, "kernel_layer_name"):
+ self.registered_layer_names[name] = module.kernel_layer_name
+
+ def sanitize_kernel_mapping(self, model):
+ """
+ Validates the kernel_mapping to ensure that:
+ 1. Each layer_name in the mapping is registered in the model (i.e., the model contains a module with a matching kernel_layer_name).
+ 2. Each kernel value is either a string of the form 'org/repo:layer_name' or a dict mapping device types ("cuda", "rocm", "xpu", "npu") to such strings.
+ 3. Each device key in a dict is one of "cuda", "rocm", "xpu", or "npu".
+ 4. Each repo_name is a valid repository and layer name in the format 'org/repo:layer_name' (i.e., a string containing both a slash and a colon).
+ 5. If a local path is detected, it should be in the format '/abs/path:layer_name'. The absolute path must include the `package_name`, like "/home/user/layer_norm".
+
+ Args:
+ model: The model instance whose modules are checked for registered kernel_layer_name attributes.
+
+ Raises:
+ ValueError: If a layer_name is not registered in the model, if a device is not supported,
+ or if a repo_name is not a valid 'org/repo:layer_name' string.
+ """
+ MAPPING_FORMAT = """
+ For single device form remote
+ {
+ "RMSNorm":
+ "kernels-community/layer_norm:LlamaRMSNorm",
+ ...
+ },
+ For multiple devices form remote
+ {
+ "RMSNorm": {
+ "cuda":
+ "kernels-community/layer_norm:LlamaRMSNorm",
+ "rocm":
+ "kernels-community/layer_norm:LlamaRMSNorm",
+ ...
+ },
+ ...
+ }
+ For single device form local
+ {
+ "RMSNorm":
+ "/abs/path:LlamaRMSNorm",
+ ...
+ },
+ For multiple devices form local
+ {
+ "RMSNorm": {
+ "cuda":
+ "/abs/path:LlamaRMSNorm",
+ "rocm":
+ "/abs/path:LlamaRMSNorm",
+ ...
+ },
+ ...
+ }
+ """
+ self.store_registered_layer_names(model)
+ # Validate that the kernel mapping is a dict
+ if not isinstance(self.kernel_mapping, dict):
+ raise ValueError(
+ f"Kernel mapping must be a dict of the following format: {MAPPING_FORMAT}, got: {type(self.kernel_mapping)}"
+ )
+
+ for layer_name, kernel in self.kernel_mapping.items():
+ if layer_name not in self.registered_layer_names.values():
+ raise ValueError(
+ f"Layer {layer_name} is not registered in the model, please register it first using use_kernel_forward_from_hub"
+ )
+
+ if isinstance(kernel, str):
+ if "/" not in kernel or ":" not in kernel:
+ raise ValueError(
+ f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name' or '/abs/path:layer_name'), got: {kernel}"
+ )
+
+ elif isinstance(kernel, dict):
+ for device, repo_name in kernel.items():
+ if device not in ["cuda", "rocm", "xpu", "npu", "neuron"]:
+ raise ValueError(f"Only cuda, rocm, xpu, npu and neuron devices supported, got: {device}")
+
+ if not isinstance(repo_name, str) or "/" not in repo_name or ":" not in repo_name:
+ raise ValueError(
+ f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name' or '/abs/path:layer_name'), got: {repo_name}"
+ )
+ else:
+ raise ValueError(f"Kernel mapping must follow the format: {MAPPING_FORMAT}, got: {kernel}")
+
+ def create_compatible_mapping(self, model, compile=False):
+ """
+ Transforms a simple kernel_mapping of the form:
+ {
+ "RMSNorm":
+ "kernels-community/layer_norm:LlamaRMSNorm",
+ ...
+ },
+
+ or for local path:
+
+ {
+ "RMSNorm":
+ "/home/user/liger_kernels:LigerRMSNorm",
+ ...
+ },
+
+ into a nested mapping:
+
+ {
+ "RMSNorm": {
+ "cuda": {
+ Mode.INFERENCE: LayerRepository(
+ repo_id="kernels-community/layer_norm",
+ layer_name="LlamaRMSNorm",
+ )
+ }
+ }
+ }
+
+ or for local path:
+
+ {
+ "RMSNorm": {
+ "cuda": {
+ Mode.INFERENCE: LocalLayerRepository(
+ repo_path=Path("/home/user/liger_kernels"),
+ package_name="liger_kernels",
+ layer_name="LigerRMSNorm",
+ )
+ }
+ }
+ }
+
+ that's compatible with the kernels library.
+
+ The device is inferred from the model's parameters if not provided.
+ The Mode is inferred from the model's training state.
+ """
+ from kernels import Mode
+
+ compatible_mapping = {}
+ current_device = infer_device(model)
+ for layer_name, kernel in self.kernel_mapping.items():
+ # Infer Mode: use Mode.TRAINING if model is training, else use Mode.INFERENCE
+ mode = Mode.TRAINING if model.training else Mode.INFERENCE
+ if compile:
+ mode = mode | Mode.TORCH_COMPILE
+
+ if isinstance(kernel, str):
+ repo_name = kernel
+ if not self.use_local_kernel:
+ add_to_mapping(layer_name, current_device, repo_name, mode, compatible_mapping)
+ else:
+ add_to_mapping_local(layer_name, current_device, repo_name, mode, compatible_mapping)
+ elif isinstance(kernel, dict):
+ for device, repo_name in kernel.items():
+ if device != current_device:
+ continue
+ if not self.use_local_kernel:
+ add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping)
+ else:
+ add_to_mapping_local(layer_name, device, repo_name, mode, compatible_mapping)
+
+ self.kernel_mapping = compatible_mapping
diff --git a/third_party/transformers/src/transformers/utils/loading_report.py b/third_party/transformers/src/transformers/utils/loading_report.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e6ffcb77da4211bad1e8074cf28ed46181fc402
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/loading_report.py
@@ -0,0 +1,280 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+import logging
+import re
+import shutil
+import sys
+from collections import OrderedDict, defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+
+_DIGIT_RX = re.compile(r"(?<=\.)(\d+)(?=\.|$)") # numbers between dots or at the end
+
+
+def _pattern_of(key: str) -> str:
+ """Replace every dot-delimited integer with '*' to get the structure."""
+ return _DIGIT_RX.sub("*", key)
+
+
+def _fmt_indices(values: list[int], cutoff=10) -> str:
+ """Format a list of ints as single number, {a, ..., b}, or first...last."""
+ if len(values) == 1:
+ return str(values[0])
+ values = sorted(values)
+ if len(values) > cutoff:
+ return f"{values[0]}...{values[-1]}"
+ return ", ".join(map(str, values))
+
+
+def update_key_name(mapping: dict[str, Any]) -> dict[str, Any]:
+ """
+ Merge keys like 'layers.0.x', 'layers.1.x' into 'layers.{0, 1}.x'
+ BUT only merge together keys that have the exact same value.
+ Returns a new dict {merged_key: value}.
+ """
+ # (pattern, value) -> list[set[int]] (per-star index values)
+ not_mapping = False
+ if not isinstance(mapping, dict):
+ mapping = {k: k for k in mapping}
+ not_mapping = True
+
+ bucket: dict[str, list[set[int] | Any]] = defaultdict(list)
+ for key, val in mapping.items():
+ digs = _DIGIT_RX.findall(key)
+ patt = _pattern_of(key)
+ for i, d in enumerate(digs):
+ if len(bucket[patt]) <= i:
+ bucket[patt].append(set())
+ bucket[patt][i].add(int(d))
+ bucket[patt].append(val)
+
+ out_items = {}
+ for patt, values in bucket.items():
+ sets, val = values[:-1], values[-1]
+ parts = patt.split("*") # stars are between parts
+ final = parts[0]
+ for i in range(1, len(parts)):
+ if i - 1 < len(sets) and sets[i - 1]:
+ insert = _fmt_indices(sorted(sets[i - 1]))
+ if len(sets[i - 1]) > 1:
+ final += "{" + insert + "}"
+ else:
+ final += insert
+ else:
+ final += "*"
+ final += parts[i]
+
+ out_items[final] = val
+ out = OrderedDict(out_items)
+ if not_mapping:
+ return out.keys()
+ return out
+
+
+_ansi_re = re.compile(r"\x1b\[[0-9;]*m")
+
+
+def _strip_ansi(s: str) -> str:
+ return _ansi_re.sub("", str(s))
+
+
+def _pad(text, width):
+ t = str(text)
+ pad = max(0, width - len(_strip_ansi(t)))
+ return t + " " * pad
+
+
+def _make_table(rows, headers):
+ # compute display widths while ignoring ANSI codes
+ cols = list(zip(*([headers] + rows))) if rows else [headers]
+ widths = [max(len(_strip_ansi(x)) for x in col) for col in cols]
+ header_line = " | ".join(_pad(h, w) for h, w in zip(headers, widths))
+ sep_line = "-+-".join("-" * w for w in widths)
+ body = [" | ".join(_pad(c, w) for c, w in zip(r, widths)) for r in rows]
+ return "\n".join([header_line, sep_line] + body)
+
+
+PALETTE = {
+ "reset": "[0m",
+ "red": "[31m",
+ "yellow": "[33m",
+ "orange": "[38;5;208m",
+ "purple": "[35m",
+ "bold": "[1m",
+ "italic": "[3m",
+ "dim": "[2m",
+}
+
+
+def _style(s, color):
+ """Return color/style-formatted input `s` if `sys.stdout` is interactive, e.g. connected to a terminal."""
+ if sys.stdout.isatty():
+ return f"{PALETTE[color]}{s}{PALETTE['reset']}"
+ else:
+ return s
+
+
+def _get_terminal_width(default=80):
+ try:
+ return shutil.get_terminal_size().columns
+ except Exception:
+ return default
+
+
+@dataclass
+class LoadStateDictInfo:
+ """
+ Mutable container for state-dict loading results and diagnostics. Each entry in this structure is mutable,
+ and will usually be mutated in-place during the loading pipeline.
+
+ Attributes:
+ missing_keys (`set[str]`):
+ Keys that are missing from the loaded checkpoints but expected in the model's architecture.
+ unexpected_keys (`set[str]`):
+ Keys that are found in the checkpoints, but not expected in the model's architecture.
+ mismatched_keys (`set[tuple[str, tuple[int], tuple[int]]]`):
+ Keys that are found in the checkpoints and are expected in the model's architecture, but with a different shape.
+ error_msgs ( `list[str]`):
+ Some potential error messages.
+ conversion_errors (`dict[str, str]`):
+ Errors happening during the on-the-fly weight conversion process.
+ """
+
+ missing_keys: set[str]
+ unexpected_keys: set[str]
+ mismatched_keys: set[tuple[str, tuple[int], tuple[int]]]
+ error_msgs: list[str]
+ conversion_errors: dict[str, str]
+
+ def missing_and_mismatched(self):
+ """Return all effective missing keys, including `missing` and `mismatched` keys."""
+ return self.missing_keys | {k[0] for k in self.mismatched_keys}
+
+ def to_dict(self):
+ # Does not include the `conversion_errors` to be coherent with legacy reporting in the tests
+ return {
+ "missing_keys": self.missing_keys,
+ "unexpected_keys": self.unexpected_keys,
+ "mismatched_keys": self.mismatched_keys,
+ "error_msgs": self.error_msgs,
+ }
+
+ def create_loading_report(self) -> str | None:
+ """Generate the minimal table of a loading report."""
+ term_w = _get_terminal_width()
+
+ rows = []
+ tips = "\n\nNotes:"
+ if self.unexpected_keys:
+ tips += f"\n- {_style('UNEXPECTED:', 'orange')}\t" + _style(
+ "can be ignored when loading from different task/architecture; not ok if you expect identical arch.",
+ "italic",
+ )
+ for k in update_key_name(self.unexpected_keys):
+ status = _style("UNEXPECTED", "orange")
+ rows.append([k, status, "", ""])
+
+ if self.missing_keys:
+ tips += f"\n- {_style('MISSING:', 'red')}\t" + _style(
+ "those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.",
+ "italic",
+ )
+ for k in update_key_name(self.missing_keys):
+ status = _style("MISSING", "red")
+ rows.append([k, status, ""])
+
+ if self.mismatched_keys:
+ tips += f"\n- {_style('MISMATCH:', 'yellow')}\t" + _style(
+ "ckpt weights were loaded, but they did not match the original empty weight shapes.", "italic"
+ )
+ iterator = {a: (b, c) for a, b, c in self.mismatched_keys}
+ for key, (shape_ckpt, shape_model) in update_key_name(iterator).items():
+ status = _style("MISMATCH", "yellow")
+ data = [
+ key,
+ status,
+ f"Reinit due to size mismatch - ckpt: {str(shape_ckpt)} vs model:{str(shape_model)}",
+ ]
+ rows.append(data)
+
+ if self.conversion_errors:
+ tips += f"\n- {_style('CONVERSION:', 'purple')}\t" + _style(
+ "originate from the conversion scheme", "italic"
+ )
+ for k, v in update_key_name(self.conversion_errors).items():
+ status = _style("CONVERSION", "purple")
+ _details = f"\n\n{v}\n\n"
+ rows.append([k, status, _details])
+
+ # If nothing is wrong, return None
+ if len(rows) == 0:
+ return None
+
+ headers = ["Key", "Status"]
+ if term_w > 200:
+ headers += ["Details"]
+ else:
+ headers += ["", ""]
+ table = _make_table(rows, headers=headers)
+ report = table + tips
+
+ return report
+
+
+def log_state_dict_report(
+ model,
+ pretrained_model_name_or_path: str,
+ ignore_mismatched_sizes: bool,
+ loading_info: LoadStateDictInfo,
+ logger: logging.Logger | None = None,
+):
+ """
+ Log a readable report about state_dict loading issues.
+
+ This version is terminal-size aware: for very small terminals it falls back to a compact
+ Key | Status view so output doesn't wrap badly.
+ """
+ if logger is None:
+ logger = logging.getLogger(__name__)
+
+ # Re-raise errors early if needed
+ if loading_info.error_msgs:
+ error_msg = "\n\t".join(loading_info.error_msgs)
+ if "size mismatch" in error_msg:
+ error_msg += (
+ "\n\tYou may consider adding `ignore_mismatched_sizes=True` to `from_pretrained(...)` if appropriate."
+ )
+ raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
+
+ # Create the report table
+ report = loading_info.create_loading_report()
+ if report is None:
+ return
+
+ prelude = f"{PALETTE['bold']}{model.__class__.__name__} LOAD REPORT{PALETTE['reset']} from: {pretrained_model_name_or_path}\n"
+
+ # Log the report as warning
+ logger.warning(prelude + report)
+
+ # Re-raise in those case, after the report
+ if loading_info.conversion_errors:
+ raise RuntimeError(
+ "We encountered some issues during automatic conversion of the weights. For details look at the `CONVERSION` entries of "
+ "the above report!"
+ )
+ if not ignore_mismatched_sizes and loading_info.mismatched_keys:
+ raise RuntimeError(
+ "You set `ignore_mismatched_sizes` to `False`, thus raising an error. For details look at the above report!"
+ )
diff --git a/third_party/transformers/src/transformers/utils/logging.py b/third_party/transformers/src/transformers/utils/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc0e36aa8769ab500169a62acfda9d4bf8527397
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/logging.py
@@ -0,0 +1,434 @@
+# Copyright 2020 Optuna, Hugging Face
+#
+# 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.
+"""Logging utilities."""
+
+import functools
+import logging
+import os
+import sys
+import threading
+from collections.abc import Callable
+from logging import (
+ CRITICAL, # NOQA
+ DEBUG,
+ ERROR,
+ FATAL, # NOQA
+ INFO,
+ NOTSET, # NOQA
+ WARN, # NOQA
+ WARNING,
+)
+from logging import captureWarnings as _captureWarnings
+from typing import Any
+
+import huggingface_hub.utils as hf_hub_utils
+from tqdm import auto as tqdm_lib
+
+from .._typing import TransformersLogger
+
+
+_lock = threading.Lock()
+_default_handler: logging.Handler | None = None
+
+log_levels = {
+ "detail": logging.DEBUG, # will also print filename and line number
+ "debug": logging.DEBUG,
+ "info": logging.INFO,
+ "warning": logging.WARNING,
+ "error": logging.ERROR,
+ "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.WARNING
+
+_tqdm_active = not hf_hub_utils.are_progress_bars_disabled()
+_tqdm_hook: Callable[[Callable[..., Any], tuple[Any, ...], dict[str, Any]], Any] | None = None
+
+
+def _get_default_logging_level():
+ """
+ If TRANSFORMERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+ not - fall back to `_default_log_level`
+ """
+ env_level_str = os.getenv("TRANSFORMERS_VERBOSITY", None)
+ if env_level_str:
+ if env_level_str in log_levels:
+ return log_levels[env_level_str]
+ else:
+ logging.getLogger().warning(
+ f"Unknown option TRANSFORMERS_VERBOSITY={env_level_str}, "
+ f"has to be one of: {', '.join(log_levels.keys())}"
+ )
+ return _default_log_level
+
+
+def _get_library_name() -> str:
+ return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+ return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if _default_handler:
+ # This library has already configured the library root logger.
+ return
+ _default_handler = logging.StreamHandler() # Set sys.stderr as stream.
+ # set defaults based on https://github.com/pyinstaller/pyinstaller/issues/7334#issuecomment-1357447176
+ if sys.stderr is None:
+ sys.stderr = open(os.devnull, "w")
+
+ _default_handler.flush = sys.stderr.flush
+
+ # Apply our default configuration to the library root logger.
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.addHandler(_default_handler)
+ library_root_logger.setLevel(_get_default_logging_level())
+ # if logging level is debug, we add pathname and lineno to formatter for easy debugging
+ if os.getenv("TRANSFORMERS_VERBOSITY", None) == "detail":
+ formatter = logging.Formatter("[%(levelname)s|%(pathname)s:%(lineno)s] %(asctime)s >> %(message)s")
+ _default_handler.setFormatter(formatter)
+
+ ci = os.getenv("CI")
+ is_ci = ci is not None and ci.upper() in {"1", "ON", "YES", "TRUE"}
+ library_root_logger.propagate = is_ci
+
+
+def _reset_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if not _default_handler:
+ return
+
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.removeHandler(_default_handler)
+ library_root_logger.setLevel(logging.NOTSET)
+ _default_handler = None
+
+
+def get_log_levels_dict():
+ return log_levels
+
+
+def captureWarnings(capture):
+ """
+ Calls the `captureWarnings` method from the logging library to enable management of the warnings emitted by the
+ `warnings` library.
+
+ Read more about this method here:
+ https://docs.python.org/3/library/logging.html#integration-with-the-warnings-module
+
+ All warnings will be logged through the `py.warnings` logger.
+
+ Careful: this method also adds a handler to this logger if it does not already have one, and updates the logging
+ level of that logger to the library's root logger.
+ """
+ logger = get_logger("py.warnings")
+
+ if not logger.handlers:
+ logger.addHandler(_default_handler)
+
+ logger.setLevel(_get_library_root_logger().level)
+
+ _captureWarnings(capture)
+
+
+def get_logger(name: str | None = None) -> TransformersLogger:
+ """
+ Return a logger with the specified name.
+
+ This function is not supposed to be directly accessed unless you are writing a custom transformers module.
+ """
+
+ if name is None:
+ name = _get_library_name()
+
+ _configure_library_root_logger()
+ return logging.getLogger(name)
+
+
+def get_verbosity() -> int:
+ """
+ Return the current level for the ๐ค Transformers's root logger as an int.
+
+ Returns:
+ `int`: The logging level.
+
+
+
+ ๐ค Transformers has following logging levels:
+
+ - 50: `transformers.logging.CRITICAL` or `transformers.logging.FATAL`
+ - 40: `transformers.logging.ERROR`
+ - 30: `transformers.logging.WARNING` or `transformers.logging.WARN`
+ - 20: `transformers.logging.INFO`
+ - 10: `transformers.logging.DEBUG`
+
+ """
+
+ _configure_library_root_logger()
+ return _get_library_root_logger().getEffectiveLevel()
+
+
+def set_verbosity(verbosity: int) -> None:
+ """
+ Set the verbosity level for the ๐ค Transformers's root logger.
+
+ Args:
+ verbosity (`int`):
+ Logging level, e.g., one of:
+
+ - `transformers.logging.CRITICAL` or `transformers.logging.FATAL`
+ - `transformers.logging.ERROR`
+ - `transformers.logging.WARNING` or `transformers.logging.WARN`
+ - `transformers.logging.INFO`
+ - `transformers.logging.DEBUG`
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().setLevel(verbosity)
+
+
+def set_verbosity_info():
+ """Set the verbosity to the `INFO` level."""
+ return set_verbosity(INFO)
+
+
+def set_verbosity_warning():
+ """Set the verbosity to the `WARNING` level."""
+ return set_verbosity(WARNING)
+
+
+def set_verbosity_debug():
+ """Set the verbosity to the `DEBUG` level."""
+ return set_verbosity(DEBUG)
+
+
+def set_verbosity_error():
+ """Set the verbosity to the `ERROR` level."""
+ return set_verbosity(ERROR)
+
+
+def disable_default_handler() -> None:
+ """Disable the default handler of the HuggingFace Transformers's root logger."""
+
+ _configure_library_root_logger()
+
+ assert _default_handler is not None
+ _get_library_root_logger().removeHandler(_default_handler)
+
+
+def enable_default_handler() -> None:
+ """Enable the default handler of the HuggingFace Transformers's root logger."""
+
+ _configure_library_root_logger()
+
+ assert _default_handler is not None
+ _get_library_root_logger().addHandler(_default_handler)
+
+
+def add_handler(handler: logging.Handler) -> None:
+ """adds a handler to the HuggingFace Transformers's root logger."""
+
+ _configure_library_root_logger()
+
+ assert handler is not None
+ _get_library_root_logger().addHandler(handler)
+
+
+def remove_handler(handler: logging.Handler) -> None:
+ """removes given handler from the HuggingFace Transformers's root logger."""
+
+ _configure_library_root_logger()
+
+ assert handler is not None and handler not in _get_library_root_logger().handlers
+ _get_library_root_logger().removeHandler(handler)
+
+
+def disable_propagation() -> None:
+ """
+ Disable propagation of the library log outputs. Note that log propagation is disabled by default.
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().propagate = False
+
+
+def enable_propagation() -> None:
+ """
+ Enable propagation of the library log outputs. Please disable the HuggingFace Transformers's default handler to
+ prevent double logging if the root logger has been configured.
+ """
+
+ _configure_library_root_logger()
+ _get_library_root_logger().propagate = True
+
+
+def enable_explicit_format() -> None:
+ """
+ Enable explicit formatting for every HuggingFace Transformers's logger. The explicit formatter is as follows:
+ ```
+ [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
+ ```
+ All handlers currently bound to the root logger are affected by this method.
+ """
+ handlers = _get_library_root_logger().handlers
+
+ for handler in handlers:
+ formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s")
+ handler.setFormatter(formatter)
+
+
+def reset_format() -> None:
+ """
+ Resets the formatting for HuggingFace Transformers's loggers.
+
+ All handlers currently bound to the root logger are affected by this method.
+ """
+ handlers = _get_library_root_logger().handlers
+
+ for handler in handlers:
+ handler.setFormatter(None)
+
+
+def warning_advice(self, *args, **kwargs):
+ """
+ This method is identical to `logger.warning()`, but if env var TRANSFORMERS_NO_ADVISORY_WARNINGS=1 is set, this
+ warning will not be printed
+ """
+ no_advisory_warnings = os.getenv("TRANSFORMERS_NO_ADVISORY_WARNINGS")
+ if no_advisory_warnings:
+ return
+ self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_advice = warning_advice # type: ignore[unresolved-attribute]
+
+
+@functools.lru_cache(None)
+def warning_once(self, *args, **kwargs):
+ """
+ This method is identical to `logger.warning()`, but will emit the warning with the same message only once
+
+ Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache.
+ The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to
+ another type of cache that includes the caller frame information in the hashing function.
+ """
+ self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_once = warning_once # type: ignore[unresolved-attribute]
+
+
+@functools.lru_cache(None)
+def info_once(self, *args, **kwargs):
+ """
+ This method is identical to `logger.info()`, but will emit the info with the same message only once
+
+ Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache.
+ The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to
+ another type of cache that includes the caller frame information in the hashing function.
+ """
+ self.info(*args, **kwargs)
+
+
+logging.Logger.info_once = info_once # type: ignore[unresolved-attribute]
+
+
+class EmptyTqdm:
+ """Dummy tqdm which doesn't do anything."""
+
+ def __init__(self, *args, **kwargs): # pylint: disable=unused-argument
+ self._iterator = args[0] if args else None
+
+ def __iter__(self):
+ return iter(self._iterator)
+
+ def __getattr__(self, _):
+ """Return empty function."""
+
+ def empty_fn(*args, **kwargs): # pylint: disable=unused-argument
+ return
+
+ return empty_fn
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, type_, value, traceback):
+ return
+
+
+class _tqdm_cls:
+ def __call__(self, *args, **kwargs):
+ factory = tqdm_lib.tqdm if _tqdm_active else EmptyTqdm
+ if _tqdm_hook is not None:
+ return _tqdm_hook(factory, args, kwargs)
+ return factory(*args, **kwargs)
+
+ def set_lock(self, *args, **kwargs):
+ self._lock = None
+ if _tqdm_active:
+ return tqdm_lib.tqdm.set_lock(*args, **kwargs)
+
+ def get_lock(self):
+ if _tqdm_active:
+ return tqdm_lib.tqdm.get_lock()
+
+
+tqdm = _tqdm_cls()
+
+
+def is_progress_bar_enabled() -> bool:
+ """Return a boolean indicating whether tqdm progress bars are enabled."""
+ return bool(_tqdm_active)
+
+
+def enable_progress_bar():
+ """Enable tqdm progress bar."""
+ global _tqdm_active
+ _tqdm_active = True
+ hf_hub_utils.enable_progress_bars()
+
+
+def disable_progress_bar():
+ """Disable tqdm progress bar."""
+ global _tqdm_active
+ _tqdm_active = False
+ hf_hub_utils.disable_progress_bars()
+
+
+def set_tqdm_hook(hook: Callable[[Callable[..., Any], tuple[Any, ...], dict[str, Any]], Any] | None):
+ """
+ Set a hook that customizes tqdm creation.
+
+ The hook is called with the tqdm factory to use (either `tqdm.auto.tqdm` or an empty shim), along with the
+ positional and keyword arguments that would have been passed to tqdm. The hook should return an object compatible
+ with tqdm (i.e. implementing the methods your code relies on, such as `update`, `close`, context manager methods,
+ etc.).
+
+ Passing `None` clears the hook.
+
+ Returns:
+ The previous hook, which can be restored later.
+ """
+ global _tqdm_hook
+ previous_hook = _tqdm_hook
+ _tqdm_hook = hook
+ return previous_hook
diff --git a/third_party/transformers/src/transformers/utils/metrics.py b/third_party/transformers/src/transformers/utils/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..e62fa846ae022b1ee287ad02752ca0ee7910387f
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/metrics.py
@@ -0,0 +1,405 @@
+import functools
+import logging
+import time
+from collections.abc import Callable
+from enum import Enum
+from typing import Any
+
+from .import_utils import is_opentelemetry_available
+
+
+class RequestStatus(Enum):
+ """Status of a generation request through its lifecycle."""
+
+ PENDING = "pending"
+ PREFILLING = "prefilling"
+ PREFILLING_SPLIT = "prefilling_split"
+ SPLIT_PENDING_REMAINDER = "split_pending_remainder"
+ DECODING = "decoding"
+ FINISHED = "finished"
+ FAILED = "failed"
+
+
+if is_opentelemetry_available():
+ from opentelemetry import metrics
+ from opentelemetry.trace import Status, StatusCode, get_tracer
+
+ _has_opentelemetry = True
+else:
+ _has_opentelemetry = False
+
+
+def attach_tracer(tracer_name_template=None):
+ """
+ Decorator that attaches a tracer to a class.
+
+ This decorator should be applied to classes that need OpenTelemetry tracing.
+ It adds a tracer attribute to the class instance that can be used by the traced decorator.
+
+ Args:
+ tracer_name_template: Optional template string for the tracer name.
+ If provided, it should contain {module} which will be replaced with the class's full module path
+ and {class_name} for the class name.
+ If None, a default naming scheme will be used where:
+ - If the module already starts with "transformers.", it will use that directly
+ - Otherwise, it will prepend "transformers." to the module name
+
+ Returns:
+ Class decorator function
+ """
+ if not _has_opentelemetry:
+ return lambda cls: cls
+
+ def decorator(cls):
+ original_init = cls.__init__
+
+ @functools.wraps(original_init)
+ def init_with_tracer(self, *args, **kwargs):
+ original_init(self, *args, **kwargs)
+
+ module_name = cls.__module__
+ class_name = cls.__qualname__
+
+ if tracer_name_template is None:
+ if module_name.startswith("transformers."):
+ tracer_name = f"{module_name}.{class_name}"
+ else:
+ tracer_name = f"transformers.{module_name}.{class_name}"
+ else:
+ tracer_name = tracer_name_template.format(module=module_name, class_name=class_name)
+
+ self.tracer = get_tracer(tracer_name)
+
+ cls.__init__ = init_with_tracer
+ return cls
+
+ return decorator
+
+
+def traced(
+ func=None,
+ *,
+ span_name=None,
+ standalone=False,
+ additional_attributes: list[tuple[str, str, Any | Callable[[Any], Any]]] | None = None,
+):
+ """
+ Decorator to trace function calls with OpenTelemetry.
+
+ Can be used as @traced or @traced(span_name="custom_name")
+
+ Args:
+ func: The function to trace
+ span_name: Optional custom name for the span (defaults to function name)
+ standalone: If True, creates a parentless span
+ additional_attributes: Optional list of additional attributes to set on the span.
+ Each item is a tuple of (instance_attribute_name, span_attribute_key, value_or_transform_function)
+ where:
+ - instance_attribute_name: Name of the attribute to get from the class instance
+ - span_attribute_key: Key to use when setting the attribute on the span
+ - value_or_transform_function: Either a raw value to use directly, or a function to transform
+ the attribute value before setting it on the span
+
+ Returns:
+ Decorated function with tracing
+ """
+
+ def decorator(func):
+ if not _has_opentelemetry:
+ return func
+
+ @functools.wraps(func)
+ def wrapper(*args, **kwargs):
+ instance = args[0] if args and (hasattr(func, "__self__") and func.__self__ is not None) else None
+ is_method = instance is not None
+
+ if is_method and hasattr(instance, "tracer"):
+ tracer = instance.tracer
+ else:
+ tracer = get_tracer(f"transformers.{func.__module__}.{func.__name__}")
+
+ name = span_name or func.__name__
+ span_fn = tracer.start_span if standalone else tracer.start_as_current_span
+ with span_fn(name) as span:
+ span.set_attribute("function.name", func.__name__)
+ span.set_attribute("function.module", func.__module__)
+ span.set_attribute("function.is_method", is_method)
+
+ if args:
+ for i, arg in enumerate(args):
+ if isinstance(arg, (str, int, float, bool)) or arg is None:
+ span.set_attribute(f"args.{i}", str(arg))
+ else:
+ span.set_attribute(f"args.{i}", str(type(arg)))
+ if kwargs:
+ for key, value in kwargs.items():
+ if isinstance(value, (str, int, float, bool)) or value is None:
+ span.set_attribute(f"kwargs.{key}", str(value))
+ else:
+ span.set_attribute(f"kwargs.{key}", str(type(value)))
+
+ if additional_attributes and is_method:
+ for attr_config in additional_attributes:
+ instance_attribute_name, span_attribute_key, value_or_transform_function = attr_config
+ if hasattr(instance, instance_attribute_name):
+ attribute_value = getattr(instance, instance_attribute_name)
+ if callable(value_or_transform_function):
+ transformed_value = value_or_transform_function(attribute_value)
+ else:
+ transformed_value = value_or_transform_function
+ span.set_attribute(span_attribute_key, transformed_value)
+
+ try:
+ result = func(*args, **kwargs)
+ return result
+ except Exception as e:
+ span.set_status(Status(StatusCode.ERROR))
+ span.record_exception(e)
+ raise
+
+ return wrapper
+
+ if func is None:
+ return decorator
+ return decorator(func)
+
+
+logger = logging.getLogger(__name__)
+
+
+@attach_tracer()
+class ContinuousBatchProcessorMetrics:
+ """Metrics collection for ContinuousBatchProcessor."""
+
+ def __init__(self, max_batch_tokens: int):
+ """Initialize metrics for continuous batch processor.
+
+ Args:
+ max_batch_tokens: Maximum number of tokens in a batch
+ """
+ self.max_batch_tokens = max_batch_tokens
+
+ self._setup_metrics()
+
+ def _setup_metrics(self):
+ """Initialize OpenTelemetry metrics and tracing if the library is available."""
+
+ if not _has_opentelemetry:
+ logger.info(
+ "OpenTelemetry is not installed. Metrics and tracing will not be recorded."
+ "You can install it with `pip install opentelemetry-api>=1.30.0`"
+ )
+ return
+
+ self.meter = metrics.get_meter("transformers.generation.continuous_batch_processor")
+
+ # Define appropriate buckets for TTFT (typically ranges from ~50ms to several seconds)
+ ttft_buckets = [10, 25, 50, 75, 100, 150, 200, 300, 500, 750, 1000, 2000, 5000, 10000]
+
+ self.ttft_histogram = self.meter.create_histogram(
+ name="ttft_milliseconds",
+ description="Time to first token in milliseconds",
+ unit="ms",
+ explicit_bucket_boundaries_advisory=ttft_buckets,
+ )
+
+ self.active_requests_gauge = self.meter.create_gauge(
+ name="active_requests_count",
+ description="Number of active requests currently being processed",
+ unit="requests",
+ )
+
+ self.waiting_requests_gauge = self.meter.create_gauge(
+ name="waiting_requests_count",
+ description="Number of requests waiting to be processed",
+ unit="requests",
+ )
+
+ # Define appropriate buckets for request latency (similar to TTFT but with higher upper bounds)
+ latency_buckets = [50, 100, 250, 500, 1000, 2000, 5000, 10000, 20000, 30000, 60000]
+
+ self.request_latency_histogram = self.meter.create_histogram(
+ name="request_latency_milliseconds",
+ description="End-to-end latency for completed requests in milliseconds",
+ unit="ms",
+ explicit_bucket_boundaries_advisory=latency_buckets,
+ )
+
+ self.decode_prefill_ratio_gauge = self.meter.create_gauge(
+ name="decode_prefill_ratio",
+ description="Ratio of decode tokens to prefill tokens in a batch",
+ unit="ratio",
+ )
+
+ self.prefill_tokens_counter = self.meter.create_counter(
+ name="prefill_tokens_processed",
+ description="Number of prefill tokens processed",
+ unit="tokens",
+ )
+
+ self.decode_tokens_counter = self.meter.create_counter(
+ name="decode_tokens_processed",
+ description="Number of decode tokens processed",
+ unit="tokens",
+ )
+
+ # Define appropriate buckets for batch fill percentage (0-100%)
+ batch_fill_buckets = [5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 95, 98, 100]
+
+ self.batch_fill_percentage_histogram = self.meter.create_histogram(
+ name="batch_fill_percentage",
+ description="Percentage of max_batch_tokens utilized in each batch",
+ unit="percent",
+ explicit_bucket_boundaries_advisory=batch_fill_buckets,
+ )
+
+ self.kv_cache_free_memory_gauge = self.meter.create_gauge(
+ name="kv_cache_free_memory_bytes",
+ description="Free memory of the PagedAttentionCache in bytes",
+ unit="bytes",
+ )
+
+ self.kv_cache_memory_gauge = self.meter.create_gauge(
+ name="kv_cache_memory_bytes",
+ description="Memory usage of the PagedAttentionCache in bytes",
+ unit="bytes",
+ )
+
+ @traced
+ def record_ttft_metric(self, created_time: float, request_id: str) -> None:
+ """Record Time to First Token (TTFT).
+
+ Args:
+ created_time: The time the request was created
+ request_id: The ID of the request
+ """
+ if not _has_opentelemetry:
+ return
+
+ ttft_ms = (time.time() - created_time) * 1000.0
+
+ try:
+ self.ttft_histogram.record(ttft_ms)
+ logger.debug(f"Recorded TTFT for request {request_id}: {ttft_ms:.2f}ms")
+ except Exception as e:
+ logger.warning(f"Failed to record TTFT metric: {e}")
+
+ @traced
+ def record_batch_metrics(self, future_states: list) -> None:
+ """Record metrics about the batch composition including decode/prefill ratio and batch fill percentage.
+
+ Args:
+ requests_in_batch: List of request states in the current batch
+ """
+ if not _has_opentelemetry or not future_states:
+ return
+
+ decode_tokens = 0
+ prefill_tokens = 0
+
+ for future_state in future_states:
+ state = future_state.state
+ if state.status == RequestStatus.DECODING:
+ decode_tokens += 1
+ elif state.status in [RequestStatus.PREFILLING, RequestStatus.PREFILLING_SPLIT]:
+ prefill_tokens += len(state.prompt_ids)
+
+ total_batch_tokens = decode_tokens + prefill_tokens
+
+ try:
+ if prefill_tokens > 0:
+ self.prefill_tokens_counter.add(prefill_tokens)
+
+ if decode_tokens > 0:
+ self.decode_tokens_counter.add(decode_tokens)
+
+ if prefill_tokens > 0:
+ ratio = decode_tokens / prefill_tokens
+ self.decode_prefill_ratio_gauge.set(ratio)
+
+ fill_percentage = (total_batch_tokens / self.max_batch_tokens) * 100.0
+ self.batch_fill_percentage_histogram.record(fill_percentage)
+ logger.debug(
+ f"Batch metrics: {decode_tokens} decode tokens, {prefill_tokens} prefill tokens, "
+ f"batch fill: {fill_percentage:.2f}% ({total_batch_tokens}/{self.max_batch_tokens})"
+ )
+ except Exception as e:
+ logger.warning(f"Failed to record batch metrics: {e}")
+
+ @traced
+ def record_kv_cache_memory_metrics(self, cache) -> None:
+ """Record memory usage of the PagedAttentionCache without GPU synchronization.
+
+ This calculates the theoretical memory usage based on cache configuration
+ and the number of blocks currently in use.
+
+ Args:
+ cache: The PagedAttentionCache object to measure
+ """
+ if not _has_opentelemetry:
+ return
+
+ try:
+ # Retrieve the memory footprint of the cache
+ page_size = cache.head_dim * cache.num_key_value_heads
+ page_mem_in_bytes = page_size * cache.dtype.itemsize
+ # When a block is allocated, it is for both K and V, so we multiply by 2
+ # It's also allocated across all cache tensors, so we multiply by the nb of tensors: len(cache.key_cache)
+ block_mem_in_bytes = 2 * len(cache.key_cache) * cache.block_size * page_mem_in_bytes
+
+ # Retrieve the number of used and free blocks
+ free_blocks = cache.get_num_free_blocks()
+ used_blocks = cache.num_blocks - free_blocks
+
+ # Convert that into used and free memory in bytes
+ used_memory_bytes = used_blocks * block_mem_in_bytes
+ free_memory_bytes = free_blocks * block_mem_in_bytes
+
+ # Update the telemetry gauges and add a message in the logs
+ self.kv_cache_memory_gauge.set(used_memory_bytes)
+ self.kv_cache_free_memory_gauge.set(free_memory_bytes)
+ logger.debug(
+ f"KV Cache memory: {used_memory_bytes / (1024 * 1024):.2f}MB, "
+ f"Used blocks: {used_blocks}/{cache.num_blocks} "
+ f"({used_blocks / cache.num_blocks * 100:.1f}%)"
+ )
+ except Exception as e:
+ logger.warning(f"Failed to record KV cache memory metrics: {e}")
+
+ @traced
+ def record_queue_metrics(self, active_requests: int, waiting_requests: int) -> None:
+ """Record metrics about active and waiting requests.
+
+ Args:
+ active_requests: Number of active requests
+ waiting_requests: Number of waiting requests
+ """
+ if not _has_opentelemetry:
+ return
+
+ try:
+ self.active_requests_gauge.set(active_requests)
+ self.waiting_requests_gauge.set(waiting_requests)
+ logger.debug(f"Queue metrics: {active_requests} active requests, {waiting_requests} waiting requests")
+ except Exception as e:
+ logger.warning(f"Failed to record queue metrics: {e}")
+
+ @traced
+ def record_request_completion(self, created_time: float, request_id: str) -> None:
+ """Record metrics about a completed request.
+
+ Args:
+ created_time: The time the request was created
+ request_id: The ID of the request
+ """
+ if not _has_opentelemetry:
+ return
+
+ latency_ms = (time.time() - created_time) * 1000.0
+
+ try:
+ self.request_latency_histogram.record(latency_ms)
+
+ logger.debug(f"Recorded request completion for {request_id}: {latency_ms:.2f}ms")
+ except Exception as e:
+ logger.warning(f"Failed to record request completion metric: {e}")
diff --git a/third_party/transformers/src/transformers/utils/network_logging.py b/third_party/transformers/src/transformers/utils/network_logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..92f74ccd6d181bfa8b0fcf8a6fbfd8a39a19b4b8
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/network_logging.py
@@ -0,0 +1,485 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from __future__ import annotations
+
+import inspect
+import json
+import os
+import threading
+import time
+from collections import defaultdict
+from functools import wraps
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from .generic import strtobool
+
+
+class _NetworkRequestTrace:
+ def __init__(self, request: httpx.Request):
+ self.request = request
+ self.started_at = time.perf_counter()
+ self.phase_started_at = {}
+ self.phases_ms = defaultdict(float)
+
+ def trace(self, name: str, info: dict[str, Any]) -> None:
+ parts = name.rsplit(".", 2)
+ if len(parts) != 3:
+ return
+
+ _, phase, state = parts
+ now = time.perf_counter()
+ if state == "started":
+ self.phase_started_at[phase] = now
+ elif state in {"complete", "failed"}:
+ phase_started_at = self.phase_started_at.pop(phase, None)
+ if phase_started_at is not None:
+ self.phases_ms[phase] += (now - phase_started_at) * 1000
+
+ def build_record(
+ self,
+ *,
+ response: httpx.Response | None = None,
+ error: BaseException | None = None,
+ stream: bool = False,
+ ) -> dict[str, Any]:
+ total_ms = (time.perf_counter() - self.started_at) * 1000
+ url = self.request.url
+ host = url.host or ""
+ port = url.port
+ default_port = {"http": 80, "https": 443}.get(url.scheme)
+ host_display = host if port in (None, default_port) else f"{host}:{port}"
+
+ http_version = None
+ status_code = None
+ bytes_downloaded = None
+ response_complete = False
+ if response is not None:
+ status_code = response.status_code
+ response_complete = response.is_closed
+ raw_http_version = response.extensions.get("http_version")
+ if isinstance(raw_http_version, bytes):
+ http_version = raw_http_version.decode("ascii", errors="replace")
+ elif raw_http_version is not None:
+ http_version = str(raw_http_version)
+
+ if response_complete:
+ try:
+ bytes_downloaded = len(response.content)
+ except httpx.ResponseNotRead:
+ pass
+
+ return {
+ "method": self.request.method,
+ "scheme": url.scheme,
+ "host": host,
+ "host_display": host_display,
+ "port": port,
+ "path": url.path,
+ "has_query": bool(url.query),
+ "url": f"{url.scheme}://{host_display}{url.path}{'?...' if url.query else ''}",
+ "request_id": self.request.headers.get("x-amzn-trace-id") or self.request.headers.get("x-request-id"),
+ "status_code": status_code,
+ "http_version": http_version,
+ "bytes_downloaded": bytes_downloaded,
+ "total_ms": total_ms,
+ "stream": stream,
+ "response_complete": response_complete,
+ "phases_ms": dict(sorted(self.phases_ms.items())),
+ "error": None if error is None else f"{type(error).__name__}: {error}",
+ }
+
+
+class _NetworkDebugProfiler:
+ def __init__(self):
+ self._records = []
+ self._lock = threading.Lock()
+ self._enabled = False
+ self._output_path = None
+ self._original_client_send = None
+ self._original_async_client_send = None
+ self._shared_dir = None
+
+ @property
+ def enabled(self) -> bool:
+ return self._enabled
+
+ def clear(self) -> None:
+ with self._lock:
+ self._records = []
+
+ def enable(self, output_path: str | os.PathLike | None = None) -> None:
+ if self._enabled:
+ self._output_path = None if output_path is None else os.fspath(output_path)
+ self.clear()
+ return
+
+ self._output_path = None if output_path is None else os.fspath(output_path)
+ self.clear()
+
+ profiler = self
+ self._original_client_send = httpx.Client.send
+ self._original_async_client_send = httpx.AsyncClient.send
+
+ @wraps(self._original_client_send)
+ def patched_client_send(client, request, *args, **kwargs):
+ return profiler._send_with_trace(profiler._original_client_send, client, request, *args, **kwargs)
+
+ @wraps(self._original_async_client_send)
+ async def patched_async_client_send(client, request, *args, **kwargs):
+ return await profiler._async_send_with_trace(
+ profiler._original_async_client_send, client, request, *args, **kwargs
+ )
+
+ httpx.Client.send = patched_client_send
+ httpx.AsyncClient.send = patched_async_client_send
+ self._enabled = True
+
+ def setup_shared_dir(self) -> str | None:
+ """Create a shared temp directory for xdist workers to dump records into."""
+ if self._shared_dir is None:
+ import tempfile
+
+ self._shared_dir = tempfile.mkdtemp(prefix="network_debug_")
+ return self._shared_dir
+
+ def set_shared_dir(self, shared_dir: str) -> None:
+ """Set the shared directory (called in xdist workers)."""
+ self._shared_dir = shared_dir
+
+ def dump_worker_records(self, worker_id: str | None = None) -> None:
+ """Write this process's records to a file in the shared directory (called in workers)."""
+ if not self._shared_dir or not self._records:
+ return
+ worker_id = worker_id or f"pid{os.getpid()}"
+ dump_path = os.path.join(self._shared_dir, f"records_{worker_id}.json")
+ with self._lock:
+ records = [{**record, "phases_ms": dict(record["phases_ms"])} for record in self._records]
+ Path(dump_path).write_text(json.dumps(records), encoding="utf-8")
+
+ def load_worker_records(self) -> None:
+ """Load all worker record files from the shared directory (called in controller)."""
+ if not self._shared_dir or not os.path.isdir(self._shared_dir):
+ return
+ import glob as glob_module
+
+ for record_file in glob_module.glob(os.path.join(self._shared_dir, "records_*.json")):
+ try:
+ records = json.loads(Path(record_file).read_text(encoding="utf-8"))
+ with self._lock:
+ for record in records:
+ record["phases_ms"] = defaultdict(float, record.get("phases_ms", {}))
+ self._records.append(record)
+ except (OSError, json.JSONDecodeError):
+ pass
+
+ def cleanup_shared_dir(self) -> None:
+ """Remove the shared temp directory."""
+ if self._shared_dir and os.path.isdir(self._shared_dir):
+ import shutil
+
+ shutil.rmtree(self._shared_dir, ignore_errors=True)
+ self._shared_dir = None
+
+ def disable(self) -> None:
+ if not self._enabled:
+ return
+
+ httpx.Client.send = self._original_client_send
+ httpx.AsyncClient.send = self._original_async_client_send
+ self._enabled = False
+ self._original_client_send = None
+ self._original_async_client_send = None
+ self._output_path = None
+ self.clear()
+
+ def _append_record(self, record: dict[str, Any]) -> None:
+ with self._lock:
+ self._records.append(record)
+
+ def _wrap_trace_callback(self, request: httpx.Request, trace: _NetworkRequestTrace):
+ existing_trace = request.extensions.get("trace")
+
+ def wrapped_trace(name: str, info: dict[str, Any]) -> Any:
+ trace.trace(name, info)
+ if existing_trace is not None:
+ return existing_trace(name, info)
+ return None
+
+ return wrapped_trace
+
+ async def _awrap_trace_callback(self, request: httpx.Request, trace: _NetworkRequestTrace):
+ existing_trace = request.extensions.get("trace")
+
+ async def wrapped_trace(name: str, info: dict[str, Any]) -> Any:
+ trace.trace(name, info)
+ if existing_trace is not None:
+ result = existing_trace(name, info)
+ if inspect.isawaitable(result):
+ return await result
+ return result
+ return None
+
+ return wrapped_trace
+
+ def _send_with_trace(self, original_send, client, request: httpx.Request, *args, **kwargs):
+ trace = _NetworkRequestTrace(request)
+ request.extensions = dict(request.extensions)
+ request.extensions["trace"] = self._wrap_trace_callback(request, trace)
+
+ try:
+ response = original_send(client, request, *args, **kwargs)
+ except Exception as error:
+ self._append_record(trace.build_record(error=error, stream=kwargs.get("stream", False)))
+ raise
+
+ self._append_record(trace.build_record(response=response, stream=kwargs.get("stream", False)))
+ return response
+
+ async def _async_send_with_trace(self, original_send, client, request: httpx.Request, *args, **kwargs):
+ trace = _NetworkRequestTrace(request)
+ request.extensions = dict(request.extensions)
+ request.extensions["trace"] = await self._awrap_trace_callback(request, trace)
+
+ try:
+ response = await original_send(client, request, *args, **kwargs)
+ except Exception as error:
+ self._append_record(trace.build_record(error=error, stream=kwargs.get("stream", False)))
+ raise
+
+ self._append_record(trace.build_record(response=response, stream=kwargs.get("stream", False)))
+ return response
+
+ def build_report(self) -> dict[str, Any]:
+ with self._lock:
+ records = [
+ {
+ **record,
+ "phases_ms": dict(record["phases_ms"]),
+ }
+ for record in self._records
+ ]
+
+ phase_totals_ms = defaultdict(float)
+ route_totals = {}
+ for record in records:
+ for phase, duration_ms in record["phases_ms"].items():
+ phase_totals_ms[phase] += duration_ms
+
+ route_key = (record["method"], record["host_display"], record["path"])
+ route_total = route_totals.setdefault(
+ route_key,
+ {
+ "method": record["method"],
+ "host_display": record["host_display"],
+ "path": record["path"],
+ "count": 0,
+ "failures": 0,
+ "total_ms": 0.0,
+ "phase_totals_ms": defaultdict(float),
+ },
+ )
+ route_total["count"] += 1
+ route_total["total_ms"] += record["total_ms"]
+ route_total["failures"] += int(record["error"] is not None)
+ for phase, duration_ms in record["phases_ms"].items():
+ route_total["phase_totals_ms"][phase] += duration_ms
+
+ routes = []
+ for route_total in route_totals.values():
+ route_total["avg_ms"] = route_total["total_ms"] / route_total["count"]
+ route_total["phase_totals_ms"] = dict(sorted(route_total["phase_totals_ms"].items()))
+ routes.append(route_total)
+
+ routes.sort(key=lambda route: route["total_ms"], reverse=True)
+ total_time_ms = sum(record["total_ms"] for record in records)
+ return {
+ "enabled": self._enabled,
+ "output_path": self._output_path,
+ "total_requests": len(records),
+ "failed_requests": sum(int(record["error"] is not None) for record in records),
+ "total_time_ms": total_time_ms,
+ "phase_totals_ms": dict(sorted(phase_totals_ms.items())),
+ "requests": records,
+ "routes": routes,
+ }
+
+ def maybe_write_report(self) -> str | None:
+ if self._output_path is None:
+ return None
+
+ report_path = Path(self._output_path)
+ report_path.parent.mkdir(parents=True, exist_ok=True)
+ report_path.write_text(json.dumps(self.build_report(), indent=2, sort_keys=True), encoding="utf-8")
+ return str(report_path)
+
+
+_NETWORK_DEBUG_PROFILER = _NetworkDebugProfiler()
+
+
+_DEFAULT_REPORT_PATH = "network_debug_report.json"
+
+
+def _parse_network_debug_env() -> tuple[bool, str]:
+ enabled_raw = os.environ.get("NETWORK_DEBUG_REPORT", "").strip()
+ try:
+ enabled = bool(strtobool(enabled_raw)) if enabled_raw else False
+ except ValueError:
+ enabled = False
+
+ output_path = os.environ.get("NETWORK_DEBUG_REPORT_PATH", "").strip() or _DEFAULT_REPORT_PATH
+ return enabled, output_path
+
+
+def _enable_network_debug_report(output_path: str | os.PathLike | None = None) -> None:
+ _NETWORK_DEBUG_PROFILER.enable(output_path=output_path)
+
+
+def _disable_network_debug_report() -> None:
+ _NETWORK_DEBUG_PROFILER.disable()
+
+
+def _clear_network_debug_report() -> None:
+ _NETWORK_DEBUG_PROFILER.clear()
+
+
+def _get_network_debug_report() -> dict[str, Any]:
+ return _NETWORK_DEBUG_PROFILER.build_report()
+
+
+def _enable_network_debug_report_from_env() -> bool:
+ enabled, output_path = _parse_network_debug_env()
+ if not enabled:
+ return False
+
+ _enable_network_debug_report(output_path=output_path)
+ return True
+
+
+def _format_network_debug_report(max_requests: int = 20, max_routes: int = 10) -> str:
+ report = _get_network_debug_report()
+ if report["total_requests"] == 0:
+ return "Network debug report: no httpx requests captured."
+
+ lines = [
+ "Network debug report",
+ f"Requests captured: {report['total_requests']}",
+ f"Failed requests: {report['failed_requests']}",
+ f"Cumulative request time: {report['total_time_ms']:.1f} ms",
+ ]
+
+ if report["phase_totals_ms"]:
+ phase_summary = ", ".join(
+ f"{phase}={duration_ms:.1f} ms"
+ for phase, duration_ms in sorted(report["phase_totals_ms"].items(), key=lambda item: item[1], reverse=True)
+ )
+ lines.append(f"Phase totals: {phase_summary}")
+
+ lines.append("")
+ lines.append("Slowest requests:")
+ for idx, record in enumerate(
+ sorted(report["requests"], key=lambda request: request["total_ms"], reverse=True)[:max_requests],
+ start=1,
+ ):
+ status = record["error"] or f"status={record['status_code']}"
+ phase_bits = []
+ for phase in ("connect_tcp", "start_tls", "receive_response_headers", "receive_response_body"):
+ duration_ms = record["phases_ms"].get(phase)
+ if duration_ms is not None:
+ phase_bits.append(f"{phase}={duration_ms:.1f} ms")
+ phase_suffix = f" ({', '.join(phase_bits)})" if phase_bits else ""
+ incomplete_suffix = " incomplete" if record["stream"] and not record["response_complete"] else ""
+ lines.append(
+ f"{idx:>2}. {record['method']} {record['url']} {record['total_ms']:.1f} ms {status}{incomplete_suffix}{phase_suffix}"
+ )
+
+ lines.append("")
+ lines.append("Slowest routes:")
+ for idx, route in enumerate(report["routes"][:max_routes], start=1):
+ lines.append(
+ f"{idx:>2}. {route['method']} {route['host_display']}{route['path']} count={route['count']} "
+ f"total={route['total_ms']:.1f} ms avg={route['avg_ms']:.1f} ms failures={route['failures']}"
+ )
+
+ return "\n".join(lines)
+
+
+class NetworkDebugPlugin:
+ """Pytest plugin that handles all network debug orchestration including xdist coordination."""
+
+ def pytest_configure(self, config):
+ _enable_network_debug_report_from_env()
+ if not _NETWORK_DEBUG_PROFILER.enabled:
+ return
+
+ # xdist controller: create shared dir for workers to dump network records
+ if not hasattr(config, "workerinput"):
+ shared_dir = _NETWORK_DEBUG_PROFILER.setup_shared_dir()
+ if shared_dir:
+ config._network_debug_shared_dir = shared_dir
+ else:
+ # xdist worker: receive shared dir from controller
+ shared_dir = config.workerinput.get("network_debug_shared_dir")
+ if shared_dir:
+ _NETWORK_DEBUG_PROFILER.set_shared_dir(shared_dir)
+
+ def pytest_configure_node(self, node):
+ """xdist hook: called on the controller to configure each worker node."""
+ shared_dir = getattr(node.config, "_network_debug_shared_dir", None)
+ if shared_dir:
+ node.workerinput["network_debug_shared_dir"] = shared_dir
+
+ def pytest_sessionfinish(self, session, exitstatus):
+ # xdist worker: dump network debug records for the controller to aggregate
+ if hasattr(session.config, "workerinput"):
+ worker_id = session.config.workerinput.get("workerid", f"pid{os.getpid()}")
+ _NETWORK_DEBUG_PROFILER.dump_worker_records(worker_id=worker_id)
+
+ def pytest_terminal_summary(self, terminalreporter):
+ if not _NETWORK_DEBUG_PROFILER.enabled:
+ return
+
+ # Skip report generation in xdist worker processes; only the controller should aggregate and report.
+ if hasattr(terminalreporter.config, "workerinput"):
+ return
+
+ # Aggregate worker records if running under xdist.
+ _NETWORK_DEBUG_PROFILER.load_worker_records()
+
+ report_path = None
+ try:
+ report_path = _NETWORK_DEBUG_PROFILER.maybe_write_report()
+ except OSError as error:
+ report_path = f"Failed to write JSON report: {error}"
+
+ terminalreporter.section("Network debug", sep="=")
+ for line in _format_network_debug_report().splitlines():
+ terminalreporter.write_line(line)
+ if report_path is not None:
+ terminalreporter.write_line(f"JSON report: {report_path}")
+
+ _NETWORK_DEBUG_PROFILER.cleanup_shared_dir()
+
+
+def register_network_debug_plugin(config) -> None:
+ """Register the network debug pytest plugin. Single entry point for conftest.py."""
+ config.pluginmanager.register(NetworkDebugPlugin(), "network_debug")
+
+
+__all__ = [
+ "register_network_debug_plugin",
+]
diff --git a/third_party/transformers/src/transformers/utils/notebook.py b/third_party/transformers/src/transformers/utils/notebook.py
new file mode 100644
index 0000000000000000000000000000000000000000..ecbe8271fe139896c318fe69cbfc13a82c8a4269
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/notebook.py
@@ -0,0 +1,397 @@
+# Copyright 2020 Hugging Face
+#
+# 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.
+
+import os
+import re
+import time
+from typing import Optional, TypeVar
+
+import IPython.display as disp
+
+from ..trainer_callback import TrainerCallback
+from ..trainer_utils import IntervalStrategy, has_length
+
+
+_T = TypeVar("_T")
+
+
+def _require(x: _T | None, msg: str) -> _T:
+ if x is None:
+ raise RuntimeError(msg)
+ return x
+
+
+def format_time(t):
+ "Format `t` (in seconds) to (h):mm:ss"
+ t = int(t)
+ h, m, s = t // 3600, (t // 60) % 60, t % 60
+ return f"{h}:{m:02d}:{s:02d}" if h != 0 else f"{m:02d}:{s:02d}"
+
+
+def html_progress_bar(value, total, prefix, label, width=300):
+ # docstyle-ignore
+ return f"""
+
+ """
+
+
+def text_to_html_table(items):
+ "Put the texts in `items` in an HTML table."
+ html_code = """\n"""
+ html_code += """ \n \n"""
+ for i in items[0]:
+ html_code += f" | {i} | \n"
+ html_code += "
\n \n \n"
+ for line in items[1:]:
+ html_code += " \n"
+ for elt in line:
+ elt = f"{elt:.6f}" if isinstance(elt, float) else str(elt)
+ html_code += f" | {elt} | \n"
+ html_code += "
\n"
+ html_code += " \n
"
+ return html_code
+
+
+class NotebookProgressBar:
+ """
+ A progress par for display in a notebook.
+
+ Class attributes (overridden by derived classes)
+
+ - **warmup** (`int`) -- The number of iterations to do at the beginning while ignoring `update_every`.
+ - **update_every** (`float`) -- Since calling the time takes some time, we only do it every presumed
+ `update_every` seconds. The progress bar uses the average time passed up until now to guess the next value
+ for which it will call the update.
+
+ Args:
+ total (`int`):
+ The total number of iterations to reach.
+ prefix (`str`, *optional*):
+ A prefix to add before the progress bar.
+ leave (`bool`, *optional*, defaults to `True`):
+ Whether or not to leave the progress bar once it's completed. You can always call the
+ [`~utils.notebook.NotebookProgressBar.close`] method to make the bar disappear.
+ parent ([`~notebook.NotebookTrainingTracker`], *optional*):
+ A parent object (like [`~utils.notebook.NotebookTrainingTracker`]) that spawns progress bars and handle
+ their display. If set, the object passed must have a `display()` method.
+ width (`int`, *optional*, defaults to 300):
+ The width (in pixels) that the bar will take.
+
+ Example:
+
+ ```python
+ import time
+
+ pbar = NotebookProgressBar(100)
+ for val in range(100):
+ pbar.update(val)
+ time.sleep(0.07)
+ pbar.update(100)
+ ```"""
+
+ warmup = 5
+ update_every = 0.2
+
+ def __init__(
+ self,
+ total: int,
+ prefix: str | None = None,
+ leave: bool = True,
+ parent: Optional["NotebookTrainingTracker"] = None,
+ width: int = 300,
+ ):
+ self.total = total
+ self.prefix = "" if prefix is None else prefix
+ self.leave = leave
+ self.parent = parent
+ self.width = width
+ self.last_value = None
+ self.comment = None
+ self.output = None
+ self.value = None
+ self.label = None
+ if "VSCODE_PID" in os.environ:
+ self.update_every = 0.5 # Adjusted for smooth updated as html rending is slow on VS Code
+ # This is the only adjustment required to optimize training html rending
+
+ def update(self, value: int, force_update: bool = False, comment: str | None = None):
+ """
+ The main method to update the progress bar to `value`.
+
+ Args:
+ value (`int`):
+ The value to use. Must be between 0 and `total`.
+ force_update (`bool`, *optional*, defaults to `False`):
+ Whether or not to force and update of the internal state and display (by default, the bar will wait for
+ `value` to reach the value it predicted corresponds to a time of more than the `update_every` attribute
+ since the last update to avoid adding boilerplate).
+ comment (`str`, *optional*):
+ A comment to add on the left of the progress bar.
+ """
+ self.value = value
+ if comment is not None:
+ self.comment = comment
+ if self.last_value is None:
+ self.start_time = self.last_time = time.time()
+ self.start_value = self.last_value = value
+ self.elapsed_time = self.predicted_remaining = None
+ self.first_calls = self.warmup
+ self.wait_for = 1
+ self.update_bar(value)
+ elif value <= self.last_value and not force_update:
+ return
+ elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for, self.total):
+ if self.first_calls > 0:
+ self.first_calls -= 1
+ current_time = time.time()
+ self.elapsed_time = current_time - self.start_time
+ # We could have value = self.start_value if the update is called twixe with the same start value.
+ if value > self.start_value:
+ self.average_time_per_item = self.elapsed_time / (value - self.start_value)
+ else:
+ self.average_time_per_item = None
+ if value >= self.total:
+ value = self.total
+ self.predicted_remaining = None
+ if not self.leave:
+ self.close()
+ elif self.average_time_per_item is not None:
+ self.predicted_remaining = self.average_time_per_item * (self.total - value)
+ self.update_bar(value)
+ self.last_value = value
+ self.last_time = current_time
+ if (self.average_time_per_item is None) or (self.average_time_per_item == 0):
+ self.wait_for = 1
+ else:
+ self.wait_for = max(int(self.update_every / self.average_time_per_item), 1)
+
+ def update_bar(self, value, comment=None):
+ spaced_value = " " * (len(str(self.total)) - len(str(value))) + str(value)
+ if self.elapsed_time is None:
+ self.label = f"[{spaced_value}/{self.total} : < :"
+ elif self.predicted_remaining is None:
+ self.label = f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)}"
+ else:
+ self.label = (
+ f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)} <"
+ f" {format_time(self.predicted_remaining)}"
+ )
+ if self.average_time_per_item == 0:
+ self.label += ", +inf it/s"
+ else:
+ self.label += f", {1 / self.average_time_per_item:.2f} it/s"
+
+ self.label += "]" if self.comment is None or len(self.comment) == 0 else f", {self.comment}]"
+ self.display()
+
+ def display(self):
+ self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width)
+ if self.parent is not None:
+ # If this is a child bar, the parent will take care of the display.
+ self.parent.display()
+ return
+ if self.output is None:
+ self.output = disp.display(disp.HTML(self.html_code), display_id=True)
+ else:
+ self.output.update(disp.HTML(self.html_code))
+
+ def close(self):
+ "Closes the progress bar."
+ if self.parent is None and self.output is not None:
+ self.output.update(disp.HTML(""))
+
+
+class NotebookTrainingTracker(NotebookProgressBar):
+ """
+ An object tracking the updates of an ongoing training with progress bars and a nice table reporting metrics.
+
+ Args:
+ num_steps (`int`): The number of steps during training. column_names (`list[str]`, *optional*):
+ The list of column names for the metrics table (will be inferred from the first call to
+ [`~utils.notebook.NotebookTrainingTracker.write_line`] if not set).
+ """
+
+ def __init__(self, num_steps, column_names=None):
+ super().__init__(num_steps)
+ self.inner_table = None if column_names is None else [column_names]
+ self.child_bar = None
+
+ def display(self):
+ self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width)
+ if self.inner_table is not None:
+ self.html_code += text_to_html_table(self.inner_table)
+ if self.child_bar is not None:
+ self.html_code += self.child_bar.html_code
+ if self.output is None:
+ self.output = disp.display(disp.HTML(self.html_code), display_id=True)
+ else:
+ self.output.update(disp.HTML(self.html_code))
+
+ def write_line(self, values):
+ """
+ Write the values in the inner table.
+
+ Args:
+ values (`dict[str, float]`): The values to display.
+ """
+ if self.inner_table is None:
+ self.inner_table = [list(values.keys()), list(values.values())]
+ else:
+ columns = self.inner_table[0]
+ for key in values:
+ if key not in columns:
+ columns.append(key)
+ self.inner_table[0] = columns
+ if len(self.inner_table) > 1:
+ last_values = self.inner_table[-1]
+ first_column = self.inner_table[0][0]
+ if last_values[0] != values[first_column]:
+ # write new line
+ self.inner_table.append([values.get(c, "No Log") for c in columns])
+ else:
+ # update last line
+ new_values = values
+ for c in columns:
+ if c not in new_values:
+ new_values[c] = last_values[columns.index(c)]
+ self.inner_table[-1] = [new_values[c] for c in columns]
+ else:
+ self.inner_table.append([values[c] for c in columns])
+
+ def add_child(self, total, prefix=None, width=300):
+ """
+ Add a child progress bar displayed under the table of metrics. The child progress bar is returned (so it can be
+ easily updated).
+
+ Args:
+ total (`int`): The number of iterations for the child progress bar.
+ prefix (`str`, *optional*): A prefix to write on the left of the progress bar.
+ width (`int`, *optional*, defaults to 300): The width (in pixels) of the progress bar.
+ """
+ self.child_bar = NotebookProgressBar(total, prefix=prefix, parent=self, width=width)
+ return self.child_bar
+
+ def remove_child(self):
+ """
+ Closes the child progress bar.
+ """
+ self.child_bar = None
+ self.display()
+
+
+class NotebookProgressCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that displays the progress of training or evaluation, optimized for Jupyter Notebooks or
+ Google colab.
+ """
+
+ def __init__(self):
+ self.training_tracker = None
+ self.prediction_bar = None
+ self._force_next_update = False
+
+ def on_train_begin(self, args, state, control, **kwargs):
+ self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
+ self.training_loss = 0
+ self.last_log = 0
+ column_names = [self.first_column] + ["Training Loss"]
+ if args.eval_strategy != IntervalStrategy.NO:
+ column_names.append("Validation Loss")
+ self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
+
+ def on_step_end(self, args, state, control, **kwargs):
+ epoch = int(state.epoch) if int(state.epoch) == state.epoch else f"{state.epoch:.2f}"
+ tt = _require(self.training_tracker, "on_train_begin must be called before on_step_end")
+ tt.update(
+ state.global_step + 1,
+ comment=f"Epoch {epoch}/{state.num_train_epochs}",
+ force_update=self._force_next_update,
+ )
+ self._force_next_update = False
+
+ def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
+ if not has_length(eval_dataloader):
+ return
+ if self.prediction_bar is None:
+ if self.training_tracker is not None:
+ self.prediction_bar = self.training_tracker.add_child(len(eval_dataloader))
+ else:
+ self.prediction_bar = NotebookProgressBar(len(eval_dataloader))
+ self.prediction_bar.update(1)
+ else:
+ self.prediction_bar.update(self.prediction_bar.value + 1)
+
+ def on_predict(self, args, state, control, **kwargs):
+ if self.prediction_bar is not None:
+ self.prediction_bar.close()
+ self.prediction_bar = None
+
+ def on_log(self, args, state, control, logs=None, **kwargs):
+ # Only for when there is no evaluation
+ if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
+ tt = _require(self.training_tracker, "on_train_begin must be called before on_log")
+ values = {"Training Loss": logs["loss"]}
+ # First column is necessarily Step sine we're not in epoch eval strategy
+ values["Step"] = state.global_step
+ tt.write_line(values)
+
+ def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+ tt = _require(self.training_tracker, "on_train_begin must be called before on_evaluate")
+
+ values = {"Training Loss": "No log", "Validation Loss": "No log"}
+ for log in reversed(state.log_history):
+ if "loss" in log:
+ values["Training Loss"] = log["loss"]
+ break
+
+ if self.first_column == "Epoch":
+ values["Epoch"] = int(state.epoch)
+ else:
+ values["Step"] = state.global_step
+ if metrics is None:
+ metrics = {}
+ metric_key_prefix = "eval"
+ for k in metrics:
+ if k.endswith("_loss"):
+ metric_key_prefix = re.sub(r"\_loss$", "", k)
+ _ = metrics.pop("total_flos", None)
+ _ = metrics.pop("epoch", None)
+ _ = metrics.pop(f"{metric_key_prefix}_runtime", None)
+ _ = metrics.pop(f"{metric_key_prefix}_samples_per_second", None)
+ _ = metrics.pop(f"{metric_key_prefix}_steps_per_second", None)
+ for k, v in metrics.items():
+ splits = k.split("_")
+ name = " ".join([part.capitalize() for part in splits[1:]])
+ if name == "Loss":
+ # Single dataset
+ name = "Validation Loss"
+ values[name] = v
+ tt.write_line(values)
+ tt.remove_child()
+ self.prediction_bar = None
+ # Evaluation takes a long time so we should force the next update.
+ self._force_next_update = True
+
+ def on_train_end(self, args, state, control, **kwargs):
+ tt = _require(self.training_tracker, "on_train_begin must be called before on_train_end")
+ tt.update(
+ state.global_step,
+ comment=f"Epoch {int(state.epoch)}/{state.num_train_epochs}",
+ force_update=True,
+ )
+ self.training_tracker = None
diff --git a/third_party/transformers/src/transformers/utils/output_capturing.py b/third_party/transformers/src/transformers/utils/output_capturing.py
new file mode 100644
index 0000000000000000000000000000000000000000..aea31ab985ea12e9f3a34658d33ef36c7a0b0875
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/output_capturing.py
@@ -0,0 +1,285 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""
+Contains the logic for automatic additional output capture with our forward decorators.
+This mostly describe the hooks used and the logic to make capture thread/context safe.
+"""
+
+from __future__ import annotations
+
+import threading
+from contextvars import ContextVar
+from dataclasses import dataclass
+from functools import wraps
+from typing import TYPE_CHECKING
+
+from .import_utils import is_torchdynamo_compiling, requires
+
+
+if TYPE_CHECKING:
+ from torch import nn
+
+ from ..modeling_utils import PreTrainedModel
+
+
+_CAN_RECORD_REGISTRY = {}
+
+
+@dataclass
+@requires(backends=("torch",))
+class OutputRecorder:
+ """
+ Configuration for recording outputs from a model via hooks.
+
+ Attributes:
+ target_class (Type): The class (e.g., nn.Module) to which the hook will be attached.
+ index (Optional[int]): If the output is a tuple/list, optionally record only at a specific index.
+ layer_name (Optional[str]): Name of the submodule to target (if needed), e.g., "transformer.layer.3.attn".
+ class_name (Optional[str]): Name of the class to which the hook will be attached. Could be the suffix of class name in some cases.
+ """
+
+ target_class: type[nn.Module]
+ index: int = 0
+ layer_name: str | None = None
+ class_name: str | None = None
+
+
+class CompileableContextVar:
+ """
+ Convenience wrapper around a ContextVar for usage with `torch.compile`.
+ This behaves exactly as a `ContextVar`, except when compilation is triggered in which case it behaves as a simple
+ global variable. This is useful as `torch.compile` cannot trace the `get` method of `ContextVar`. This however means
+ that the access to the underlying variable is not thread-safe when compilation is triggered.
+ """
+
+ def __init__(self, name):
+ self.context_var = ContextVar(name, default=None)
+ self.global_var = None
+ self.compiling = False
+
+ def get(self):
+ # Set was called before and compilation was already detected
+ if self.compiling:
+ return self.global_var
+ else:
+ return self.context_var.get()
+
+ def set(self, value):
+ if is_torchdynamo_compiling():
+ self.global_var = value
+ self.compiling = True
+ return None
+ else:
+ return self.context_var.set(value)
+
+ def reset(self, token):
+ if self.compiling or token is None:
+ self.global_var = None
+ self.compiling = False
+ else:
+ self.context_var.reset(token)
+
+
+# Thread/context-safe global variable
+_active_collector = CompileableContextVar("output_collector")
+
+
+def install_output_capuring_hook(module: nn.Module, key: str, index: int) -> None:
+ """Install the forward hook needed to capture the output described by `key` and `index` in `module`."""
+
+ def output_capturing_hook(module, args, output):
+ # Get the current thread-local collector
+ collected_outputs = _active_collector.get()
+ # If it's None or not a key we want to capture, simply return, the hook is inactive
+ if collected_outputs is None or key not in collected_outputs.keys():
+ return
+
+ if key == "hidden_states" and len(collected_outputs[key]) == 0:
+ collected_outputs[key].append(args[0])
+ if not isinstance(output, tuple):
+ collected_outputs[key].append(output)
+ elif output[index] is not None:
+ collected_outputs[key].append(output[index])
+
+ module.register_forward_hook(output_capturing_hook)
+
+
+def recursively_install_hooks(
+ parent_module: nn.Module, module_name: str, capture_tasks: list[tuple[str, OutputRecorder]]
+) -> None:
+ """
+ Recursively install all output capturing hooks on all submodules of `parent_module`.
+ Note that we need to use this recursive approach instead of simply iterating over all modules, because we want
+ to respect the `capture_tasks` of all individual submodels (`PreTrainedModel` instances) in the graph. That is, once
+ we reach a submodel in the graph, its children should use this submodel's `capture_tasks`, but other parts of the graph
+ should not.
+ """
+ from ..modeling_utils import PreTrainedModel
+
+ # First dispatch to children if needed
+ for name, module in parent_module.named_children():
+ # Keep dispatching the same `capture_tasks`
+ if not isinstance(module, PreTrainedModel):
+ recursively_install_hooks(module, f"{module_name}.{name}", capture_tasks)
+ # New Submodel: we need to dispatch its own `capture_tasks`
+ else:
+ install_all_output_capturing_hooks(module, prefix=f"{module_name}.{name}")
+
+ # Potentially install the hook on current `parent_module`
+ for key, specs in capture_tasks:
+ # The second check is for multimodals where only backbone layer suffix is available
+ if (specs.target_class is not None and isinstance(parent_module, specs.target_class)) or (
+ specs.class_name is not None and module_name.endswith(specs.class_name)
+ ):
+ if specs.layer_name is not None and specs.layer_name not in module_name:
+ continue
+ install_output_capuring_hook(parent_module, key, specs.index)
+
+
+def install_all_output_capturing_hooks(model: PreTrainedModel, prefix: str | None = None) -> None:
+ """
+ Install the output recording hooks on all the modules in `model`. Tis will take care of correctly dispatching
+ the `_can_record_outputs` property of each individual submodels in case of composite models.
+ """
+ # _can_record_outputs is None by default
+ capture_flags = _CAN_RECORD_REGISTRY.get(str(model.__class__)) or {} # there is a weak ref for executorch
+
+ capture_tasks = []
+ for key, layer_specs in capture_flags.items():
+ if not isinstance(layer_specs, list):
+ layer_specs = [layer_specs]
+ for specs in layer_specs:
+ if not isinstance(specs, OutputRecorder):
+ index = 0 if "hidden_states" in key else 1
+ class_name = None if not isinstance(specs, str) else specs
+ target_class = specs if not isinstance(specs, str) else None
+ specs = OutputRecorder(target_class=target_class, index=index, class_name=class_name)
+ capture_tasks.append((key, specs))
+
+ # Install the hooks
+ prefix = prefix if prefix is not None else ""
+ recursively_install_hooks(model, prefix, capture_tasks)
+ # Mark the model as already hooked
+ setattr(model, "_output_capturing_hooks_installed", True)
+
+
+# We need this to make sure we don't have race conditions when installing hooks, resulting in them being installed
+# several times
+_hook_installation_lock = threading.Lock()
+
+
+def maybe_install_capturing_hooks(model: PreTrainedModel) -> None:
+ """
+ Check if the model already has output capturing hooks installed, and install them if it is not already the
+ case.
+ Note that this is thread-safe, in case 2 (or more) threads want to install them concurrently.
+ """
+ # First check
+ if getattr(model, "_output_capturing_hooks_installed", False):
+ return
+
+ with _hook_installation_lock:
+ # Second check, in case several threads entered this function concurrently and did not return on the
+ # previous check
+ if getattr(model, "_output_capturing_hooks_installed", False):
+ return
+ # This will install the hooks and mark the model as hooked
+ install_all_output_capturing_hooks(model)
+
+
+def capture_outputs(func=None, *, tie_last_hidden_states=True):
+ """
+ Decorator to intercept specific layer outputs through hooks. The hooks are installed only once and lazily,
+ the first time output capture is requested with the `output_xxx` kwargs/config.
+ The implementation is fully context/thread safe, except when using `torch.compile`, as dynamo is unable to trace
+ through `ContextVar` methods.
+
+ Args:
+ tie_last_hidden_states (`bool`, *optional*, defaults to `True`):
+ Whether to overwrite `out.hidden_states[-1]` with the `out.last_hidden_state`.
+ This is true for all language models and should be toggled off only if
+ `out.hidden_states[-1]` has to be the hidden state before last layer norm, which
+ is needed for some vision models (e.g. CLIP, SigLIP)
+ """
+
+ def wrapped_fn(func):
+ @wraps(func)
+ def wrapper(self, *args, **kwargs):
+ # Pop it so that internal modules always return a dict even if False is requested
+ return_dict = kwargs.pop("return_dict", getattr(self.config, "return_dict", True))
+
+ # _can_record_outputs is None by default
+ capturable_flags = _CAN_RECORD_REGISTRY.get(str(self.__class__)) or {}
+ recordable_keys = {
+ f"output_{k}": kwargs.get(f"output_{k}", getattr(self.config, f"output_{k}", False))
+ for k in capturable_flags
+ }
+ # For BC as cross-attentions used to be captured with `output_attentions`
+ if "cross_attentions" in capturable_flags:
+ recordable_keys["output_cross_attentions"] = kwargs.get(
+ "output_attentions", getattr(self.config, "output_attentions", False)
+ )
+ # The sam model variants need this annoying exception as well...
+ if "mask_decoder_attentions" in capturable_flags:
+ recordable_keys["output_mask_decoder_attentions"] = kwargs.get(
+ "output_attentions", getattr(self.config, "output_attentions", False)
+ )
+
+ collected_outputs = {k.replace("output_", ""): [] for k, v in recordable_keys.items() if v}
+ # Make sure hooks are installed if we need to collect outputs
+ if len(collected_outputs) > 0:
+ maybe_install_capturing_hooks(self)
+ # Let's activate the output collector hooks if needed!
+ output_token = _active_collector.set(collected_outputs)
+
+ # Run the forward
+ try:
+ outputs = func(self, *args, **kwargs)
+ # Reset the states
+ finally:
+ _active_collector.reset(output_token)
+
+ # Inject collected outputs into model output (return everything as tuples for BC)
+ for key in collected_outputs:
+ if key == "hidden_states":
+ if not tie_last_hidden_states:
+ pass
+ elif hasattr(outputs, "vision_hidden_states"):
+ collected_outputs[key] = collected_outputs[key][:-1]
+ collected_outputs[key].append(outputs.vision_hidden_states)
+ elif hasattr(outputs, "last_hidden_state"):
+ collected_outputs[key] = collected_outputs[key][:-1]
+ collected_outputs[key].append(outputs.last_hidden_state)
+
+ outputs[key] = tuple(collected_outputs[key])
+ elif key == "attentions":
+ # In this case, the second item are cross attentions
+ if isinstance(capturable_flags[key], list) and len(capturable_flags[key]) == 2:
+ outputs[key] = tuple(collected_outputs[key][0::2])
+ outputs["cross_" + key] = tuple(collected_outputs[key][1::2])
+ else:
+ outputs[key] = tuple(collected_outputs[key])
+ else:
+ outputs[key] = tuple(collected_outputs[key])
+
+ if return_dict is False:
+ outputs = outputs.to_tuple()
+
+ return outputs
+
+ return wrapper
+
+ if func is not None:
+ return wrapped_fn(func)
+ return wrapped_fn
diff --git a/third_party/transformers/src/transformers/utils/peft_utils.py b/third_party/transformers/src/transformers/utils/peft_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1ec093b4e672c362e9fcaecd4363d29339d4b6e
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/peft_utils.py
@@ -0,0 +1,117 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# 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.
+import importlib
+import importlib.metadata
+import os
+
+from packaging import version
+
+from .hub import cached_file
+from .import_utils import is_peft_available
+
+
+ADAPTER_CONFIG_NAME = "adapter_config.json"
+ADAPTER_WEIGHTS_NAME = "adapter_model.bin"
+ADAPTER_SAFE_WEIGHTS_NAME = "adapter_model.safetensors"
+
+
+def find_adapter_config_file(
+ model_id: str,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ subfolder: str = "",
+ _commit_hash: str | None = None,
+) -> str | None:
+ r"""
+ Simply checks if the model stored on the Hub or locally is an adapter model or not, return the path of the adapter
+ config file if it is, None otherwise.
+
+ Args:
+ model_id (`str`):
+ The identifier of the model to look for, can be either a local path or an id to the repository on the Hub.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+
+
+
+ To test a pull request you made on the Hub, you can pass `revision="refs/pr/".
+
+
+
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ subfolder (`str`, *optional*, defaults to `""`):
+ In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
+ specify the folder name here.
+ """
+ adapter_cached_filename = None
+ if model_id is None:
+ return None
+ elif os.path.isdir(model_id):
+ list_remote_files = os.listdir(model_id)
+ if ADAPTER_CONFIG_NAME in list_remote_files:
+ adapter_cached_filename = os.path.join(model_id, ADAPTER_CONFIG_NAME)
+ else:
+ adapter_cached_filename = cached_file(
+ model_id,
+ ADAPTER_CONFIG_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ subfolder=subfolder,
+ _commit_hash=_commit_hash,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ )
+
+ return adapter_cached_filename
+
+
+def check_peft_version(min_version: str) -> None:
+ r"""
+ Checks if the version of PEFT is compatible.
+
+ Args:
+ version (`str`):
+ The version of PEFT to check against.
+ """
+ if not is_peft_available():
+ raise ValueError("PEFT is not installed. Please install it with `pip install peft`")
+
+ is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) >= version.parse(min_version)
+
+ if not is_peft_version_compatible:
+ raise ValueError(f"The version of PEFT you are using is not compatible, please use a version >= {min_version}")
diff --git a/third_party/transformers/src/transformers/utils/pytest_helpers.py b/third_party/transformers/src/transformers/utils/pytest_helpers.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f22e01ba5081318e38f591d15f29ad833a7aa2d
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/pytest_helpers.py
@@ -0,0 +1,111 @@
+import argparse
+import json
+import re
+from collections import Counter
+from pathlib import Path
+
+
+def _base_test_name(nodeid: str) -> str:
+ # Strip parameters like [param=..] from the last component
+ name = nodeid.split("::")[-1]
+ return re.sub(r"\[.*\]$", "", name)
+
+
+def _class_name(nodeid: str) -> str | None:
+ parts = nodeid.split("::")
+ # nodeid can be: file::Class::test or file::test
+ if len(parts) >= 3:
+ return parts[-2]
+ return None
+
+
+def _file_path(nodeid: str) -> str:
+ return nodeid.split("::")[0]
+
+
+def _modeling_key(file_path: str) -> str | None:
+ # Extract "xxx" from test_modeling_xxx.py
+ m = re.search(r"test_modeling_([A-Za-z0-9_]+)\.py$", file_path)
+ if m:
+ return m.group(1)
+ return None
+
+
+def summarize(report_path: str):
+ p = Path(report_path)
+ if not p.exists():
+ raise FileNotFoundError(f"Report file not found: {p.resolve()}")
+
+ data = json.loads(p.read_text())
+ tests = data.get("tests", [])
+
+ # Overall counts
+ outcomes = Counter(t.get("outcome", "unknown") for t in tests)
+
+ # Filter failures (pytest-json-report uses "failed" and may have "error")
+ failed = [t for t in tests if t.get("outcome") in ("failed", "error")]
+
+ # 1) Failures per test file
+ failures_per_file = Counter(_file_path(t.get("nodeid", "")) for t in failed)
+
+ # 2) Failures per class (if any; otherwise "NO_CLASS")
+ failures_per_class = Counter((_class_name(t.get("nodeid", "")) or "NO_CLASS") for t in failed)
+
+ # 3) Failures per base test name (function), aggregating parametrized cases
+ failures_per_testname = Counter(_base_test_name(t.get("nodeid", "")) for t in failed)
+
+ # 4) Failures per test_modeling_xxx (derived from filename)
+ failures_per_modeling_key = Counter()
+ for t in failed:
+ key = _modeling_key(_file_path(t.get("nodeid", "")))
+ if key:
+ failures_per_modeling_key[key] += 1
+
+ return {
+ "outcomes": outcomes,
+ "failures_per_file": failures_per_file,
+ "failures_per_class": failures_per_class,
+ "failures_per_testname": failures_per_testname,
+ "failures_per_modeling_key": failures_per_modeling_key,
+ }
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Summarize pytest JSON report failures")
+ parser.add_argument(
+ "--report", default="report.json", help="Path to pytest JSON report file (default: report.json)"
+ )
+ args = parser.parse_args()
+
+ try:
+ summary = summarize(args.report)
+ except FileNotFoundError as e:
+ print(str(e))
+ return
+
+ outcomes = summary["outcomes"]
+ print("=== Overall ===")
+ total = sum(outcomes.values())
+ print(f"Total tests: {total}")
+ for k in sorted(outcomes):
+ print(f"{k:>10}: {outcomes[k]}")
+
+ def _print_counter(title, counter: Counter, label=""):
+ print(f"\n=== {title} ===")
+ if not counter:
+ print("None")
+ return
+ for key, cnt in sorted(counter.items(), key=lambda x: (x[1], x[0])):
+ if label:
+ print(f"{cnt:4d} {label}{key}")
+ else:
+ print(f"{cnt:4d} {key}")
+
+ _print_counter("Failures per test class", summary["failures_per_class"], label="class ")
+ _print_counter("Failures per test_modeling_xxx", summary["failures_per_modeling_key"], label="model ")
+ _print_counter("Failures per test file", summary["failures_per_file"])
+ _print_counter("Failures per test name (base)", summary["failures_per_testname"])
+
+
+if __name__ == "__main__":
+ main()
diff --git a/third_party/transformers/src/transformers/utils/quantization_config.py b/third_party/transformers/src/transformers/utils/quantization_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..908fb69fa2f85d84311c96593d75136ad7678d86
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/quantization_config.py
@@ -0,0 +1,1990 @@
+#!/usr/bin/env python
+
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# 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.
+import copy
+import importlib.metadata
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional, Union
+
+from packaging import version
+
+from ..utils import (
+ is_compressed_tensors_available,
+ is_hqq_available,
+ is_quark_available,
+ is_torch_available,
+ is_torchao_available,
+ logging,
+)
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class QuantizationMethod(str, Enum):
+ BITS_AND_BYTES = "bitsandbytes"
+ GPTQ = "gptq"
+ AWQ = "awq"
+ AQLM = "aqlm"
+ VPTQ = "vptq"
+ QUANTO = "quanto"
+ EETQ = "eetq"
+ HIGGS = "higgs"
+ HQQ = "hqq"
+ COMPRESSED_TENSORS = "compressed-tensors"
+ FBGEMM_FP8 = "fbgemm_fp8"
+ TORCHAO = "torchao"
+ BITNET = "bitnet"
+ SPQR = "spqr"
+ FP8 = "fp8"
+ QUARK = "quark"
+ FPQUANT = "fp_quant"
+ AUTOROUND = "auto-round"
+ MXFP4 = "mxfp4"
+ METAL = "metal"
+ FOUR_OVER_SIX = "fouroversix"
+ SINQ = "sinq"
+
+
+class AwqFormat(str, Enum):
+ GEMM = "gemm"
+ GEMV = "gemv"
+ GEMV_FAST = "gemv_fast"
+ LLM_AWQ = "llm-awq"
+
+
+class AwqBackend(str, Enum):
+ LEGACY_AWQ = "autoawq"
+ AUTO = "auto"
+ AUTO_TRAINABLE = "auto_trainable"
+ MACHETE = "machete"
+ MARLIN = "marlin"
+ EXLLAMA_V2 = "exllama_v2"
+ EXLLAMA_V1 = "exllama_v1"
+ GEMM = "gemm"
+ GEMM_TRITON = "gemm_triton"
+ GEMV = "gemv"
+ GEMV_FAST = "gemv_fast"
+ TORCH_AWQ = "torch_awq"
+ TORCH_FUSED_AWQ = "torch_fused_awq"
+
+
+@dataclass
+class QuantizationConfigMixin:
+ """
+ Mixin class for quantization config
+ """
+
+ quant_method: QuantizationMethod
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """
+ Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters.
+
+ Args:
+ config_dict (`dict[str, Any]`):
+ Dictionary that will be used to instantiate the configuration object.
+ return_unused_kwargs (`bool`,*optional*, defaults to `False`):
+ Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
+ `PreTrainedModel`.
+ kwargs (`dict[str, Any]`):
+ Additional parameters from which to initialize the configuration object.
+
+ Returns:
+ [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters.
+ """
+ config = cls(**config_dict)
+
+ to_remove = []
+ for key, value in kwargs.items():
+ if hasattr(config, key):
+ setattr(config, key, value)
+ to_remove.append(key)
+ for key in to_remove:
+ kwargs.pop(key, None)
+
+ if return_unused_kwargs:
+ return config, kwargs
+ else:
+ return config
+
+ def to_json_file(self, json_file_path: str | os.PathLike):
+ """
+ Save this instance to a JSON file.
+
+ Args:
+ json_file_path (`str` or `os.PathLike`):
+ Path to the JSON file in which this configuration instance's parameters will be saved.
+ use_diff (`bool`, *optional*, defaults to `True`):
+ If set to `True`, only the difference between the config instance and the default
+ `QuantizationConfig()` is serialized to JSON file.
+ """
+ with open(json_file_path, "w", encoding="utf-8") as writer:
+ config_dict = self.to_dict()
+ json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+ writer.write(json_string)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ return copy.deepcopy(self.__dict__)
+
+ def __iter__(self):
+ """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin"""
+ yield from copy.deepcopy(self.__dict__).items()
+
+ def __repr__(self):
+ return f"{self.__class__.__name__} {self.to_json_string()}"
+
+ def to_diff_dict(self) -> dict[str, Any]:
+ """
+ Default behavior: no diffing implemented for this config.
+ """
+ return self.to_dict()
+
+ def to_json_string(self, use_diff: bool = True) -> str:
+ """
+ Serializes this instance to a JSON string.
+
+ Args:
+ use_diff (`bool`, *optional*, defaults to `True`):
+ If set to `True`, only the difference between the config instance and the default `PreTrainedConfig()`
+ is serialized to JSON string.
+
+ Returns:
+ `str`: String containing all the attributes that make up this configuration instance in JSON format.
+ """
+ config_dict = self.to_diff_dict() if use_diff else self.to_dict()
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+ def update(self, **kwargs):
+ """
+ Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes,
+ returning all the unused kwargs.
+
+ Args:
+ kwargs (`dict[str, Any]`):
+ Dictionary of attributes to tentatively update this class.
+
+ Returns:
+ `dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
+ """
+ to_remove = []
+ for key, value in kwargs.items():
+ if hasattr(self, key):
+ setattr(self, key, value)
+ to_remove.append(key)
+
+ # Remove all the attributes that were updated, without modifying the input dict
+ unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove}
+ return unused_kwargs
+
+
+@dataclass
+class AutoRoundConfig(QuantizationConfigMixin):
+ """This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded AutoRound quantization.
+
+ Args:
+ bits (`int`, *optional*, defaults to 4):
+ The number of bits to quantize to, supported numbers are (2, 3, 4, 8).
+ group_size (`int`, *optional*, defaults to 128): Group-size value
+ sym (`bool`, *optional*, defaults to `True`): Symmetric quantization or not
+ backend (`str`, *optional*, defaults to `"auto"`): The kernel to use, e.g., ipex,marlin, exllamav2, triton, etc. Ref. https://github.com/intel/auto-round?tab=readme-ov-file#specify-backend
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ group_size: int = 128,
+ sym: bool = True,
+ backend: str = "auto",
+ **kwargs,
+ ):
+ self.bits = bits
+ self.group_size = group_size
+ self.sym = sym
+ self.backend = backend
+ self.packing_format = "auto_round:gptq"
+ if kwargs is not None:
+ for key, value in kwargs.items():
+ setattr(self, key, value)
+ self.quant_method = QuantizationMethod.AUTOROUND
+ self.post_init()
+
+ def post_init(self):
+ r"""Safety checker that arguments are correct."""
+ if self.bits not in [2, 3, 4, 8]:
+ raise ValueError(f"Only support quantization to [2,3,4,8] bits but found {self.bits}")
+ if self.group_size != -1 and self.group_size <= 0:
+ raise ValueError("group_size must be greater than 0 or equal to -1")
+
+ def get_loading_attributes(self):
+ loading_attributes_dict = {"backend": self.backend}
+ return loading_attributes_dict
+
+ def to_dict(self):
+ config_dict = super().to_dict()
+ return config_dict
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ quant_method = config_dict["quant_method"]
+ if "auto-round" not in quant_method and "gptq" not in quant_method and "awq" not in quant_method:
+ raise NotImplementedError(
+ "Failed to convert to auto_round format. Only `gptqv1`, `awq`, and `auto-round` formats are supported."
+ )
+
+ if "gptq" in quant_method and "meta" in config_dict:
+ raise NotImplementedError("Failed to convert gptq format to auto_round format. Only supports `gptqv1`")
+
+ if "awq" in quant_method and config_dict.get("version", "gemm") != "gemm":
+ raise NotImplementedError(
+ "Failed to convert awq format to auto_round format. Only supports awq format with gemm version"
+ )
+
+ if "auto-round" not in quant_method:
+ config_dict["packing_format"] = f"auto_round:{quant_method}"
+
+ return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs)
+
+
+@dataclass
+class HqqConfig(QuantizationConfigMixin):
+ """
+ This is wrapper around hqq's BaseQuantizeConfig.
+
+ Args:
+ nbits (`int`, *optional*, defaults to 4):
+ Number of bits. Supported values are (8, 4, 3, 2, 1).
+ group_size (`int`, *optional*, defaults to 64):
+ Group-size value. Supported values are any value that is divisible by weight.shape[axis]).
+ view_as_float (`bool`, *optional*, defaults to `False`):
+ View the quantized weight as float (used in distributed training) if set to `True`.
+ axis (`Optional[int]`, *optional*):
+ Axis along which grouping is performed. Supported values are 0 or 1.
+ dynamic_config (dict, *optional*):
+ Parameters for dynamic configuration. The key is the name tag of the layer and the value is a quantization config.
+ If set, each layer specified by its id will use its dedicated quantization configuration.
+ skip_modules (`list[str]`, *optional*, defaults to `['lm_head']`):
+ List of `nn.Linear` layers to skip.
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ nbits: int = 4,
+ group_size: int = 64,
+ view_as_float: bool = False,
+ axis: int | None = None,
+ dynamic_config: dict | None = None,
+ skip_modules: list[str] = ["lm_head"],
+ **kwargs,
+ ):
+ if is_hqq_available():
+ from hqq.core.quantize import BaseQuantizeConfig as HQQBaseQuantizeConfig
+ else:
+ raise ImportError(
+ "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`."
+ )
+
+ if axis is None:
+ axis = 1
+ logger.info("Setting axis=1 as faster backends such as TorchAO or BitBlas are only compatible with it.")
+
+ if axis not in [0, 1]:
+ raise ValueError("Invalid axis value. Only 0 and 1 are allowed.")
+
+ if dynamic_config is not None:
+ self.quant_config = {}
+ for key in dynamic_config:
+ self.quant_config[key] = HQQBaseQuantizeConfig(**dynamic_config[key])
+ else:
+ self.quant_config = HQQBaseQuantizeConfig(
+ nbits=nbits, group_size=group_size, view_as_float=view_as_float, axis=axis
+ )
+
+ self.quant_method = QuantizationMethod.HQQ
+ self.skip_modules = skip_modules
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+
+ @classmethod
+ def from_dict(cls, config: dict[str, Any]):
+ """
+ Override from_dict, used in AutoQuantizationConfig.from_dict in quantizers/auto.py
+ """
+ instance = cls()
+ instance.quant_config = config["quant_config"]
+ instance.skip_modules = config["skip_modules"]
+ return instance
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ return {
+ "quant_config": self.quant_config,
+ "quant_method": self.quant_method,
+ "skip_modules": self.skip_modules,
+ }
+
+ def __repr__(self):
+ config_dict = self.to_dict()
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
+
+ def to_diff_dict(self) -> dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+ Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = HqqConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+
+@dataclass
+class BitsAndBytesConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `bitsandbytes`.
+
+ Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`,
+ then more arguments will be added to this class.
+
+ Args:
+ load_in_8bit (`bool`, *optional*, defaults to `False`):
+ This flag is used to enable 8-bit quantization with LLM.int8().
+ load_in_4bit (`bool`, *optional*, defaults to `False`):
+ This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from
+ `bitsandbytes`.
+ llm_int8_threshold (`float`, *optional*, defaults to 6.0):
+ This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix
+ Multiplication for Transformers at Scale` paper: https://huggingface.co/papers/2208.07339 Any hidden states value
+ that is above this threshold will be considered an outlier and the operation on those values will be done
+ in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but
+ there are some exceptional systematic outliers that are very differently distributed for large models.
+ These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
+ magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
+ but a lower threshold might be needed for more unstable models (small models, fine-tuning).
+ llm_int8_skip_modules (`list[str]`, *optional*):
+ An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
+ Jukebox that has several heads in different places and not necessarily at the last position. For example
+ for `CausalLM` models, the last `lm_head` is kept in its original `dtype`.
+ llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`):
+ This flag is used for advanced use cases and users that are aware of this feature. If you want to split
+ your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use
+ this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8
+ operations will not be run on CPU.
+ llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`):
+ This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not
+ have to be converted back and forth for the backward pass.
+ bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`):
+ This sets the computational type which might be different than the input type. For example, inputs might be
+ fp32, but computation can be set to bf16 for speedups.
+ bnb_4bit_quant_type (`str`, *optional*, defaults to `"fp4"`):
+ This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types
+ which are specified by `fp4` or `nf4`.
+ bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`):
+ This flag is used for nested quantization where the quantization constants from the first quantization are
+ quantized again.
+ bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`):
+ This sets the storage type to pack the quantized 4-bit params.
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ load_in_8bit=False,
+ load_in_4bit=False,
+ llm_int8_threshold=6.0,
+ llm_int8_skip_modules=None,
+ llm_int8_enable_fp32_cpu_offload=False,
+ llm_int8_has_fp16_weight=False,
+ bnb_4bit_compute_dtype=None,
+ bnb_4bit_quant_type="fp4",
+ bnb_4bit_use_double_quant=False,
+ bnb_4bit_quant_storage=None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.BITS_AND_BYTES
+
+ if load_in_4bit and load_in_8bit:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+
+ self._load_in_8bit = load_in_8bit
+ self._load_in_4bit = load_in_4bit
+ self.llm_int8_threshold = llm_int8_threshold
+ self.llm_int8_skip_modules = llm_int8_skip_modules
+ self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
+ self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
+ self.bnb_4bit_quant_type = bnb_4bit_quant_type
+ self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
+
+ if bnb_4bit_compute_dtype is None:
+ self.bnb_4bit_compute_dtype = torch.float32
+ elif isinstance(bnb_4bit_compute_dtype, str):
+ self.bnb_4bit_compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
+ elif isinstance(bnb_4bit_compute_dtype, torch.dtype):
+ self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
+ else:
+ raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype")
+
+ if bnb_4bit_quant_storage is None:
+ self.bnb_4bit_quant_storage = torch.uint8
+ elif isinstance(bnb_4bit_quant_storage, str):
+ if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]:
+ raise ValueError(
+ "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') "
+ )
+ self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage)
+ elif isinstance(bnb_4bit_quant_storage, torch.dtype):
+ self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
+ else:
+ raise ValueError("bnb_4bit_quant_storage must be a string or a torch.dtype")
+
+ if kwargs:
+ logger.info(f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__}.")
+
+ self.post_init()
+
+ @property
+ def load_in_4bit(self):
+ return self._load_in_4bit
+
+ @load_in_4bit.setter
+ def load_in_4bit(self, value: bool):
+ if not isinstance(value, bool):
+ raise TypeError("load_in_4bit must be a boolean")
+
+ if self.load_in_8bit and value:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+ self._load_in_4bit = value
+
+ @property
+ def load_in_8bit(self):
+ return self._load_in_8bit
+
+ @load_in_8bit.setter
+ def load_in_8bit(self, value: bool):
+ if not isinstance(value, bool):
+ raise TypeError("load_in_8bit must be a boolean")
+
+ if self.load_in_4bit and value:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+ self._load_in_8bit = value
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.load_in_4bit, bool):
+ raise TypeError("load_in_4bit must be a boolean")
+
+ if not isinstance(self.load_in_8bit, bool):
+ raise TypeError("load_in_8bit must be a boolean")
+
+ if not isinstance(self.llm_int8_threshold, float):
+ raise TypeError("llm_int8_threshold must be a float")
+
+ if self.llm_int8_skip_modules is not None and not isinstance(self.llm_int8_skip_modules, list):
+ raise TypeError("llm_int8_skip_modules must be a list of strings")
+ if not isinstance(self.llm_int8_enable_fp32_cpu_offload, bool):
+ raise TypeError("llm_int8_enable_fp32_cpu_offload must be a boolean")
+
+ if not isinstance(self.llm_int8_has_fp16_weight, bool):
+ raise TypeError("llm_int8_has_fp16_weight must be a boolean")
+
+ if self.bnb_4bit_compute_dtype is not None and not isinstance(self.bnb_4bit_compute_dtype, torch.dtype):
+ raise TypeError("bnb_4bit_compute_dtype must be torch.dtype")
+
+ if not isinstance(self.bnb_4bit_quant_type, str):
+ raise TypeError("bnb_4bit_quant_type must be a string")
+
+ if not isinstance(self.bnb_4bit_use_double_quant, bool):
+ raise TypeError("bnb_4bit_use_double_quant must be a boolean")
+
+ def is_quantizable(self):
+ r"""
+ Returns `True` if the model is quantizable, `False` otherwise.
+ """
+ return self.load_in_8bit or self.load_in_4bit
+
+ def quantization_method(self):
+ r"""
+ This method returns the quantization method used for the model. If the model is not quantizable, it returns
+ `None`.
+ """
+ if self.load_in_8bit:
+ return "llm_int8"
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "fp4":
+ return "fp4"
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "nf4":
+ return "nf4"
+ else:
+ return None
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ output = copy.deepcopy(self.__dict__)
+ output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1]
+ output["bnb_4bit_quant_storage"] = str(output["bnb_4bit_quant_storage"]).split(".")[1]
+ output["load_in_4bit"] = self.load_in_4bit
+ output["load_in_8bit"] = self.load_in_8bit
+
+ return output
+
+ def __repr__(self):
+ config_dict = self.to_dict()
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
+
+ def to_diff_dict(self) -> dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+
+ Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = BitsAndBytesConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+
+class ExllamaVersion(int, Enum):
+ ONE = 1
+ TWO = 2
+
+
+@dataclass
+class GPTQConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `optimum` api for GPTQ quantization relying on the gptqmodel backend.
+
+ Args:
+ bits (`int`):
+ The number of bits to quantize to, supported numbers are (2, 3, 4, 8).
+ tokenizer (`str` or `PreTrainedTokenizerBase`, *optional*):
+ The tokenizer used to process the dataset. You can pass either:
+ - A custom tokenizer object.
+ - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
+ using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+ dataset (`Union[list[str]]`, *optional*):
+ The dataset used for quantization. You can provide your own dataset in a list of string or just use the
+ original datasets used in GPTQ paper ['wikitext2','c4','c4-new']
+ group_size (`int`, *optional*, defaults to 128):
+ The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
+ damp_percent (`float`, *optional*, defaults to 0.1):
+ The percent of the average Hessian diagonal to use for dampening. Recommended value is 0.1.
+ desc_act (`bool`, *optional*, defaults to `False`):
+ Whether to quantize columns in order of decreasing activation size. Setting it to False can significantly
+ speed up inference but the perplexity may become slightly worse. Also known as act-order.
+ act_group_aware (`bool`, *optional*, defaults to `True`):
+ Use GAR (group aware activation order) during quantization. Has measurable positive impact on quantization
+ quality. Only applicable when `desc_act = False`. Will forced to be `False` when `desc_act = True`.
+ sym (`bool`, *optional*, defaults to `True`):
+ Whether to use symmetric quantization.
+ true_sequential (`bool`, *optional*, defaults to `True`):
+ Whether to perform sequential quantization even within a single Transformer block. Instead of quantizing
+ the entire block at once, we perform layer-wise quantization. As a result, each layer undergoes
+ quantization using inputs that have passed through the previously quantized layers.
+ format (`str`, *optional*, defaults to `"gptq"`):
+ GPTQ weight format. `gptq` (v1) is supported by gptqmodel. `gptq_v2` is gptqmodel only.
+ meta (`dict[str, any]`, *optional*):
+ Properties, such as tooling:version, that do not directly contributes to quantization or quant inference are stored in meta.
+ i.e. `meta.quantizer`: ["optimum:_version_", "gptqmodel:_version_"]
+ backend (`str`, *optional*):
+ Controls which kernel to use. Valid values for gptqmodel are `auto`, `auto_trainable` and more. Ref gptqmodel backends:
+ https://github.com/ModelCloud/GPTQModel/blob/main/gptqmodel/utils/backend.py
+ model_seqlen (`int`, *optional*):
+ The maximum sequence length that the model can take.
+ block_name_to_quantize (`str`, *optional*):
+ The transformers block name to quantize. If None, we will infer the block name using common patterns (e.g. model.layers)
+ module_name_preceding_first_block (`list[str]`, *optional*):
+ The layers that are preceding the first Transformer block.
+ batch_size (`int`, *optional*, defaults to 1):
+ The batch size used when processing the dataset
+ pad_token_id (`int`, *optional*):
+ The pad token id. Needed to prepare the dataset when `batch_size` > 1.
+ max_input_length (`int`, *optional*):
+ The maximum input length. This is needed to initialize a buffer that depends on the maximum expected input
+ length. It is specific to the exllama backend with act-order.
+ cache_block_outputs (`bool`, *optional*, defaults to `True`):
+ Whether to cache block outputs to reuse as inputs for the succeeding block.
+ modules_in_block_to_quantize (`list[list[str]]`, *optional*):
+ List of list of module names to quantize in the specified block. This argument is useful to exclude certain linear modules from being quantized.
+ The block to quantize can be specified by setting `block_name_to_quantize`. We will quantize each list sequentially. If not set, we will quantize all linear layers.
+ Example: `modules_in_block_to_quantize =[["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"], ["self_attn.o_proj"]]`.
+ In this example, we will first quantize the q,k,v layers simultaneously since they are independent.
+ Then, we will quantize `self_attn.o_proj` layer with the q,k,v layers quantized. This way, we will get
+ better results since it reflects the real input `self_attn.o_proj` will get when the model is quantized.
+ """
+
+ def __init__(
+ self,
+ bits: int,
+ tokenizer: Any = None,
+ dataset: list[str] | str | None = None,
+ group_size: int = 128,
+ damp_percent: float = 0.1,
+ desc_act: bool = False,
+ act_group_aware: bool = True,
+ sym: bool = True,
+ true_sequential: bool = True,
+ format: str = "gptq",
+ meta: dict[str, Any] | None = None,
+ backend: str | None = None,
+ model_seqlen: int | None = None,
+ block_name_to_quantize: str | None = None,
+ module_name_preceding_first_block: list[str] | None = None,
+ batch_size: int = 1,
+ pad_token_id: int | None = None,
+ max_input_length: int | None = None,
+ cache_block_outputs: bool = True,
+ modules_in_block_to_quantize: list[list[str]] | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.GPTQ
+ self.bits = bits
+ self.tokenizer = tokenizer
+ self.dataset = dataset
+ self.group_size = group_size
+ self.damp_percent = damp_percent
+ self.desc_act = desc_act
+ self.act_group_aware = act_group_aware
+ self.sym = sym
+ self.true_sequential = true_sequential
+ self.format = format.lower()
+ # Compatible with legacy field: checkpoint_format
+ if kwargs.get("checkpoint_format") is not None:
+ self.format = kwargs.pop("checkpoint_format").lower()
+ self.meta = meta
+ self.backend = backend.lower() if isinstance(backend, str) else backend
+ self.model_seqlen = model_seqlen
+ self.block_name_to_quantize = block_name_to_quantize
+ self.module_name_preceding_first_block = module_name_preceding_first_block
+ self.batch_size = batch_size
+ self.pad_token_id = pad_token_id
+ self.max_input_length = max_input_length
+ self.cache_block_outputs = cache_block_outputs
+ self.modules_in_block_to_quantize = modules_in_block_to_quantize
+ self.post_init()
+
+ def get_loading_attributes(self):
+ attributes_dict = copy.deepcopy(self.__dict__)
+ loading_attributes = ["max_input_length", "backend"]
+ loading_attributes_dict = {i: j for i, j in attributes_dict.items() if i in loading_attributes}
+ return loading_attributes_dict
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ if self.bits not in [2, 3, 4, 8]:
+ raise ValueError(f"Only support quantization to [2,3,4,8] bits but found {self.bits}")
+ if self.group_size != -1 and self.group_size <= 0:
+ raise ValueError("group_size must be greater than 0 or equal to -1")
+ if not (0 < self.damp_percent < 1):
+ raise ValueError("damp_percent must between 0 and 1.")
+ if self.dataset is not None:
+ if isinstance(self.dataset, str):
+ if self.dataset not in ["wikitext2", "c4", "c4-new"]:
+ raise ValueError(
+ f"""You have entered a string value for dataset. You can only choose between
+ ['wikitext2','c4','c4-new'], but we found {self.dataset}"""
+ )
+ elif not isinstance(self.dataset, list):
+ raise ValueError(
+ f"""dataset needs to be either a list of string or a value in
+ ['wikitext2','c4','c4-new'], but we found {self.dataset}"""
+ )
+
+ # act_group_order is only applicable when `desc_act = False`
+ if self.desc_act and self.act_group_aware:
+ self.act_group_aware = False
+ logger.warning("`act_group_aware` has been auto-disabled as it is not compatible with `desc_act = True`.")
+
+ # make sure backend default stays consistent with gptqmodel expectations
+ if self.backend is None:
+ self.backend = "auto"
+ if self.modules_in_block_to_quantize is not None:
+ optimum_version = version.parse(importlib.metadata.version("optimum"))
+ if optimum_version < version.parse("1.15.0"):
+ raise ValueError(
+ "You current version of `optimum` does not support `modules_in_block_to_quantize` quantization argument, please upgrade `optimum` package to a version superior than 1.15.0 ."
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ config_dict = super().to_dict()
+ # Compatible with legacy field: checkpoint_format
+ config_dict["checkpoint_format"] = self.format
+ return config_dict
+
+ def to_dict_optimum(self):
+ """
+ Get compatible dict for optimum gptq config
+ """
+ return self.to_dict()
+
+ @classmethod
+ def from_dict_optimum(cls, config_dict):
+ """
+ Get compatible class with optimum gptq config dict
+ """
+
+ config = cls(**config_dict)
+ return config
+
+
+@dataclass
+class AwqConfig(GPTQConfig):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `auto-awq` library awq quantization relying on auto_awq backend.
+
+ Args:
+ bits (`int`, *optional*, defaults to 4):
+ The number of bits to quantize to.
+ group_size (`int`, *optional*, defaults to 128):
+ The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
+ zero_point (`bool`, *optional*, defaults to `True`):
+ Whether to use zero point quantization.
+ backend (`AwqBackend`, *optional*, defaults to `AwqBackend.AUTO`):
+ The quantization backend.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ Note you cannot quantize directly with transformers, please refer to `AutoAWQ` documentation for quantizing HF models.
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ group_size: int = 128,
+ zero_point: bool = True,
+ backend: AwqBackend = AwqBackend.AUTO,
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ format = kwargs.pop("format", AwqFormat.GEMM)
+ # Compatible with legacy field: version
+ if kwargs.get("version") is not None:
+ format = kwargs.pop("version").lower()
+ # Compatible with legacy backend
+ if backend == AwqBackend.LEGACY_AWQ:
+ backend = AwqBackend.AUTO
+ self.zero_point = zero_point
+ self.modules_to_not_convert = modules_to_not_convert
+
+ super().__init__(bits=bits, group_size=group_size, backend=backend, format=format, **kwargs)
+ self.quant_method = QuantizationMethod.AWQ
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+
+ if self.backend == "llm-awq":
+ self.format = AwqFormat.LLM_AWQ
+ self.backend = AwqBackend.AUTO
+
+ if self.format not in AwqFormat.__members__.values():
+ raise ValueError(f"Invalid format '{self.format}'. Must be one of: {[b.value for b in AwqFormat]}")
+
+ if self.backend not in AwqBackend.__members__.values():
+ raise ValueError(f"Invalid backend '{self.backend}'. Must be one of: {[b.value for b in AwqBackend]}")
+
+ def to_dict(self) -> dict[str, Any]:
+ config_dict = super().to_dict()
+ config_dict.pop("checkpoint_format")
+ # Compatible with legacy field: version
+ config_dict["version"] = self.format
+ return config_dict
+
+
+@dataclass
+class AqlmConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `aqlm` parameters.
+
+ Args:
+ in_group_size (`int`, *optional*, defaults to 8):
+ The group size along the input dimension.
+ out_group_size (`int`, *optional*, defaults to 1):
+ The group size along the output dimension. It's recommended to always use 1.
+ num_codebooks (`int`, *optional*, defaults to 1):
+ Number of codebooks for the Additive Quantization procedure.
+ nbits_per_codebook (`int`, *optional*, defaults to 16):
+ Number of bits encoding a single codebook vector. Codebooks size is 2**nbits_per_codebook.
+ linear_weights_not_to_quantize (`Optional[list[str]]`, *optional*):
+ List of full paths of `nn.Linear` weight parameters that shall not be quantized.
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ in_group_size: int = 8,
+ out_group_size: int = 1,
+ num_codebooks: int = 1,
+ nbits_per_codebook: int = 16,
+ linear_weights_not_to_quantize: list[str] | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.AQLM
+ self.in_group_size = in_group_size
+ self.out_group_size = out_group_size
+ self.num_codebooks = num_codebooks
+ self.nbits_per_codebook = nbits_per_codebook
+ self.linear_weights_not_to_quantize = linear_weights_not_to_quantize
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.in_group_size, int):
+ raise TypeError("in_group_size must be a float")
+ if not isinstance(self.out_group_size, int):
+ raise TypeError("out_group_size must be a float")
+ if not isinstance(self.num_codebooks, int):
+ raise TypeError("num_codebooks must be a float")
+ if not isinstance(self.nbits_per_codebook, int):
+ raise TypeError("nbits_per_codebook must be a float")
+
+ if self.linear_weights_not_to_quantize is not None and not isinstance(
+ self.linear_weights_not_to_quantize, list
+ ):
+ raise ValueError("linear_weights_not_to_quantize must be a list of strings")
+
+ if self.linear_weights_not_to_quantize is None:
+ self.linear_weights_not_to_quantize = []
+
+
+@dataclass
+class VptqLayerConfig(QuantizationConfigMixin):
+ """
+ This is used to explain vptq config params for each layer
+ Args:
+ enable_norm (`bool`, *optional*, defaults to `True`): to control if we have scale/bias for fp-weight
+ enable_perm (`bool`, *optional*, defaults to `True`): to perm input_channel or not
+ group_num (`int`, *optional*, defaults to `1`): how many single groups for vector-quantization
+ group_size (`int`, *optional*, defaults to `-1`): depends on out-features
+ indices_as_float (`bool`, *optional*, defaults to `False`): for Finetuning
+ is_indice_packed (`bool`, *optional*, defaults to `True`): should always be True
+ num_centroids (`list`, *optional*, defaults to `[-1, -1]`): centroid numbers of clusters
+ num_res_centroids (`list`, *optional*, defaults to `[-1, -1]`): ditto for residual
+ outlier_size (`int`, *optional*, defaults to `1`): outliers
+ vector_lens (`list`, *optional*, defaults to `[-1, -1]`): centroid vector length in quantization
+ """
+
+ def __init__(
+ self,
+ enable_norm: bool = True,
+ enable_perm: bool = True,
+ group_num: int = 1,
+ group_size: int = -1,
+ in_features: int = -1,
+ indices_as_float: bool = False,
+ is_indice_packed: bool = True,
+ num_centroids: list = [-1, -1],
+ num_res_centroids: list = [-1, -1],
+ out_features: int = -1,
+ outlier_size: int = 0,
+ vector_lens: list = [-1, -1],
+ **kwargs,
+ ):
+ self.enable_norm = enable_norm
+ self.enable_perm = enable_perm
+ self.group_num = group_num
+ self.group_size = group_size
+ self.in_features = in_features
+ self.indices_as_float = indices_as_float
+ self.is_indice_packed = is_indice_packed
+ self.num_centroids = num_centroids
+ self.num_res_centroids = num_res_centroids
+ self.out_features = out_features
+ self.outlier_size = outlier_size
+ self.vector_lens = vector_lens
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ if self.is_indice_packed is False:
+ raise ValueError("is_indice_packed should always be True")
+
+
+@dataclass
+class VptqConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `vptq` parameters.
+
+ Args:
+ enable_proxy_error (`bool`, *optional*, defaults to `False`): calculate proxy error for each layer
+ config_for_layers (`Dict`, *optional*, defaults to `{}`): quantization params for each layer
+ shared_layer_config (`Dict`, *optional*, defaults to `{}`): shared quantization params among layers
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ enable_proxy_error: bool = False,
+ config_for_layers: dict[str, Any] = {},
+ shared_layer_config: dict[str, Any] = {},
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.VPTQ
+ self.enable_proxy_error = enable_proxy_error
+ self.config_for_layers: dict[str, Any] = config_for_layers
+ self.shared_layer_config: dict[str, Any] = shared_layer_config
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ for layer_param in self.config_for_layers.values():
+ VptqLayerConfig(**layer_param)
+ if self.enable_proxy_error is True:
+ raise ValueError("enable_proxy_error should always be False until we support training")
+
+
+@dataclass
+class QuantoConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `quanto`.
+
+ Args:
+ weights (`str`, *optional*, defaults to `"int8"`):
+ The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2")
+ activations (`str`, *optional*):
+ The target dtype for the activations after quantization. Supported values are (None,"int8","float8")
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ """
+
+ def __init__(
+ self,
+ weights="int8",
+ activations=None,
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.QUANTO
+ self.weights = weights
+ self.activations = activations
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ accepted_weights = ["float8", "int8", "int4", "int2"]
+ accepted_activations = [None, "int8", "float8"]
+ if self.weights not in accepted_weights:
+ raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}")
+ if self.activations not in accepted_activations:
+ raise ValueError(f"Only support weights in {accepted_activations} but found {self.activations}")
+
+
+@dataclass
+class EetqConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `eetq`.
+
+ Args:
+ weights (`str`, *optional*, defaults to `"int8"`):
+ The target dtype for the weights. Supported value is only "int8"
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ """
+
+ def __init__(
+ self,
+ weights: str = "int8",
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.EETQ
+ self.weights = weights
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ accepted_weights = ["int8"]
+ if self.weights not in accepted_weights:
+ raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}")
+
+
+class CompressedTensorsConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class that handles compressed-tensors quantization config options.
+ It is a wrapper around `compressed_tensors.QuantizationConfig`
+ Args:
+ config_groups (`typing.dict[str, typing.Union[ForwardRef('QuantizationScheme'), typing.list[str]]]`, *optional*):
+ dictionary mapping group name to a quantization scheme definition
+ format (`str`, *optional*, defaults to `"dense"`):
+ format the model is represented as. Set `run_compressed` True to execute model as the
+ compressed format if not `dense`
+ quantization_status (`QuantizationStatus`, *optional*, defaults to `"initialized"`):
+ status of model in the quantization lifecycle, ie 'initialized', 'calibration', 'frozen'
+ kv_cache_scheme (`typing.Union[QuantizationArgs, NoneType]`, *optional*):
+ specifies quantization of the kv cache. If None, kv cache is not quantized.
+ global_compression_ratio (`typing.Union[float, NoneType]`, *optional*):
+ 0-1 float percentage of model compression
+ ignore (`typing.Union[typing.list[str], NoneType]`, *optional*):
+ layer names or types to not quantize, supports regex prefixed by 're:'
+ sparsity_config (`typing.dict[str, typing.Any]`, *optional*):
+ configuration for sparsity compression
+ quant_method (`str`, *optional*, defaults to `"compressed-tensors"`):
+ do not override, should be compressed-tensors
+ run_compressed (`bool`, *optional*, defaults to `True`): alter submodules (usually linear) in order to
+ emulate compressed model execution if True, otherwise use default submodule
+ """
+
+ def __init__(
+ self,
+ config_groups: dict[str, Union["QuantizationScheme", list[str]]] | None = None, # noqa: F821
+ format: str = "dense",
+ quantization_status: "QuantizationStatus" = "initialized", # noqa: F821
+ kv_cache_scheme: Optional["QuantizationArgs"] = None, # noqa: F821
+ global_compression_ratio: float | None = None,
+ ignore: list[str] | None = None,
+ sparsity_config: dict[str, Any] | None = None,
+ quant_method: str = "compressed-tensors",
+ run_compressed: bool = True,
+ **kwargs,
+ ):
+ if is_compressed_tensors_available():
+ from compressed_tensors.config import SparsityCompressionConfig
+ from compressed_tensors.quantization import QuantizationConfig
+ else:
+ raise ImportError(
+ "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`."
+ )
+ self.quantization_config = None
+ self.sparsity_config = None
+
+ self.run_compressed = run_compressed
+
+ # parse from dict to load nested QuantizationScheme objects
+ if config_groups or kv_cache_scheme:
+ self.quantization_config = QuantizationConfig.model_validate(
+ {
+ "config_groups": config_groups,
+ "quant_method": quant_method,
+ "format": format,
+ "quantization_status": quantization_status,
+ "kv_cache_scheme": kv_cache_scheme,
+ "global_compression_ratio": global_compression_ratio,
+ "ignore": ignore,
+ **kwargs,
+ }
+ )
+
+ if sparsity_config:
+ self.sparsity_config = SparsityCompressionConfig.load_from_registry(
+ sparsity_config.get("format"), **sparsity_config
+ )
+
+ self.quant_method = QuantizationMethod.COMPRESSED_TENSORS
+
+ def post_init(self):
+ if self.run_compressed:
+ if self.is_sparsification_compressed:
+ logger.warning(
+ "`run_compressed` is only supported for quantized_compressed models"
+ " and not for sparsified models. Setting `run_compressed=False`"
+ )
+ self.run_compressed = False
+ elif not self.is_quantization_compressed:
+ logger.warning(
+ "`run_compressed` is only supported for compressed models. Setting `run_compressed=False`"
+ )
+ self.run_compressed = False
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """
+ Instantiates a [`CompressedTensorsConfig`] from a Python dictionary of parameters.
+ Optionally unwraps any args from the nested quantization_config
+
+ Args:
+ config_dict (`dict[str, Any]`):
+ Dictionary that will be used to instantiate the configuration object.
+ return_unused_kwargs (`bool`,*optional*, defaults to `False`):
+ Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
+ `PreTrainedModel`.
+ kwargs (`dict[str, Any]`):
+ Additional parameters from which to initialize the configuration object.
+
+ Returns:
+ [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters.
+
+ """
+
+ if "quantization_config" in config_dict:
+ config_dict = dict(
+ sparsity_config=config_dict.get("sparsity_config"),
+ **config_dict["quantization_config"],
+ )
+
+ return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Quantization config to be added to config.json
+
+ Serializes this instance to a Python dictionary. Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ quantization_config = {}
+ if self.quantization_config is not None:
+ quantization_config = self.quantization_config.model_dump()
+ else:
+ quantization_config["quant_method"] = QuantizationMethod.COMPRESSED_TENSORS
+
+ if self.sparsity_config is not None:
+ quantization_config["sparsity_config"] = self.sparsity_config.model_dump()
+ else:
+ quantization_config["sparsity_config"] = {}
+
+ return quantization_config
+
+ def to_diff_dict(self) -> dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+ Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = CompressedTensorsConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if key not in default_config_dict or value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+ def get_loading_attributes(self):
+ return {"run_compressed": self.run_compressed}
+
+ @property
+ def is_quantized(self):
+ return bool(self.quantization_config) and bool(self.quantization_config.config_groups)
+
+ @property
+ def is_quantization_compressed(self):
+ from compressed_tensors.quantization import QuantizationStatus
+
+ qc = self.quantization_config
+ return self.is_quantized and (qc is not None and qc.quantization_status == QuantizationStatus.COMPRESSED)
+
+ @property
+ def is_sparsification_compressed(self):
+ from compressed_tensors.config import (
+ CompressionFormat,
+ SparsityCompressionConfig,
+ )
+
+ return (
+ isinstance(self.sparsity_config, SparsityCompressionConfig)
+ and self.sparsity_config.format != CompressionFormat.dense.value
+ )
+
+
+@dataclass
+class FbgemmFp8Config(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using fbgemm fp8 quantization.
+
+ Args:
+ activation_scale_ub (`float`, *optional*, defaults to 1200.0):
+ The activation scale upper bound. This is used when quantizing the input activation.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ """
+
+ def __init__(
+ self,
+ activation_scale_ub: float = 1200.0,
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.FBGEMM_FP8
+ self.activation_scale_ub = activation_scale_ub
+ self.modules_to_not_convert = modules_to_not_convert
+
+ def get_loading_attributes(self):
+ attributes_dict = copy.deepcopy(self.__dict__)
+ loading_attributes = ["activation_scale_ub"]
+ loading_attributes_dict = {i: j for i, j in attributes_dict.items() if i in loading_attributes}
+ return loading_attributes_dict
+
+
+@dataclass
+class HiggsConfig(QuantizationConfigMixin):
+ """
+ HiggsConfig is a configuration class for quantization using the HIGGS method.
+
+ Args:
+ bits (int, *optional*, defaults to 4):
+ Number of bits to use for quantization. Can be 2, 3 or 4. Default is 4.
+ p (int, *optional*, defaults to 2):
+ Quantization grid dimension. 1 and 2 are supported. 2 is always better in practice. Default is 2.
+ modules_to_not_convert (`list`, *optional*, default to ["lm_head"]):
+ List of linear layers that should not be quantized.
+ hadamard_size (int, *optional*, defaults to 512):
+ Hadamard size for the HIGGS method. Default is 512. Input dimension of matrices is padded to this value. Decreasing this below 512 will reduce the quality of the quantization.
+ group_size (int, *optional*, defaults to 256):
+ Group size for the HIGGS method. Can be 64, 128 or 256. Decreasing it barely affects the performance. Default is 256. Must be a divisor of hadamard_size.
+ tune_metadata ('dict', *optional*, defaults to {}):
+ Module-wise metadata (gemm block shapes, GPU metadata, etc.) for saving the kernel tuning results. Default is an empty dictionary. Is set automatically during tuning.
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ p: int = 2,
+ modules_to_not_convert: list[str] | None = None,
+ hadamard_size: int = 512,
+ group_size: int = 256,
+ tune_metadata: dict[str, Any] | None = None,
+ **kwargs,
+ ):
+ if tune_metadata is None:
+ tune_metadata = {}
+ self.quant_method = QuantizationMethod.HIGGS
+ self.bits = bits
+ self.p = p
+ self.modules_to_not_convert = modules_to_not_convert
+ self.hadamard_size = hadamard_size
+ self.group_size = group_size
+ self.tune_metadata = tune_metadata
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if self.bits not in [2, 3, 4]:
+ raise ValueError("bits must be 2, 3, or 4")
+ if self.p not in [1, 2]:
+ raise ValueError("p must be 1 or 2. 2 is always better in practice")
+ if self.group_size not in [64, 128, 256]:
+ raise ValueError("group_size must be 64, 128, or 256")
+ if self.hadamard_size % self.group_size != 0:
+ raise ValueError("hadamard_size must be divisible by group_size")
+
+
+@dataclass
+class FPQuantConfig(QuantizationConfigMixin):
+ """
+ FPQuantConfig is a configuration class for quantization using the FPQuant method.
+
+ Args:
+ forward_dtype (`str`, *optional*, defaults to `"nvfp4"`):
+ The dtype to use for the forward pass.
+ forward_method (`str`, *optional*, defaults to `"abs_max"`):
+ The scaling to use for the forward pass. Can be `"abs_max"` or `"quest"`. `"abs_max"` is better for PTQ, `"quest"` is better for QAT.
+ backward_dtype (`str`, *optional*, defaults to `"bf16"`):
+ The dtype to use for the backward pass.
+ store_master_weights (`bool`, *optional*, defaults to `False`):
+ Whether to store the master weights. Needed for QAT over layer weights.
+ hadamard_group_size (`int`, *optional*):
+ The group size for the hadamard transform before quantization for `"quest"` it matches the MXFP4 group size (32). If `None`, it will be set to 16 for `"nvfp4"` and 32 for `"mxfp4"`.
+ pseudoquantization (`bool`, *optional*, defaults to `False`):
+ Whether to use Triton-based pseudo-quantization. Is mandatory for non-Blackwell GPUs. Doesn't provide any speedup. For debugging purposes.
+ transform_init (`str`, *optional*, defaults to `"hadamard"`): a method to initialize the pre-processing matrix with. Can be `"hadamard"`, `"identity"` or `"gsr"`.
+ modules_to_not_convert (`list`, *optional*):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ """
+
+ def __init__(
+ self,
+ forward_dtype: str = "nvfp4",
+ forward_method: str = "abs_max",
+ backward_dtype: str = "bf16",
+ store_master_weights: bool = False,
+ hadamard_group_size: int | None = None,
+ pseudoquantization: bool = False,
+ transform_init: str = "hadamard",
+ modules_to_not_convert: list[str] | None = None,
+ **kwargs,
+ ):
+ self.forward_dtype = forward_dtype
+ self.forward_method = forward_method
+ self.backward_dtype = backward_dtype
+ self.store_master_weights = store_master_weights
+ self.hadamard_group_size = hadamard_group_size
+ self.pseudoquantization = pseudoquantization
+ self.transform_init = transform_init
+ self.modules_to_not_convert = modules_to_not_convert
+
+ self.quant_method = QuantizationMethod.FPQUANT
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+
+ if self.hadamard_group_size is None:
+ if self.forward_dtype == "nvfp4":
+ self.hadamard_group_size = 16
+ else:
+ self.hadamard_group_size = 32
+
+ if self.forward_dtype == "mxfp4":
+ if self.forward_method not in ["abs_max", "quest"]:
+ raise ValueError("Only 'abs_max' and 'quest' are supported for forward_method for 'mxfp4'.")
+ if self.hadamard_group_size is None:
+ self.hadamard_group_size = 32
+ if self.hadamard_group_size not in [32, 64, 128]:
+ raise ValueError("Only a `hadamard_group_size` of [32, 64, 128] is supported for 'mxfp4'.")
+ elif self.forward_dtype == "nvfp4":
+ if self.forward_method != "abs_max":
+ raise ValueError("Only 'abs_max' is supported for forward_method for 'nvfp4'.")
+ if self.hadamard_group_size is None:
+ self.hadamard_group_size = 16
+ if self.hadamard_group_size not in [16, 32, 64, 128]:
+ raise ValueError("Only a `hadamard_group_size` of [16, 32, 64, 128] is supported for 'nvfp4'.")
+ else:
+ raise ValueError("Only 'mxfp4' and 'nvfp4' are supported for forward_dtype for now.")
+
+ if self.backward_dtype not in ["bf16", "mxfp8", "mxfp4"]:
+ raise ValueError("Only 'bf16', 'mxfp8' and 'mxfp4' are supported for backward_dtype for now.")
+
+ if self.backward_dtype != "bf16" and self.forward_dtype != "mxfp4":
+ raise ValueError("Only 'mxfp4' forward is compatible with non-bf16 backwards for now.")
+
+ if self.transform_init not in ["hadamard", "identity", "gsr"]:
+ raise ValueError("Only 'hadamard', 'identity' and 'gsr' are supported for transform_init.")
+
+ if self.modules_to_not_convert is None:
+ self.modules_to_not_convert = ["lm_head"]
+
+
+@dataclass
+class TorchAoConfig(QuantizationConfigMixin):
+ """Config class for torchao quantization/sparsity techniques.
+
+ Args:
+ quant_type (`AOBaseConfig`):
+ A torchao `AOBaseConfig` instance specifying the quantization type, e.g.
+ `Int4WeightOnlyConfig(group_size=32)`, `Int8WeightOnlyConfig()`,
+ `Int8DynamicActivationInt8WeightConfig()`, `Float8WeightOnlyConfig()`, etc.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ include_input_output_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether to include embedding in quantization or not, input embedding will be removed from
+ the module_not_to_convert list as well if this flag is set.
+ untie_embedding_weights (`bool`, *optional*, defaults to `False`):
+ Whether to untie the weights when we are quantizing input embedding weights that is tied
+ to other weights.
+
+ Example:
+
+ ```python
+ from torchao.quantization import Int4WeightOnlyConfig
+
+ quantization_config = TorchAoConfig(Int4WeightOnlyConfig(group_size=32))
+ model = AutoModelForCausalLM.from_pretrained(
+ model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config
+ )
+ ```
+ """
+
+ quant_method: QuantizationMethod
+ quant_type: "AOBaseConfig" # noqa: F821
+ modules_to_not_convert: list | None
+ include_input_output_embeddings: bool
+ untie_embedding_weights: bool
+
+ def __init__(
+ self,
+ quant_type: "AOBaseConfig", # noqa: F821
+ modules_to_not_convert: list | None = None,
+ include_input_output_embeddings: bool = False,
+ untie_embedding_weights: bool = False,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.TORCHAO
+ self.quant_type = quant_type
+ self.modules_to_not_convert = modules_to_not_convert
+ self.include_input_output_embeddings = include_input_output_embeddings
+ self.untie_embedding_weights = untie_embedding_weights
+ self.post_init()
+
+ def post_init(self):
+ """Validate configuration and set defaults."""
+ if not is_torchao_available():
+ raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")
+
+ if isinstance(self.quant_type, str):
+ raise ValueError(
+ f"String-based quantization type '{self.quant_type}' is no longer supported. "
+ f"Please use the corresponding Config object directly, e.g. "
+ f"TorchAoConfig(Int4WeightOnlyConfig(group_size=32)) instead of "
+ f"TorchAoConfig('int4_weight_only', group_size=32)."
+ )
+
+ from torchao.quantization.quant_api import AOBaseConfig
+
+ if not isinstance(self.quant_type, AOBaseConfig):
+ raise TypeError(f"quant_type must be an AOBaseConfig instance, got {type(self.quant_type)}")
+
+ def get_apply_tensor_subclass(self):
+ """Return the quantization config to apply."""
+ return self.quant_type
+
+ def to_dict(self):
+ """Convert configuration to a dictionary."""
+ d = super().to_dict()
+
+ from torchao.core.config import config_to_dict
+
+ d["quant_type"] = {"default": config_to_dict(self.quant_type)}
+
+ return d
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """Create configuration from a dictionary."""
+ from torchao.core.config import config_from_dict
+
+ config_dict = config_dict.copy()
+ quant_type = config_dict.pop("quant_type")
+
+ # Check if we only have one key which is "default"
+ # In the future we may update this
+ assert len(quant_type) == 1 and "default" in quant_type, (
+ "Expected only one key 'default' in quant_type dictionary"
+ )
+ quant_type = quant_type["default"]
+ quant_type = config_from_dict(quant_type)
+
+ return cls(quant_type=quant_type, **config_dict)
+
+
+@dataclass
+class BitNetQuantConfig(QuantizationConfigMixin):
+ """
+ Configuration class for applying BitNet quantization.
+
+ Args:
+ modules_to_not_convert (`Optional[List]`, *optional*):
+ Optionally, provides a list of full paths of `nn.Linear` weight parameters
+ that shall not be quantized. Defaults to None.
+ linear_class (`str`, *optional*, defaults to `"bitlinear"`):
+ The type of linear class to use. Can be either `bitlinear` or `autobitlinear`.
+ quantization_mode (`str`, *optional*, defaults to `"offline"`):
+ The quantization mode to use. Can be either `online` or `offline`.
+ In `online` mode, the weight quantization parameters are calculated dynamically
+ during each forward pass (e.g., based on the current weight values). This can
+ adapt to weight changes during training (Quantization-Aware Training - QAT).
+ In `offline` mode, quantization parameters are pre-calculated *before* inference.
+ These parameters are then fixed and loaded into the quantized model. This
+ generally results in lower runtime overhead compared to online quantization.
+ use_rms_norm (`bool`, *optional*, defaults to `False`):
+ Whether to apply RMSNorm on the activations before quantization. This matches the original BitNet paper's approach
+ of normalizing activations before quantization/packing.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon value used in the RMSNorm layer for numerical stability.
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional keyword arguments that may be used by specific quantization
+ backends or future versions.
+ """
+
+ def __init__(
+ self,
+ modules_to_not_convert: list | None = None,
+ linear_class: str = "bitlinear",
+ quantization_mode: str = "offline",
+ use_rms_norm: bool = False,
+ rms_norm_eps: float | None = 1e-6,
+ **kwargs,
+ ):
+ if linear_class not in ["bitlinear", "autobitlinear"]:
+ raise ValueError(f"linear_class must be either 'bitlinear' or 'autobitlinear', but got {linear_class}")
+ if quantization_mode not in ["online", "offline"]:
+ raise ValueError(f"quantization_mode must be either 'online' or 'offline', but got {quantization_mode}")
+ self.quant_method = QuantizationMethod.BITNET
+ self.modules_to_not_convert = modules_to_not_convert
+ self.linear_class = linear_class
+ self.quantization_mode = quantization_mode
+ self.use_rms_norm = use_rms_norm
+ self.rms_norm_eps = rms_norm_eps
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+
+
+@dataclass
+class SpQRConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `spqr` parameters. Refer to the original publication for more details.
+
+ Args:
+ bits (`int`, *optional*, defaults to 3):
+ Specifies the bit count for the weights and first order zero-points and scales.
+ Currently only bits = 3 is supported.
+ beta1 (`int`, *optional*, defaults to 16):
+ SpQR tile width. Currently only beta1 = 16 is supported.
+ beta2 (`int`, *optional*, defaults to 16):
+ SpQR tile height. Currently only beta2 = 16 is supported.
+ shapes (`Optional`, *optional*):
+ A dictionary holding the shape of each object. We need this because it's impossible
+ to deduce the exact size of the parameters just from bits, beta1, beta2.
+ modules_to_not_convert (`Optional[list[str]]`, *optional*):
+ Optionally, provides a list of full paths of `nn.Linear` weight parameters that shall not be quantized.
+ Defaults to None.
+ kwargs (`dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ bits: int = 3,
+ beta1: int = 16,
+ beta2: int = 16,
+ shapes: dict[str, int] | None = None,
+ modules_to_not_convert: list[str] | None = None,
+ **kwargs,
+ ):
+ if shapes is None:
+ shapes = {}
+ self.shapes = shapes
+ self.quant_method = QuantizationMethod.SPQR
+ self.bits = bits
+ self.beta1 = beta1
+ self.beta2 = beta2
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.bits, int):
+ raise TypeError("bits must be an int")
+ if not isinstance(self.beta1, int):
+ raise TypeError("beta1 must be an int")
+ if not isinstance(self.beta2, int):
+ raise TypeError("beta2 must be an int")
+
+ if self.bits != 3:
+ raise ValueError("SpQR currently only supports bits = 3")
+ if self.beta1 != 16:
+ raise ValueError("SpQR currently only supports beta1 = 16")
+ if self.beta2 != 16:
+ raise ValueError("SpQR currently only supports beta2 = 16")
+ if not isinstance(self.shapes, dict):
+ raise TypeError("shapes must be a dict")
+
+
+@dataclass
+class FineGrainedFP8Config(QuantizationConfigMixin):
+ """
+ FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models.
+
+ Args:
+ activation_scheme (`str`, *optional*, defaults to `"dynamic"`):
+ The scheme used for activation, the defaults and only support scheme for now is "dynamic".
+ weight_block_size (`typing.tuple[int, int]`, *optional*, defaults to `(128, 128)`):
+ The size of the weight blocks for quantization, default is (128, 128).
+ dequantize (`bool`, *optional*, defaults to `False`):
+ Whether to dequantize the model during loading.
+ modules_to_not_convert (`list`, *optional*):
+ A list of module names that should not be converted during quantization.
+ """
+
+ def __init__(
+ self,
+ activation_scheme: str = "dynamic",
+ weight_block_size: tuple[int, int] = (128, 128),
+ dequantize: bool = False,
+ modules_to_not_convert: list | None = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.FP8
+ self.modules_to_not_convert = modules_to_not_convert
+ self.activation_scheme = activation_scheme
+ self.weight_block_size = weight_block_size
+ self.dequantize = dequantize
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ self.activation_scheme = self.activation_scheme.lower()
+ if self.activation_scheme not in ["dynamic", "static"]:
+ raise ValueError(f"Activation scheme {self.activation_scheme} not supported")
+ if self.weight_block_size is not None and len(self.weight_block_size) != 2:
+ raise ValueError("weight_block_size must be a tuple of two integers")
+ if self.weight_block_size is not None and (self.weight_block_size[0] <= 0 or self.weight_block_size[1] <= 0):
+ raise ValueError("weight_block_size must be a tuple of two positive integers")
+
+ def get_loading_attributes(self):
+ return {"dequantize": self.dequantize}
+
+
+class QuarkConfig(QuantizationConfigMixin):
+ def __init__(
+ self,
+ **kwargs,
+ ):
+ if is_torch_available() and is_quark_available():
+ from quark import __version__ as quark_version
+ from quark.torch.export.config.config import JsonExporterConfig
+ from quark.torch.export.main_export.quant_config_parser import QuantConfigParser
+ from quark.torch.quantization.config.config import Config
+ else:
+ raise ImportError(
+ "Quark is not installed. Please refer to https://quark.docs.amd.com/latest/install.html."
+ )
+ # This might be e.g. `"fp8"` or `"awq"`.
+ self.custom_mode = kwargs["quant_method"]
+ self.legacy = "export" not in kwargs
+
+ if self.custom_mode in ["awq", "fp8"]:
+ # Legacy (quark<1.0) or custom export.
+ self.quant_config = QuantConfigParser.from_custom_config(kwargs, is_bias_quantized=False)
+ self.json_export_config = JsonExporterConfig()
+ else:
+ self.quant_config = Config.from_dict(kwargs)
+
+ if "export" in kwargs:
+ # TODO: Remove this check once configuration version is handled natively by Quark.
+ if "min_kv_scale" in kwargs["export"] and version.parse(quark_version) < version.parse("0.8"):
+ min_kv_scale = kwargs["export"].pop("min_kv_scale")
+ logger.warning(
+ f"The parameter `min_kv_scale={min_kv_scale}` was found in the model config.json's `quantization_config.export` configuration, but this parameter is supported only for quark>=0.8. Ignoring this configuration parameter. Please update the `amd-quark` package."
+ )
+
+ self.json_export_config = JsonExporterConfig(**kwargs["export"])
+ else:
+ # Legacy (quark<1.0) or custom export.
+ self.json_export_config = JsonExporterConfig()
+
+ self.quant_method = QuantizationMethod.QUARK
+
+
+@dataclass
+class Mxfp4Config(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using mxfp4 quantization.
+
+ Args:
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ dequantize (`bool`, *optional*, default to `False`):
+ Whether we dequantize the model to bf16 precision or not
+ """
+
+ def __init__(
+ self,
+ modules_to_not_convert: list | None = None,
+ dequantize: bool = False,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.MXFP4
+ self.modules_to_not_convert = modules_to_not_convert
+ self.dequantize = dequantize
+
+ def get_loading_attributes(self):
+ return {"dequantize": self.dequantize}
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ return {"quant_method": self.quant_method, "modules_to_not_convert": self.modules_to_not_convert}
+
+
+class MetalConfig(QuantizationConfigMixin):
+ """
+ Configuration class for Metal affine quantization targeting Apple Silicon (MPS) devices.
+
+ This quantization method uses the ``mlx-quantization-metal-kernels`` Metal kernels from the Hugging Face Hub
+ to perform affine quantization (scales + qbiases) with configurable bit-width and group size.
+ The quantized weights are packed into ``uint32`` tensors and the forward pass uses fused
+ dequantization + matmul Metal kernels.
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ group_size: int = 64,
+ modules_to_not_convert: list | None = None,
+ dequantize: bool = False,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.METAL
+ self.bits = bits
+ self.group_size = group_size
+ self.modules_to_not_convert = modules_to_not_convert
+ self.dequantize = dequantize
+ self.post_init()
+
+ def post_init(self):
+ if self.bits not in (2, 4, 8):
+ raise ValueError(f"Metal quantization only supports bits in {{2, 4, 8}}, got {self.bits}")
+ if self.group_size <= 0:
+ raise ValueError(f"group_size must be positive, got {self.group_size}")
+
+ def get_loading_attributes(self):
+ return {"dequantize": self.dequantize}
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "quant_method": self.quant_method,
+ "bits": self.bits,
+ "group_size": self.group_size,
+ "modules_to_not_convert": self.modules_to_not_convert,
+ }
+
+
+@dataclass
+class FourOverSixConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class containing all options for quantization with `fouroversix`. In brief,
+ Four Over Six is a modification to NVFP4 quantization which adaptively scales the largest value
+ in each block of 16 FP4 values to either 4 or 6. Selecting a scale of 6 uses the full range of
+ FP4 values, but selecting a scale of 4 allows for a more uniform distribution of quantization
+ error. Refer to the original publication for more details: https://arxiv.org/abs/2512.02010.
+
+ Args:
+ activation_scale_rule (`str`, *optional*):
+ Scaling rule to use when selecting a scale for blocks in activation tensors. If not
+ provided, `scale_rule` is used.
+ dtype (`str`, default "nvfp4", *optional*, defaults to `"nvfp4"`):
+ The data type to use for the layer's weights, activations, and tensors. Can be
+ `"nvfp4"` or `"mxfp4"`.
+ gradient_scale_rule (`str`, *optional*):
+ Scaling rule to use when selecting a scale for blocks in gradient tensors. If not
+ provided, `scale_rule` is used.
+ keep_master_weights (`bool`, default False, *optional*, defaults to `False`):
+ Whether to keep the master weights. If `True`, high-precision weights are kept at all
+ times and weights are quantized online in each forward pass. This is useful for
+ quantized training.
+ matmul_backend (`str`, *optional*):
+ The backend to use for matrix multiplications. Can be `"cutlass"` or `"pytorch"`. If
+ not provided, CUTLASS will be used if available and PyTorch will be used otherwise.
+ output_dtype (`str`, *optional*, defaults to `"bfloat16"`):
+ The data type to use for the output of the layer. Can be `"bfloat16"` or `"float16"`.
+ quantize_backend (`str`, *optional*):
+ The backend to use for quantization. Can be `"cuda"`, `"triton"`, or `"pytorch"`. If
+ not provided, the fastest backend will be selected based on your environment, and based
+ on the options supported by each backend. Typically, `"cuda"` will be used for
+ inference, `"triton"` will be used for training, and `"pytorch"` will be used on
+ non-CUDA devices.
+ scale_rule (`str`, default "mse", *optional*, defaults to `"mse"`):
+ Rule to use when selecting block scales. Can be `"mse"`, `"mae"`, or `"abs_max"` for
+ Four Over Six, `"static_6"` for default NVFP4 quantization, or `"static_4"` to scale
+ all blocks to a maximum value of 4.
+ weight_scale_2d (`bool`, default False, *optional*, defaults to `False`):
+ Whether to compute scale factors on weight tensors in 2D blocks. This should be done
+ during training.
+ weight_scale_rule (`str`, *optional*):
+ Scaling rule to use when selecting a scale for blocks in weight tensors. If not
+ provided, `scale_rule` is used.
+ module_config_overrides (`dict[str, dict[str, Any]]`, *optional*):
+ A dictionary of module-specific configuration overrides. Keys should be module names, and
+ values should be dictionaries containing the quantization configuration for that module.
+ This can be used to override the default configuration for specific modules.
+ modules_to_not_convert (`list[str]`, *optional*, defaults to `['lm_head']`):
+ The list of modules to exclude from quantization. By default, the `lm_head` is excluded.
+ """
+
+ def __init__(
+ self,
+ activation_scale_rule: str | None = None,
+ dtype: str = "nvfp4",
+ gradient_scale_rule: str | None = None,
+ keep_master_weights: bool = False,
+ matmul_backend: str | None = None,
+ output_dtype: str | None = "bfloat16",
+ quantize_backend: str | None = None,
+ scale_rule: str = "mse",
+ weight_scale_2d: bool = False,
+ weight_scale_rule: str | None = None,
+ module_config_overrides: dict[str, dict[str, Any]] | None = None,
+ modules_to_not_convert: list[str] | None = ["lm_head"],
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.FOUR_OVER_SIX
+
+ self.activation_scale_rule = activation_scale_rule
+ self.dtype = dtype
+ self.gradient_scale_rule = gradient_scale_rule
+ self.keep_master_weights = keep_master_weights
+ self.matmul_backend = matmul_backend
+ self.quantize_backend = quantize_backend
+ self.output_dtype = output_dtype
+ self.scale_rule = scale_rule
+ self.weight_scale_2d = weight_scale_2d
+ self.weight_scale_rule = weight_scale_rule
+ self.module_config_overrides = module_config_overrides
+ self.modules_to_not_convert = modules_to_not_convert
+
+
+class SinqConfig(QuantizationConfigMixin):
+ """
+ Quantization config for SINQ / A-SINQ.
+
+ Pass this to:
+
+ AutoModel.from_pretrained(..., quantization_config=SinqConfig(...))
+
+ Args:
+ nbits (`int`, default 4):
+ Quantization bits for weights.
+ group_size (`int`, default 64):
+ Group size used in SINQ weight quantization (must be multiple of 8).
+ tiling_mode (`str`, default "1D"):
+ Tiling mode for SINQ (typically "1D"; "2D" if supported in your backend).
+ method (`str`, default "sinq"):
+ "sinq" โ calibration-free weight-only SINQ
+ "asinq" โ A-SINQ (activation-aware), not supported in Hugging Face. Please refer to the official SINQ repository.
+ modules_to_not_convert (`list[str]`, *optional*):
+ List of module names/prefixes to keep in full precision.
+
+ **kwargs:
+ Extra user arguments (kept in `_extra_kwargs` for round-tripping).
+ """
+
+ def __init__(
+ self,
+ nbits: int = 4,
+ group_size: int = 64,
+ tiling_mode: str = "1D",
+ method: str = "sinq", # "sinq" | "asinq"
+ modules_to_not_convert: list[str] | None = None,
+ **kwargs: Any,
+ ):
+ self.quant_method = QuantizationMethod.SINQ
+
+ self.nbits = nbits
+ self.group_size = group_size
+ self.tiling_mode = tiling_mode
+ self.method = method
+
+ self.modules_to_not_convert = modules_to_not_convert
+
+ self._extra_kwargs: dict[str, Any] = dict(kwargs)
+
+ self.post_init()
+
+ def post_init(self):
+ self.nbits = int(self.nbits)
+ self.group_size = int(self.group_size)
+ self.tiling_mode = str(self.tiling_mode)
+ self.method = str(self.method).lower()
+
+ # Validation
+ if not isinstance(self.nbits, int):
+ raise TypeError("`nbits` must be convertible to an int")
+ if not isinstance(self.group_size, int):
+ raise TypeError("`group_size` must be convertible to an int")
+ if not isinstance(self.tiling_mode, str):
+ raise TypeError("`tiling_mode` must be convertible to a string")
+ if self.method not in {"sinq", "asinq"}:
+ raise ValueError(f"`method` must be either 'sinq' or 'asinq', got {self.method}")
+ if self.group_size is not None and self.group_size % 8 != 0:
+ logger.warning(
+ f"SINQ: group_size={self.group_size} is not a multiple of 8; this may be rejected by the backend."
+ )
diff --git a/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2.py b/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f063575fd7af2ae65d3d55b054b856ca8b77fe8
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2.py
@@ -0,0 +1,1511 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# source: sentencepiece_model.proto
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+ name="sentencepiece_model.proto",
+ package="sentencepiece",
+ syntax="proto2",
+ serialized_options=b"H\003",
+ create_key=_descriptor._internal_create_key,
+ serialized_pb=(
+ b'\n\x19sentencepiece_model.proto\x12\rsentencepiece"\xa1\n\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01'
+ b" \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02"
+ b" \x01(\t\x12\x41\n\nmodel_type\x18\x03"
+ b" \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04"
+ b" \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12"
+ b' \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12"\n\x12\x63haracter_coverage\x18\n'
+ b" \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b"
+ b" \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12"
+ b' \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12"\n\x16training_sentence_size\x18\r'
+ b" \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e"
+ b" \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f"
+ b" \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12"
+ b" \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10"
+ b" \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11"
+ b" \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14"
+ b" \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15"
+ b" \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17"
+ b" \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16"
+ b" \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18"
+ b" \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19"
+ b" \x01(\x08:\x05\x66\x61lse\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e"
+ b" \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$"
+ b" \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18"
+ b' \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18"'
+ b" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18)"
+ b" \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+"
+ b" \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05\x12\x16\n\tbos_piece\x18."
+ b" \x01(\t:\x03\x12\x17\n\teos_piece\x18/ \x01(\t:\x04\x12\x18\n\tpad_piece\x18\x30"
+ b" \x01(\t:\x05\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87"
+ b" \x12+\n\x1ctrain_extremely_large_corpus\x18\x31"
+ b' \x01(\x08:\x05\x66\x61lse"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01'
+ b" \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03"
+ b" \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12"
+ b" \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06"
+ b' \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01'
+ b' \x03(\x0b\x32".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01'
+ b" \x01(\t\x12\x10\n\x08\x65xpected\x18\x02"
+ b' \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01'
+ b" \x03(\x0b\x32'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02"
+ b" \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03"
+ b" \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04"
+ b" \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05"
+ b" \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01"
+ b" \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03"
+ b' \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03'
+ ),
+)
+
+
+_TRAINERSPEC_MODELTYPE = _descriptor.EnumDescriptor(
+ name="ModelType",
+ full_name="sentencepiece.TrainerSpec.ModelType",
+ filename=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ values=[
+ _descriptor.EnumValueDescriptor(
+ name="UNIGRAM",
+ index=0,
+ number=1,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="BPE",
+ index=1,
+ number=2,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="WORD",
+ index=2,
+ number=3,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="CHAR",
+ index=3,
+ number=4,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ containing_type=None,
+ serialized_options=None,
+ serialized_start=1294,
+ serialized_end=1347,
+)
+_sym_db.RegisterEnumDescriptor(_TRAINERSPEC_MODELTYPE)
+
+_MODELPROTO_SENTENCEPIECE_TYPE = _descriptor.EnumDescriptor(
+ name="Type",
+ full_name="sentencepiece.ModelProto.SentencePiece.Type",
+ filename=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ values=[
+ _descriptor.EnumValueDescriptor(
+ name="NORMAL",
+ index=0,
+ number=1,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="UNKNOWN",
+ index=1,
+ number=2,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="CONTROL",
+ index=2,
+ number=3,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="USER_DEFINED",
+ index=3,
+ number=4,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="BYTE",
+ index=4,
+ number=6,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.EnumValueDescriptor(
+ name="UNUSED",
+ index=5,
+ number=5,
+ serialized_options=None,
+ type=None,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ containing_type=None,
+ serialized_options=None,
+ serialized_start=2100,
+ serialized_end=2184,
+)
+_sym_db.RegisterEnumDescriptor(_MODELPROTO_SENTENCEPIECE_TYPE)
+
+
+_TRAINERSPEC = _descriptor.Descriptor(
+ name="TrainerSpec",
+ full_name="sentencepiece.TrainerSpec",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="input",
+ full_name="sentencepiece.TrainerSpec.input",
+ index=0,
+ number=1,
+ type=9,
+ cpp_type=9,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="input_format",
+ full_name="sentencepiece.TrainerSpec.input_format",
+ index=1,
+ number=7,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="model_prefix",
+ full_name="sentencepiece.TrainerSpec.model_prefix",
+ index=2,
+ number=2,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="model_type",
+ full_name="sentencepiece.TrainerSpec.model_type",
+ index=3,
+ number=3,
+ type=14,
+ cpp_type=8,
+ label=1,
+ has_default_value=True,
+ default_value=1,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="vocab_size",
+ full_name="sentencepiece.TrainerSpec.vocab_size",
+ index=4,
+ number=4,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=8000,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="accept_language",
+ full_name="sentencepiece.TrainerSpec.accept_language",
+ index=5,
+ number=5,
+ type=9,
+ cpp_type=9,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="self_test_sample_size",
+ full_name="sentencepiece.TrainerSpec.self_test_sample_size",
+ index=6,
+ number=6,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=0,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="character_coverage",
+ full_name="sentencepiece.TrainerSpec.character_coverage",
+ index=7,
+ number=10,
+ type=2,
+ cpp_type=6,
+ label=1,
+ has_default_value=True,
+ default_value=0.9995,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="input_sentence_size",
+ full_name="sentencepiece.TrainerSpec.input_sentence_size",
+ index=8,
+ number=11,
+ type=4,
+ cpp_type=4,
+ label=1,
+ has_default_value=True,
+ default_value=0,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="shuffle_input_sentence",
+ full_name="sentencepiece.TrainerSpec.shuffle_input_sentence",
+ index=9,
+ number=19,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="mining_sentence_size",
+ full_name="sentencepiece.TrainerSpec.mining_sentence_size",
+ index=10,
+ number=12,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=False,
+ default_value=0,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=b"\030\001",
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="training_sentence_size",
+ full_name="sentencepiece.TrainerSpec.training_sentence_size",
+ index=11,
+ number=13,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=False,
+ default_value=0,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=b"\030\001",
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="seed_sentencepiece_size",
+ full_name="sentencepiece.TrainerSpec.seed_sentencepiece_size",
+ index=12,
+ number=14,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=1000000,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="shrinking_factor",
+ full_name="sentencepiece.TrainerSpec.shrinking_factor",
+ index=13,
+ number=15,
+ type=2,
+ cpp_type=6,
+ label=1,
+ has_default_value=True,
+ default_value=0.75,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="max_sentence_length",
+ full_name="sentencepiece.TrainerSpec.max_sentence_length",
+ index=14,
+ number=18,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=4192,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="num_threads",
+ full_name="sentencepiece.TrainerSpec.num_threads",
+ index=15,
+ number=16,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=16,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="num_sub_iterations",
+ full_name="sentencepiece.TrainerSpec.num_sub_iterations",
+ index=16,
+ number=17,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=2,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="max_sentencepiece_length",
+ full_name="sentencepiece.TrainerSpec.max_sentencepiece_length",
+ index=17,
+ number=20,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=16,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="split_by_unicode_script",
+ full_name="sentencepiece.TrainerSpec.split_by_unicode_script",
+ index=18,
+ number=21,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="split_by_number",
+ full_name="sentencepiece.TrainerSpec.split_by_number",
+ index=19,
+ number=23,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="split_by_whitespace",
+ full_name="sentencepiece.TrainerSpec.split_by_whitespace",
+ index=20,
+ number=22,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="treat_whitespace_as_suffix",
+ full_name="sentencepiece.TrainerSpec.treat_whitespace_as_suffix",
+ index=21,
+ number=24,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=False,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="split_digits",
+ full_name="sentencepiece.TrainerSpec.split_digits",
+ index=22,
+ number=25,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=False,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="control_symbols",
+ full_name="sentencepiece.TrainerSpec.control_symbols",
+ index=23,
+ number=30,
+ type=9,
+ cpp_type=9,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="user_defined_symbols",
+ full_name="sentencepiece.TrainerSpec.user_defined_symbols",
+ index=24,
+ number=31,
+ type=9,
+ cpp_type=9,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="required_chars",
+ full_name="sentencepiece.TrainerSpec.required_chars",
+ index=25,
+ number=36,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="byte_fallback",
+ full_name="sentencepiece.TrainerSpec.byte_fallback",
+ index=26,
+ number=35,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=False,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="vocabulary_output_piece_score",
+ full_name="sentencepiece.TrainerSpec.vocabulary_output_piece_score",
+ index=27,
+ number=32,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="hard_vocab_limit",
+ full_name="sentencepiece.TrainerSpec.hard_vocab_limit",
+ index=28,
+ number=33,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="use_all_vocab",
+ full_name="sentencepiece.TrainerSpec.use_all_vocab",
+ index=29,
+ number=34,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=False,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="unk_id",
+ full_name="sentencepiece.TrainerSpec.unk_id",
+ index=30,
+ number=40,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=0,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="bos_id",
+ full_name="sentencepiece.TrainerSpec.bos_id",
+ index=31,
+ number=41,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=1,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="eos_id",
+ full_name="sentencepiece.TrainerSpec.eos_id",
+ index=32,
+ number=42,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=2,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="pad_id",
+ full_name="sentencepiece.TrainerSpec.pad_id",
+ index=33,
+ number=43,
+ type=5,
+ cpp_type=1,
+ label=1,
+ has_default_value=True,
+ default_value=-1,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="unk_piece",
+ full_name="sentencepiece.TrainerSpec.unk_piece",
+ index=34,
+ number=45,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=True,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="bos_piece",
+ full_name="sentencepiece.TrainerSpec.bos_piece",
+ index=35,
+ number=46,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=True,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="eos_piece",
+ full_name="sentencepiece.TrainerSpec.eos_piece",
+ index=36,
+ number=47,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=True,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="pad_piece",
+ full_name="sentencepiece.TrainerSpec.pad_piece",
+ index=37,
+ number=48,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=True,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="unk_surface",
+ full_name="sentencepiece.TrainerSpec.unk_surface",
+ index=38,
+ number=44,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=True,
+ default_value=b" \342\201\207 ".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="train_extremely_large_corpus",
+ full_name="sentencepiece.TrainerSpec.train_extremely_large_corpus",
+ index=39,
+ number=49,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=False,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[],
+ enum_types=[
+ _TRAINERSPEC_MODELTYPE,
+ ],
+ serialized_options=None,
+ is_extendable=True,
+ syntax="proto2",
+ extension_ranges=[
+ (200, 536870912),
+ ],
+ oneofs=[],
+ serialized_start=45,
+ serialized_end=1358,
+)
+
+
+_NORMALIZERSPEC = _descriptor.Descriptor(
+ name="NormalizerSpec",
+ full_name="sentencepiece.NormalizerSpec",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="name",
+ full_name="sentencepiece.NormalizerSpec.name",
+ index=0,
+ number=1,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="precompiled_charsmap",
+ full_name="sentencepiece.NormalizerSpec.precompiled_charsmap",
+ index=1,
+ number=2,
+ type=12,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"",
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="add_dummy_prefix",
+ full_name="sentencepiece.NormalizerSpec.add_dummy_prefix",
+ index=2,
+ number=3,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="remove_extra_whitespaces",
+ full_name="sentencepiece.NormalizerSpec.remove_extra_whitespaces",
+ index=3,
+ number=4,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="escape_whitespaces",
+ full_name="sentencepiece.NormalizerSpec.escape_whitespaces",
+ index=4,
+ number=5,
+ type=8,
+ cpp_type=7,
+ label=1,
+ has_default_value=True,
+ default_value=True,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="normalization_rule_tsv",
+ full_name="sentencepiece.NormalizerSpec.normalization_rule_tsv",
+ index=5,
+ number=6,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[],
+ enum_types=[],
+ serialized_options=None,
+ is_extendable=True,
+ syntax="proto2",
+ extension_ranges=[
+ (200, 536870912),
+ ],
+ oneofs=[],
+ serialized_start=1361,
+ serialized_end=1570,
+)
+
+
+_SELFTESTDATA_SAMPLE = _descriptor.Descriptor(
+ name="Sample",
+ full_name="sentencepiece.SelfTestData.Sample",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="input",
+ full_name="sentencepiece.SelfTestData.Sample.input",
+ index=0,
+ number=1,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="expected",
+ full_name="sentencepiece.SelfTestData.Sample.expected",
+ index=1,
+ number=2,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[],
+ enum_types=[],
+ serialized_options=None,
+ is_extendable=False,
+ syntax="proto2",
+ extension_ranges=[],
+ oneofs=[],
+ serialized_start=1641,
+ serialized_end=1682,
+)
+
+_SELFTESTDATA = _descriptor.Descriptor(
+ name="SelfTestData",
+ full_name="sentencepiece.SelfTestData",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="samples",
+ full_name="sentencepiece.SelfTestData.samples",
+ index=0,
+ number=1,
+ type=11,
+ cpp_type=10,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[
+ _SELFTESTDATA_SAMPLE,
+ ],
+ enum_types=[],
+ serialized_options=None,
+ is_extendable=True,
+ syntax="proto2",
+ extension_ranges=[
+ (200, 536870912),
+ ],
+ oneofs=[],
+ serialized_start=1572,
+ serialized_end=1693,
+)
+
+
+_MODELPROTO_SENTENCEPIECE = _descriptor.Descriptor(
+ name="SentencePiece",
+ full_name="sentencepiece.ModelProto.SentencePiece",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="piece",
+ full_name="sentencepiece.ModelProto.SentencePiece.piece",
+ index=0,
+ number=1,
+ type=9,
+ cpp_type=9,
+ label=1,
+ has_default_value=False,
+ default_value=b"".decode("utf-8"),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="score",
+ full_name="sentencepiece.ModelProto.SentencePiece.score",
+ index=1,
+ number=2,
+ type=2,
+ cpp_type=6,
+ label=1,
+ has_default_value=False,
+ default_value=float(0),
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="type",
+ full_name="sentencepiece.ModelProto.SentencePiece.type",
+ index=2,
+ number=3,
+ type=14,
+ cpp_type=8,
+ label=1,
+ has_default_value=True,
+ default_value=1,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[],
+ enum_types=[
+ _MODELPROTO_SENTENCEPIECE_TYPE,
+ ],
+ serialized_options=None,
+ is_extendable=True,
+ syntax="proto2",
+ extension_ranges=[
+ (200, 536870912),
+ ],
+ oneofs=[],
+ serialized_start=1985,
+ serialized_end=2195,
+)
+
+_MODELPROTO = _descriptor.Descriptor(
+ name="ModelProto",
+ full_name="sentencepiece.ModelProto",
+ filename=None,
+ file=DESCRIPTOR,
+ containing_type=None,
+ create_key=_descriptor._internal_create_key,
+ fields=[
+ _descriptor.FieldDescriptor(
+ name="pieces",
+ full_name="sentencepiece.ModelProto.pieces",
+ index=0,
+ number=1,
+ type=11,
+ cpp_type=10,
+ label=3,
+ has_default_value=False,
+ default_value=[],
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="trainer_spec",
+ full_name="sentencepiece.ModelProto.trainer_spec",
+ index=1,
+ number=2,
+ type=11,
+ cpp_type=10,
+ label=1,
+ has_default_value=False,
+ default_value=None,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="normalizer_spec",
+ full_name="sentencepiece.ModelProto.normalizer_spec",
+ index=2,
+ number=3,
+ type=11,
+ cpp_type=10,
+ label=1,
+ has_default_value=False,
+ default_value=None,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="self_test_data",
+ full_name="sentencepiece.ModelProto.self_test_data",
+ index=3,
+ number=4,
+ type=11,
+ cpp_type=10,
+ label=1,
+ has_default_value=False,
+ default_value=None,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ _descriptor.FieldDescriptor(
+ name="denormalizer_spec",
+ full_name="sentencepiece.ModelProto.denormalizer_spec",
+ index=4,
+ number=5,
+ type=11,
+ cpp_type=10,
+ label=1,
+ has_default_value=False,
+ default_value=None,
+ message_type=None,
+ enum_type=None,
+ containing_type=None,
+ is_extension=False,
+ extension_scope=None,
+ serialized_options=None,
+ file=DESCRIPTOR,
+ create_key=_descriptor._internal_create_key,
+ ),
+ ],
+ extensions=[],
+ nested_types=[
+ _MODELPROTO_SENTENCEPIECE,
+ ],
+ enum_types=[],
+ serialized_options=None,
+ is_extendable=True,
+ syntax="proto2",
+ extension_ranges=[
+ (200, 536870912),
+ ],
+ oneofs=[],
+ serialized_start=1696,
+ serialized_end=2206,
+)
+
+_TRAINERSPEC.fields_by_name["model_type"].enum_type = _TRAINERSPEC_MODELTYPE
+_TRAINERSPEC_MODELTYPE.containing_type = _TRAINERSPEC
+_SELFTESTDATA_SAMPLE.containing_type = _SELFTESTDATA
+_SELFTESTDATA.fields_by_name["samples"].message_type = _SELFTESTDATA_SAMPLE
+_MODELPROTO_SENTENCEPIECE.fields_by_name["type"].enum_type = _MODELPROTO_SENTENCEPIECE_TYPE
+_MODELPROTO_SENTENCEPIECE.containing_type = _MODELPROTO
+_MODELPROTO_SENTENCEPIECE_TYPE.containing_type = _MODELPROTO_SENTENCEPIECE
+_MODELPROTO.fields_by_name["pieces"].message_type = _MODELPROTO_SENTENCEPIECE
+_MODELPROTO.fields_by_name["trainer_spec"].message_type = _TRAINERSPEC
+_MODELPROTO.fields_by_name["normalizer_spec"].message_type = _NORMALIZERSPEC
+_MODELPROTO.fields_by_name["self_test_data"].message_type = _SELFTESTDATA
+_MODELPROTO.fields_by_name["denormalizer_spec"].message_type = _NORMALIZERSPEC
+DESCRIPTOR.message_types_by_name["TrainerSpec"] = _TRAINERSPEC
+DESCRIPTOR.message_types_by_name["NormalizerSpec"] = _NORMALIZERSPEC
+DESCRIPTOR.message_types_by_name["SelfTestData"] = _SELFTESTDATA
+DESCRIPTOR.message_types_by_name["ModelProto"] = _MODELPROTO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+TrainerSpec = _reflection.GeneratedProtocolMessageType(
+ "TrainerSpec",
+ (_message.Message,),
+ {
+ "DESCRIPTOR": _TRAINERSPEC,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.TrainerSpec)
+ },
+)
+_sym_db.RegisterMessage(TrainerSpec)
+
+NormalizerSpec = _reflection.GeneratedProtocolMessageType(
+ "NormalizerSpec",
+ (_message.Message,),
+ {
+ "DESCRIPTOR": _NORMALIZERSPEC,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.NormalizerSpec)
+ },
+)
+_sym_db.RegisterMessage(NormalizerSpec)
+
+SelfTestData = _reflection.GeneratedProtocolMessageType(
+ "SelfTestData",
+ (_message.Message,),
+ {
+ "Sample": _reflection.GeneratedProtocolMessageType(
+ "Sample",
+ (_message.Message,),
+ {
+ "DESCRIPTOR": _SELFTESTDATA_SAMPLE,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData.Sample)
+ },
+ ),
+ "DESCRIPTOR": _SELFTESTDATA,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData)
+ },
+)
+_sym_db.RegisterMessage(SelfTestData)
+_sym_db.RegisterMessage(SelfTestData.Sample)
+
+ModelProto = _reflection.GeneratedProtocolMessageType(
+ "ModelProto",
+ (_message.Message,),
+ {
+ "SentencePiece": _reflection.GeneratedProtocolMessageType(
+ "SentencePiece",
+ (_message.Message,),
+ {
+ "DESCRIPTOR": _MODELPROTO_SENTENCEPIECE,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto.SentencePiece)
+ },
+ ),
+ "DESCRIPTOR": _MODELPROTO,
+ "__module__": "sentencepiece_model_pb2",
+ # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto)
+ },
+)
+_sym_db.RegisterMessage(ModelProto)
+_sym_db.RegisterMessage(ModelProto.SentencePiece)
+
+
+DESCRIPTOR._options = None
+_TRAINERSPEC.fields_by_name["mining_sentence_size"]._options = None
+_TRAINERSPEC.fields_by_name["training_sentence_size"]._options = None
+# @@protoc_insertion_point(module_scope)
diff --git a/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2_new.py b/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2_new.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ea4f4d64ce660b3d59d60c0650f5849e4fd2251
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/sentencepiece_model_pb2_new.py
@@ -0,0 +1,47 @@
+# Generated by the protocol buffer compiler. DO NOT EDIT!
+# source: sentencepiece_model.proto
+"""Generated protocol buffer code."""
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+ b'\n\x19sentencepiece_model.proto\x12\rsentencepiece"\x80\x0c\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01 \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02 \x01(\t\x12\x41\n\nmodel_type\x18\x03 \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04 \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12 \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12*\n\x1b\x65nable_differential_privacy\x18\x32 \x01(\x08:\x05\x66\x61lse\x12+\n differential_privacy_noise_level\x18\x33 \x01(\x02:\x01\x30\x12\x32\n\'differential_privacy_clipping_threshold\x18\x34 \x01(\x04:\x01\x30\x12"\n\x12\x63haracter_coverage\x18\n \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12 \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12"\n\x16training_sentence_size\x18\r \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12 \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10 \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11 \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14 \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15 \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17 \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16 \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18 \x01(\x08:\x05\x66\x61lse\x12+\n\x1c\x61llow_whitespace_only_pieces\x18\x1a \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19 \x01(\x08:\x05\x66\x61lse\x12#\n\x19pretokenization_delimiter\x18\x35 \x01(\t:\x00\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$ \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18 \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18) \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+ \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05\x12\x16\n\tbos_piece\x18. \x01(\t:\x03\x12\x17\n\teos_piece\x18/ \x01(\t:\x04\x12\x18\n\tpad_piece\x18\x30 \x01(\t:\x05\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87 \x12+\n\x1ctrain_extremely_large_corpus\x18\x31 \x01(\x08:\x05\x66\x61lse"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03 \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12 \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01 \x03(\x0b\x32".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01 \x01(\t\x12\x10\n\x08\x65xpected\x18\x02 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01 \x03(\x0b\x32\'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02 \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04 \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01 \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03 \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "sentencepiece_model_pb2", _globals)
+if _descriptor._USE_C_DESCRIPTORS is False:
+ DESCRIPTOR._options = None
+ DESCRIPTOR._serialized_options = b"H\003"
+ # (generated by protobuf compiler, but `_TRAINERSPEC` is not defined)
+ # _TRAINERSPEC.fields_by_name["mining_sentence_size"]._options = None
+ # _TRAINERSPEC.fields_by_name["mining_sentence_size"]._serialized_options = b"\030\001"
+ # _TRAINERSPEC.fields_by_name["training_sentence_size"]._options = None
+ # _TRAINERSPEC.fields_by_name["training_sentence_size"]._serialized_options = b"\030\001"
+ _globals["_TRAINERSPEC"]._serialized_start = 45
+ _globals["_TRAINERSPEC"]._serialized_end = 1581
+ _globals["_TRAINERSPEC_MODELTYPE"]._serialized_start = 1517
+ _globals["_TRAINERSPEC_MODELTYPE"]._serialized_end = 1570
+ _globals["_NORMALIZERSPEC"]._serialized_start = 1584
+ _globals["_NORMALIZERSPEC"]._serialized_end = 1793
+ _globals["_SELFTESTDATA"]._serialized_start = 1795
+ _globals["_SELFTESTDATA"]._serialized_end = 1916
+ _globals["_SELFTESTDATA_SAMPLE"]._serialized_start = 1864
+ _globals["_SELFTESTDATA_SAMPLE"]._serialized_end = 1905
+ _globals["_MODELPROTO"]._serialized_start = 1919
+ _globals["_MODELPROTO"]._serialized_end = 2429
+ _globals["_MODELPROTO_SENTENCEPIECE"]._serialized_start = 2208
+ _globals["_MODELPROTO_SENTENCEPIECE"]._serialized_end = 2418
+ _globals["_MODELPROTO_SENTENCEPIECE_TYPE"]._serialized_start = 2323
+ _globals["_MODELPROTO_SENTENCEPIECE_TYPE"]._serialized_end = 2407
+# @@protoc_insertion_point(module_scope)
diff --git a/third_party/transformers/src/transformers/utils/type_validators.py b/third_party/transformers/src/transformers/utils/type_validators.py
new file mode 100644
index 0000000000000000000000000000000000000000..08d4697683b2eea80cda1f80e4a4a84ffbf694d0
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/type_validators.py
@@ -0,0 +1,252 @@
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import Any, Union, cast
+
+from huggingface_hub.dataclasses import as_validated_field
+
+from ..tokenization_utils_base import PaddingStrategy, TruncationStrategy
+from ..video_utils import VideoMetadataType
+from .generic import TensorType
+from .import_utils import is_torch_available, is_vision_available
+
+
+if is_vision_available():
+ from ..image_utils import PILImageResampling
+
+if is_torch_available():
+ import torch
+
+ from ..activations import ACT2FN
+else:
+ ACT2FN = {}
+
+
+def positive_any_number(value: int | float | None = None):
+ if value is not None and (not isinstance(value, (int, float)) or not value >= 0):
+ raise ValueError(f"Value must be a positive integer or floating number, got {value}")
+
+
+def positive_int(value: int | None = None):
+ if value is not None and (not isinstance(value, int) or not value >= 0):
+ raise ValueError(f"Value must be a positive integer, got {value}")
+
+
+def padding_validator(value: bool | str | PaddingStrategy | None = None):
+ possible_names = ["longest", "max_length", "do_not_pad"]
+ if value is None:
+ pass
+ elif not isinstance(value, (bool, str, PaddingStrategy)):
+ raise ValueError("Value for padding must be either a boolean, a string or a `PaddingStrategy`")
+ elif isinstance(value, str) and value not in possible_names:
+ raise ValueError(f"If padding is a string, the value must be one of {possible_names}")
+
+
+def truncation_validator(value: bool | str | TruncationStrategy | None = None):
+ possible_names = ["only_first", "only_second", "longest_first", "do_not_truncate"]
+ if value is None:
+ pass
+ elif not isinstance(value, (bool, str, TruncationStrategy)):
+ raise ValueError("Value for truncation must be either a boolean, a string or a `TruncationStrategy`")
+ elif isinstance(value, str) and value not in possible_names:
+ raise ValueError(f"If truncation is a string, value must be one of {possible_names}")
+
+
+def image_size_validator(value: int | Sequence[int] | dict[str, int] | None = None):
+ possible_keys = ["height", "width", "longest_edge", "shortest_edge", "max_height", "max_width"]
+ if value is None:
+ pass
+ elif isinstance(value, dict) and any(k not in possible_keys for k in value.keys()):
+ raise ValueError(f"Value for size must be a dict with keys {possible_keys} but got size={value}")
+
+
+def device_validator(value: str | int | None = None):
+ possible_names = ["cpu", "cuda", "xla", "xpu", "mps", "meta"]
+ if value is None:
+ pass
+ elif is_torch_available() and isinstance(value, torch.device):
+ # Convert torch.device to string for validation
+ device_str = str(value)
+ if device_str.split(":")[0] not in possible_names:
+ raise ValueError(
+ f"If device is a torch.device, the value must be one of {possible_names} but got device={value}"
+ )
+ elif isinstance(value, int) and value < 0:
+ raise ValueError(
+ f"If device is an integer, the value must be a strictly positive integer but got device={value}"
+ )
+ elif isinstance(value, str) and value.split(":")[0] not in possible_names:
+ raise ValueError(f"If device is an string, the value must be one of {possible_names} but got device={value}")
+ elif not isinstance(value, (int, str)):
+ raise ValueError(
+ f"Device must be either an integer device ID, a string (e.g., 'cpu', 'cuda:0'), or a torch.device object, but got device={value}"
+ )
+
+
+def resampling_validator(value: Union[int, "PILImageResampling"] | None = None):
+ if value is None:
+ pass
+ elif isinstance(value, int) and value not in list(range(6)):
+ raise ValueError(
+ f"The resampling should be one of {list(range(6))} when provided as integer, but got resampling={value}"
+ )
+ elif is_vision_available() and not isinstance(value, (PILImageResampling, int)):
+ raise ValueError(f"The resampling should an integer or `PIL.Image.Resampling`, but got resampling={value}")
+
+
+def video_metadata_validator(value: VideoMetadataType | None = None):
+ if value is None:
+ return
+
+ valid_keys = ["total_num_frames", "fps", "width", "height", "duration", "video_backend", "frames_indices"]
+
+ def check_dict_keys(d: dict[str, Any]) -> bool:
+ return all(key in valid_keys for key in d.keys())
+
+ if isinstance(value, Sequence) and isinstance(value[0], Sequence) and isinstance(value[0][0], dict):
+ for sublist in value:
+ for item in sublist:
+ if not check_dict_keys(item):
+ raise ValueError(
+ f"Invalid keys found in video metadata. Valid keys: {valid_keys} got: {list(item.keys())}"
+ )
+
+ elif isinstance(value, Sequence) and isinstance(value[0], dict):
+ for item in value:
+ if not check_dict_keys(item):
+ raise ValueError(
+ f"Invalid keys found in video metadata. Valid keys: {valid_keys} got: {list(cast(dict, item).keys())}"
+ )
+
+ elif isinstance(value, dict):
+ if not check_dict_keys(value):
+ raise ValueError(
+ f"Invalid keys found in video metadata. Valid keys: {valid_keys}, got: {list(value.keys())}"
+ )
+
+
+def tensor_type_validator(value: str | TensorType | None = None):
+ possible_names = ["pt", "np", "mlx"]
+ if value is None:
+ pass
+ elif not isinstance(value, str) or value not in possible_names:
+ raise ValueError(f"The tensor type should be one of {possible_names} but got tensor_type={value}")
+
+
+@as_validated_field
+def label_to_id_validation(value: str | TensorType | None = None):
+ possible_names = ["pt", "np", "mlx"]
+ if value is None:
+ pass
+ elif not isinstance(value, str) or value not in possible_names:
+ raise ValueError(f"The tensor type should be one of {possible_names} but got tensor_type={value}")
+
+
+def interval(
+ min: int | float | None = None,
+ max: int | float | None = None,
+ exclude_min: bool = False,
+ exclude_max: bool = False,
+) -> Callable:
+ """
+ Parameterized validator that ensures that `value` is within the defined interval. Optionally, the interval can be
+ open on either side. Expected usage: `interval(min=0)(default=8)`
+
+ Args:
+ min (`int` or `float`, *optional*):
+ Minimum value of the interval.
+ max (`int` or `float`, *optional*):
+ Maximum value of the interval.
+ exclude_min (`bool`, *optional*, defaults to `False`):
+ If True, the minimum value is excluded from the interval.
+ exclude_max (`bool`, *optional*, defaults to `False`):
+ If True, the maximum value is excluded from the interval.
+ """
+ error_message = "Value must be"
+ if min is not None:
+ if exclude_min:
+ error_message += f" greater than {min}"
+ else:
+ error_message += f" greater or equal to {min}"
+ if min is not None and max is not None:
+ error_message += " and"
+ if max is not None:
+ if exclude_max:
+ error_message += f" smaller than {max}"
+ else:
+ error_message += f" smaller or equal to {max}"
+ error_message += ", got {value}."
+
+ min = min or float("-inf")
+ max = max or float("inf")
+
+ @as_validated_field
+ def _inner(value: int | float):
+ min_valid = min <= value if not exclude_min else min < value
+ max_valid = value <= max if not exclude_max else value < max
+ if not (min_valid and max_valid):
+ raise ValueError(error_message.format(value=value))
+
+ return _inner
+
+
+@as_validated_field
+def probability(value: float):
+ """Ensures that `value` is a valid probability number, i.e. [0,1]."""
+ if not 0 <= value <= 1:
+ raise ValueError(f"Value must be a probability between 0.0 and 1.0, got {value}.")
+
+
+def is_divisible_by(divisor: int | float):
+ @as_validated_field
+ def _inner(value: int | float):
+ if value % divisor != 0:
+ raise ValueError(f"Value has to be divisble by {divisor} but got value={value}")
+
+ return _inner
+
+
+@as_validated_field
+def activation_fn_key(value: str):
+ """Ensures that `value` is a string corresponding to an activation function."""
+ # TODO (joao): in python 3.11+, we can build a Literal type from the keys of ACT2FN
+ if len(ACT2FN) > 0: # don't validate if we can't import ACT2FN
+ if value not in ACT2FN:
+ raise ValueError(
+ f"Value must be one of {list(ACT2FN.keys())}, got {value}. "
+ "Make sure to use a string that corresponds to an activation function."
+ )
+
+
+def tensor_shape(shape: tuple[int | str], length: int | None = None):
+ @as_validated_field
+ def validator(value: Union[Sequence["torch.Tensor"], "torch.Tensor"]):
+ if value is None:
+ return
+ elif not isinstance(value, (list, tuple)):
+ value = [value]
+ elif isinstance(length, int) and len(value) != length:
+ raise ValueError(f"Value has to be a list of length={length} but got {len(value)}")
+
+ dimensions = {}
+ for tensor in value:
+ # Ensures that `value` is a floating point tensor in any device (cpu, cuda, xpu, ...).
+ # Using `torch.FloatTensor` as a type hint is discouraged if the dataclass has a `strict`
+ # decorator, because it enforces floating tensors only on CPU.
+ if not (isinstance(tensor, torch.Tensor) and tensor.is_floating_point()):
+ raise ValueError(f"Value has to be a floating point tensor but got value={tensor}")
+
+ if len(tensor.shape) != len(shape):
+ raise ValueError(f"Expected shape {shape}, but got {tensor.shape}")
+ for dim, expected in zip(tensor.shape, shape):
+ if isinstance(expected, int) and dim != expected:
+ raise ValueError(f"Expected dimension {expected}, but got {dim}")
+ elif isinstance(expected, str):
+ if expected not in dimensions:
+ dimensions[expected] = dim
+ elif dimensions[expected] != dim:
+ raise ValueError(
+ f"Dimension '{expected}' takes different values: {dimensions[expected]} and {dim}."
+ " Please check your tensors shapes."
+ )
+
+ return partial(validator, metadata={"shape": shape, "length": length})
diff --git a/third_party/transformers/src/transformers/utils/versions.py b/third_party/transformers/src/transformers/utils/versions.py
new file mode 100644
index 0000000000000000000000000000000000000000..452f0d48a8e279c6e16fba74d9d565b59da1cfc5
--- /dev/null
+++ b/third_party/transformers/src/transformers/utils/versions.py
@@ -0,0 +1,116 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""
+Utilities for working with package versions
+"""
+
+import importlib.metadata
+import operator
+import re
+import sys
+
+from packaging import version
+
+
+ops = {
+ "<": operator.lt,
+ "<=": operator.le,
+ "==": operator.eq,
+ "!=": operator.ne,
+ ">=": operator.ge,
+ ">": operator.gt,
+}
+
+
+def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint):
+ if got_ver is None or want_ver is None:
+ raise ValueError(
+ f"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider"
+ f" reinstalling {pkg}."
+ )
+ if not ops[op](version.parse(got_ver), version.parse(want_ver)):
+ raise ImportError(
+ f"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}"
+ )
+
+
+def require_version(requirement: str, hint: str | None = None) -> None:
+ """
+ Perform a runtime check of the dependency versions, using the exact same syntax used by pip.
+
+ The installed module version comes from the *site-packages* dir via *importlib.metadata*.
+
+ Args:
+ requirement (`str`): pip style definition, e.g., "tokenizers==0.9.4", "tqdm>=4.27", "numpy"
+ hint (`str`, *optional*): what suggestion to print in case of requirements not being met
+
+ Example:
+
+ ```python
+ require_version("pandas>1.1.2")
+ require_version("numpy>1.18.5", "this is important to have for whatever reason")
+ ```"""
+
+ hint = f"\n{hint}" if hint is not None else ""
+
+ # non-versioned check
+ if re.match(r"^[\w_\-\d]+$", requirement):
+ pkg, op, want_ver = requirement, None, None
+ else:
+ match = re.findall(r"^([^!=<>\s]+)([\s!=<>]{1,2}.+)", requirement)
+ if not match:
+ raise ValueError(
+ "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"
+ f" got {requirement}"
+ )
+ pkg, want_full = match[0]
+ want_range = want_full.split(",") # there could be multiple requirements
+ wanted = {}
+ for w in want_range:
+ match = re.findall(r"^([\s!=<>]{1,2})(.+)", w)
+ if not match:
+ raise ValueError(
+ "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"
+ f" but got {requirement}"
+ )
+ op, want_ver = match[0]
+ wanted[op] = want_ver
+ if op not in ops:
+ raise ValueError(f"{requirement}: need one of {list(ops.keys())}, but got {op}")
+
+ # special case
+ if pkg == "python":
+ got_ver = ".".join([str(x) for x in sys.version_info[:3]])
+ for op, want_ver in wanted.items():
+ _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)
+ return
+
+ # check if any version is installed
+ try:
+ got_ver = importlib.metadata.version(pkg)
+ except importlib.metadata.PackageNotFoundError:
+ raise importlib.metadata.PackageNotFoundError(
+ f"The '{requirement}' distribution was not found and is required by this application. {hint}"
+ )
+
+ # check that the right version is installed if version number or a range was provided
+ if want_ver is not None:
+ for op, want_ver in wanted.items():
+ _compare_versions(op, got_ver, want_ver, requirement, pkg, hint)
+
+
+def require_version_core(requirement):
+ """require_version wrapper which emits a core-specific hint on failure"""
+ hint = "Try: `pip install transformers -U` or `pip install -e '.[dev]'` if you're working with git main"
+ return require_version(requirement, hint)
+ 
+ 
+ 
+ 
+
. We discuss the hottest trends about diffusion models, help each other with contributions, personal projects or just hang out โ.
+
+
+## Popular Tasks & Pipelines
+
+
+ 
+ 
+ 
+ 
+ 
+
+
+
+