diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7855226e4b500142deef8fb247cd33a9a991d122
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/__init__.py
@@ -0,0 +1,2 @@
+"""A package that contains models that represent entities.
+"""
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/direct_url.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/direct_url.py
new file mode 100644
index 0000000000000000000000000000000000000000..0af884bd8e38ddaf9ab0ce1af4749ff549c8e9da
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/direct_url.py
@@ -0,0 +1,235 @@
+""" PEP 610 """
+import json
+import re
+import urllib.parse
+from typing import Any, Dict, Iterable, Optional, Type, TypeVar, Union
+
+__all__ = [
+ "DirectUrl",
+ "DirectUrlValidationError",
+ "DirInfo",
+ "ArchiveInfo",
+ "VcsInfo",
+]
+
+T = TypeVar("T")
+
+DIRECT_URL_METADATA_NAME = "direct_url.json"
+ENV_VAR_RE = re.compile(r"^\$\{[A-Za-z0-9-_]+\}(:\$\{[A-Za-z0-9-_]+\})?$")
+
+
+class DirectUrlValidationError(Exception):
+ pass
+
+
+def _get(
+ d: Dict[str, Any], expected_type: Type[T], key: str, default: Optional[T] = None
+) -> Optional[T]:
+ """Get value from dictionary and verify expected type."""
+ if key not in d:
+ return default
+ value = d[key]
+ if not isinstance(value, expected_type):
+ raise DirectUrlValidationError(
+ f"{value!r} has unexpected type for {key} (expected {expected_type})"
+ )
+ return value
+
+
+def _get_required(
+ d: Dict[str, Any], expected_type: Type[T], key: str, default: Optional[T] = None
+) -> T:
+ value = _get(d, expected_type, key, default)
+ if value is None:
+ raise DirectUrlValidationError(f"{key} must have a value")
+ return value
+
+
+def _exactly_one_of(infos: Iterable[Optional["InfoType"]]) -> "InfoType":
+ infos = [info for info in infos if info is not None]
+ if not infos:
+ raise DirectUrlValidationError(
+ "missing one of archive_info, dir_info, vcs_info"
+ )
+ if len(infos) > 1:
+ raise DirectUrlValidationError(
+ "more than one of archive_info, dir_info, vcs_info"
+ )
+ assert infos[0] is not None
+ return infos[0]
+
+
+def _filter_none(**kwargs: Any) -> Dict[str, Any]:
+ """Make dict excluding None values."""
+ return {k: v for k, v in kwargs.items() if v is not None}
+
+
+class VcsInfo:
+ name = "vcs_info"
+
+ def __init__(
+ self,
+ vcs: str,
+ commit_id: str,
+ requested_revision: Optional[str] = None,
+ ) -> None:
+ self.vcs = vcs
+ self.requested_revision = requested_revision
+ self.commit_id = commit_id
+
+ @classmethod
+ def _from_dict(cls, d: Optional[Dict[str, Any]]) -> Optional["VcsInfo"]:
+ if d is None:
+ return None
+ return cls(
+ vcs=_get_required(d, str, "vcs"),
+ commit_id=_get_required(d, str, "commit_id"),
+ requested_revision=_get(d, str, "requested_revision"),
+ )
+
+ def _to_dict(self) -> Dict[str, Any]:
+ return _filter_none(
+ vcs=self.vcs,
+ requested_revision=self.requested_revision,
+ commit_id=self.commit_id,
+ )
+
+
+class ArchiveInfo:
+ name = "archive_info"
+
+ def __init__(
+ self,
+ hash: Optional[str] = None,
+ hashes: Optional[Dict[str, str]] = None,
+ ) -> None:
+ # set hashes before hash, since the hash setter will further populate hashes
+ self.hashes = hashes
+ self.hash = hash
+
+ @property
+ def hash(self) -> Optional[str]:
+ return self._hash
+
+ @hash.setter
+ def hash(self, value: Optional[str]) -> None:
+ if value is not None:
+ # Auto-populate the hashes key to upgrade to the new format automatically.
+ # We don't back-populate the legacy hash key from hashes.
+ try:
+ hash_name, hash_value = value.split("=", 1)
+ except ValueError:
+ raise DirectUrlValidationError(
+ f"invalid archive_info.hash format: {value!r}"
+ )
+ if self.hashes is None:
+ self.hashes = {hash_name: hash_value}
+ elif hash_name not in self.hashes:
+ self.hashes = self.hashes.copy()
+ self.hashes[hash_name] = hash_value
+ self._hash = value
+
+ @classmethod
+ def _from_dict(cls, d: Optional[Dict[str, Any]]) -> Optional["ArchiveInfo"]:
+ if d is None:
+ return None
+ return cls(hash=_get(d, str, "hash"), hashes=_get(d, dict, "hashes"))
+
+ def _to_dict(self) -> Dict[str, Any]:
+ return _filter_none(hash=self.hash, hashes=self.hashes)
+
+
+class DirInfo:
+ name = "dir_info"
+
+ def __init__(
+ self,
+ editable: bool = False,
+ ) -> None:
+ self.editable = editable
+
+ @classmethod
+ def _from_dict(cls, d: Optional[Dict[str, Any]]) -> Optional["DirInfo"]:
+ if d is None:
+ return None
+ return cls(editable=_get_required(d, bool, "editable", default=False))
+
+ def _to_dict(self) -> Dict[str, Any]:
+ return _filter_none(editable=self.editable or None)
+
+
+InfoType = Union[ArchiveInfo, DirInfo, VcsInfo]
+
+
+class DirectUrl:
+ def __init__(
+ self,
+ url: str,
+ info: InfoType,
+ subdirectory: Optional[str] = None,
+ ) -> None:
+ self.url = url
+ self.info = info
+ self.subdirectory = subdirectory
+
+ def _remove_auth_from_netloc(self, netloc: str) -> str:
+ if "@" not in netloc:
+ return netloc
+ user_pass, netloc_no_user_pass = netloc.split("@", 1)
+ if (
+ isinstance(self.info, VcsInfo)
+ and self.info.vcs == "git"
+ and user_pass == "git"
+ ):
+ return netloc
+ if ENV_VAR_RE.match(user_pass):
+ return netloc
+ return netloc_no_user_pass
+
+ @property
+ def redacted_url(self) -> str:
+ """url with user:password part removed unless it is formed with
+ environment variables as specified in PEP 610, or it is ``git``
+ in the case of a git URL.
+ """
+ purl = urllib.parse.urlsplit(self.url)
+ netloc = self._remove_auth_from_netloc(purl.netloc)
+ surl = urllib.parse.urlunsplit(
+ (purl.scheme, netloc, purl.path, purl.query, purl.fragment)
+ )
+ return surl
+
+ def validate(self) -> None:
+ self.from_dict(self.to_dict())
+
+ @classmethod
+ def from_dict(cls, d: Dict[str, Any]) -> "DirectUrl":
+ return DirectUrl(
+ url=_get_required(d, str, "url"),
+ subdirectory=_get(d, str, "subdirectory"),
+ info=_exactly_one_of(
+ [
+ ArchiveInfo._from_dict(_get(d, dict, "archive_info")),
+ DirInfo._from_dict(_get(d, dict, "dir_info")),
+ VcsInfo._from_dict(_get(d, dict, "vcs_info")),
+ ]
+ ),
+ )
+
+ def to_dict(self) -> Dict[str, Any]:
+ res = _filter_none(
+ url=self.redacted_url,
+ subdirectory=self.subdirectory,
+ )
+ res[self.info.name] = self.info._to_dict()
+ return res
+
+ @classmethod
+ def from_json(cls, s: str) -> "DirectUrl":
+ return cls.from_dict(json.loads(s))
+
+ def to_json(self) -> str:
+ return json.dumps(self.to_dict(), sort_keys=True)
+
+ def is_local_editable(self) -> bool:
+ return isinstance(self.info, DirInfo) and self.info.editable
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/format_control.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/format_control.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccd11272c030c2d067e1bb6d90fc744c7379a923
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/format_control.py
@@ -0,0 +1,78 @@
+from typing import FrozenSet, Optional, Set
+
+from pip._vendor.packaging.utils import canonicalize_name
+
+from pip._internal.exceptions import CommandError
+
+
+class FormatControl:
+ """Helper for managing formats from which a package can be installed."""
+
+ __slots__ = ["no_binary", "only_binary"]
+
+ def __init__(
+ self,
+ no_binary: Optional[Set[str]] = None,
+ only_binary: Optional[Set[str]] = None,
+ ) -> None:
+ if no_binary is None:
+ no_binary = set()
+ if only_binary is None:
+ only_binary = set()
+
+ self.no_binary = no_binary
+ self.only_binary = only_binary
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, self.__class__):
+ return NotImplemented
+
+ if self.__slots__ != other.__slots__:
+ return False
+
+ return all(getattr(self, k) == getattr(other, k) for k in self.__slots__)
+
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__}({self.no_binary}, {self.only_binary})"
+
+ @staticmethod
+ def handle_mutual_excludes(value: str, target: Set[str], other: Set[str]) -> None:
+ if value.startswith("-"):
+ raise CommandError(
+ "--no-binary / --only-binary option requires 1 argument."
+ )
+ new = value.split(",")
+ while ":all:" in new:
+ other.clear()
+ target.clear()
+ target.add(":all:")
+ del new[: new.index(":all:") + 1]
+ # Without a none, we want to discard everything as :all: covers it
+ if ":none:" not in new:
+ return
+ for name in new:
+ if name == ":none:":
+ target.clear()
+ continue
+ name = canonicalize_name(name)
+ other.discard(name)
+ target.add(name)
+
+ def get_allowed_formats(self, canonical_name: str) -> FrozenSet[str]:
+ result = {"binary", "source"}
+ if canonical_name in self.only_binary:
+ result.discard("source")
+ elif canonical_name in self.no_binary:
+ result.discard("binary")
+ elif ":all:" in self.only_binary:
+ result.discard("source")
+ elif ":all:" in self.no_binary:
+ result.discard("binary")
+ return frozenset(result)
+
+ def disallow_binaries(self) -> None:
+ self.handle_mutual_excludes(
+ ":all:",
+ self.no_binary,
+ self.only_binary,
+ )
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/installation_report.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/installation_report.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9c6330df32bd2b57c885156cb7f8c0c8c3e3741
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/installation_report.py
@@ -0,0 +1,56 @@
+from typing import Any, Dict, Sequence
+
+from pip._vendor.packaging.markers import default_environment
+
+from pip import __version__
+from pip._internal.req.req_install import InstallRequirement
+
+
+class InstallationReport:
+ def __init__(self, install_requirements: Sequence[InstallRequirement]):
+ self._install_requirements = install_requirements
+
+ @classmethod
+ def _install_req_to_dict(cls, ireq: InstallRequirement) -> Dict[str, Any]:
+ assert ireq.download_info, f"No download_info for {ireq}"
+ res = {
+ # PEP 610 json for the download URL. download_info.archive_info.hashes may
+ # be absent when the requirement was installed from the wheel cache
+ # and the cache entry was populated by an older pip version that did not
+ # record origin.json.
+ "download_info": ireq.download_info.to_dict(),
+ # is_direct is true if the requirement was a direct URL reference (which
+ # includes editable requirements), and false if the requirement was
+ # downloaded from a PEP 503 index or --find-links.
+ "is_direct": ireq.is_direct,
+ # is_yanked is true if the requirement was yanked from the index, but
+ # was still selected by pip to conform to PEP 592.
+ "is_yanked": ireq.link.is_yanked if ireq.link else False,
+ # requested is true if the requirement was specified by the user (aka
+ # top level requirement), and false if it was installed as a dependency of a
+ # requirement. https://peps.python.org/pep-0376/#requested
+ "requested": ireq.user_supplied,
+ # PEP 566 json encoding for metadata
+ # https://www.python.org/dev/peps/pep-0566/#json-compatible-metadata
+ "metadata": ireq.get_dist().metadata_dict,
+ }
+ if ireq.user_supplied and ireq.extras:
+ # For top level requirements, the list of requested extras, if any.
+ res["requested_extras"] = sorted(ireq.extras)
+ return res
+
+ def to_dict(self) -> Dict[str, Any]:
+ return {
+ "version": "1",
+ "pip_version": __version__,
+ "install": [
+ self._install_req_to_dict(ireq) for ireq in self._install_requirements
+ ],
+ # https://peps.python.org/pep-0508/#environment-markers
+ # TODO: currently, the resolver uses the default environment to evaluate
+ # environment markers, so that is what we report here. In the future, it
+ # should also take into account options such as --python-version or
+ # --platform, perhaps under the form of an environment_override field?
+ # https://github.com/pypa/pip/issues/11198
+ "environment": default_environment(),
+ }
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/link.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/link.py
new file mode 100644
index 0000000000000000000000000000000000000000..73041b864c33def0f5d11a5b761c222cbbc62821
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/link.py
@@ -0,0 +1,579 @@
+import functools
+import itertools
+import logging
+import os
+import posixpath
+import re
+import urllib.parse
+from dataclasses import dataclass
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Dict,
+ List,
+ Mapping,
+ NamedTuple,
+ Optional,
+ Tuple,
+ Union,
+)
+
+from pip._internal.utils.deprecation import deprecated
+from pip._internal.utils.filetypes import WHEEL_EXTENSION
+from pip._internal.utils.hashes import Hashes
+from pip._internal.utils.misc import (
+ pairwise,
+ redact_auth_from_url,
+ split_auth_from_netloc,
+ splitext,
+)
+from pip._internal.utils.models import KeyBasedCompareMixin
+from pip._internal.utils.urls import path_to_url, url_to_path
+
+if TYPE_CHECKING:
+ from pip._internal.index.collector import IndexContent
+
+logger = logging.getLogger(__name__)
+
+
+# Order matters, earlier hashes have a precedence over later hashes for what
+# we will pick to use.
+_SUPPORTED_HASHES = ("sha512", "sha384", "sha256", "sha224", "sha1", "md5")
+
+
+@dataclass(frozen=True)
+class LinkHash:
+ """Links to content may have embedded hash values. This class parses those.
+
+ `name` must be any member of `_SUPPORTED_HASHES`.
+
+ This class can be converted to and from `ArchiveInfo`. While ArchiveInfo intends to
+ be JSON-serializable to conform to PEP 610, this class contains the logic for
+ parsing a hash name and value for correctness, and then checking whether that hash
+ conforms to a schema with `.is_hash_allowed()`."""
+
+ name: str
+ value: str
+
+ _hash_url_fragment_re = re.compile(
+ # NB: we do not validate that the second group (.*) is a valid hex
+ # digest. Instead, we simply keep that string in this class, and then check it
+ # against Hashes when hash-checking is needed. This is easier to debug than
+ # proactively discarding an invalid hex digest, as we handle incorrect hashes
+ # and malformed hashes in the same place.
+ r"[#&]({choices})=([^&]*)".format(
+ choices="|".join(re.escape(hash_name) for hash_name in _SUPPORTED_HASHES)
+ ),
+ )
+
+ def __post_init__(self) -> None:
+ assert self.name in _SUPPORTED_HASHES
+
+ @classmethod
+ @functools.lru_cache(maxsize=None)
+ def find_hash_url_fragment(cls, url: str) -> Optional["LinkHash"]:
+ """Search a string for a checksum algorithm name and encoded output value."""
+ match = cls._hash_url_fragment_re.search(url)
+ if match is None:
+ return None
+ name, value = match.groups()
+ return cls(name=name, value=value)
+
+ def as_dict(self) -> Dict[str, str]:
+ return {self.name: self.value}
+
+ def as_hashes(self) -> Hashes:
+ """Return a Hashes instance which checks only for the current hash."""
+ return Hashes({self.name: [self.value]})
+
+ def is_hash_allowed(self, hashes: Optional[Hashes]) -> bool:
+ """
+ Return True if the current hash is allowed by `hashes`.
+ """
+ if hashes is None:
+ return False
+ return hashes.is_hash_allowed(self.name, hex_digest=self.value)
+
+
+@dataclass(frozen=True)
+class MetadataFile:
+ """Information about a core metadata file associated with a distribution."""
+
+ hashes: Optional[Dict[str, str]]
+
+ def __post_init__(self) -> None:
+ if self.hashes is not None:
+ assert all(name in _SUPPORTED_HASHES for name in self.hashes)
+
+
+def supported_hashes(hashes: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
+ # Remove any unsupported hash types from the mapping. If this leaves no
+ # supported hashes, return None
+ if hashes is None:
+ return None
+ hashes = {n: v for n, v in hashes.items() if n in _SUPPORTED_HASHES}
+ if not hashes:
+ return None
+ return hashes
+
+
+def _clean_url_path_part(part: str) -> str:
+ """
+ Clean a "part" of a URL path (i.e. after splitting on "@" characters).
+ """
+ # We unquote prior to quoting to make sure nothing is double quoted.
+ return urllib.parse.quote(urllib.parse.unquote(part))
+
+
+def _clean_file_url_path(part: str) -> str:
+ """
+ Clean the first part of a URL path that corresponds to a local
+ filesystem path (i.e. the first part after splitting on "@" characters).
+ """
+ # We unquote prior to quoting to make sure nothing is double quoted.
+ # Also, on Windows the path part might contain a drive letter which
+ # should not be quoted. On Linux where drive letters do not
+ # exist, the colon should be quoted. We rely on urllib.request
+ # to do the right thing here.
+ return urllib.request.pathname2url(urllib.request.url2pathname(part))
+
+
+# percent-encoded: /
+_reserved_chars_re = re.compile("(@|%2F)", re.IGNORECASE)
+
+
+def _clean_url_path(path: str, is_local_path: bool) -> str:
+ """
+ Clean the path portion of a URL.
+ """
+ if is_local_path:
+ clean_func = _clean_file_url_path
+ else:
+ clean_func = _clean_url_path_part
+
+ # Split on the reserved characters prior to cleaning so that
+ # revision strings in VCS URLs are properly preserved.
+ parts = _reserved_chars_re.split(path)
+
+ cleaned_parts = []
+ for to_clean, reserved in pairwise(itertools.chain(parts, [""])):
+ cleaned_parts.append(clean_func(to_clean))
+ # Normalize %xx escapes (e.g. %2f -> %2F)
+ cleaned_parts.append(reserved.upper())
+
+ return "".join(cleaned_parts)
+
+
+def _ensure_quoted_url(url: str) -> str:
+ """
+ Make sure a link is fully quoted.
+ For example, if ' ' occurs in the URL, it will be replaced with "%20",
+ and without double-quoting other characters.
+ """
+ # Split the URL into parts according to the general structure
+ # `scheme://netloc/path;parameters?query#fragment`.
+ result = urllib.parse.urlparse(url)
+ # If the netloc is empty, then the URL refers to a local filesystem path.
+ is_local_path = not result.netloc
+ path = _clean_url_path(result.path, is_local_path=is_local_path)
+ return urllib.parse.urlunparse(result._replace(path=path))
+
+
+class Link(KeyBasedCompareMixin):
+ """Represents a parsed link from a Package Index's simple URL"""
+
+ __slots__ = [
+ "_parsed_url",
+ "_url",
+ "_hashes",
+ "comes_from",
+ "requires_python",
+ "yanked_reason",
+ "metadata_file_data",
+ "cache_link_parsing",
+ "egg_fragment",
+ ]
+
+ def __init__(
+ self,
+ url: str,
+ comes_from: Optional[Union[str, "IndexContent"]] = None,
+ requires_python: Optional[str] = None,
+ yanked_reason: Optional[str] = None,
+ metadata_file_data: Optional[MetadataFile] = None,
+ cache_link_parsing: bool = True,
+ hashes: Optional[Mapping[str, str]] = None,
+ ) -> None:
+ """
+ :param url: url of the resource pointed to (href of the link)
+ :param comes_from: instance of IndexContent where the link was found,
+ or string.
+ :param requires_python: String containing the `Requires-Python`
+ metadata field, specified in PEP 345. This may be specified by
+ a data-requires-python attribute in the HTML link tag, as
+ described in PEP 503.
+ :param yanked_reason: the reason the file has been yanked, if the
+ file has been yanked, or None if the file hasn't been yanked.
+ This is the value of the "data-yanked" attribute, if present, in
+ a simple repository HTML link. If the file has been yanked but
+ no reason was provided, this should be the empty string. See
+ PEP 592 for more information and the specification.
+ :param metadata_file_data: the metadata attached to the file, or None if
+ no such metadata is provided. This argument, if not None, indicates
+ that a separate metadata file exists, and also optionally supplies
+ hashes for that file.
+ :param cache_link_parsing: A flag that is used elsewhere to determine
+ whether resources retrieved from this link should be cached. PyPI
+ URLs should generally have this set to False, for example.
+ :param hashes: A mapping of hash names to digests to allow us to
+ determine the validity of a download.
+ """
+
+ # The comes_from, requires_python, and metadata_file_data arguments are
+ # only used by classmethods of this class, and are not used in client
+ # code directly.
+
+ # url can be a UNC windows share
+ if url.startswith("\\\\"):
+ url = path_to_url(url)
+
+ self._parsed_url = urllib.parse.urlsplit(url)
+ # Store the url as a private attribute to prevent accidentally
+ # trying to set a new value.
+ self._url = url
+
+ link_hash = LinkHash.find_hash_url_fragment(url)
+ hashes_from_link = {} if link_hash is None else link_hash.as_dict()
+ if hashes is None:
+ self._hashes = hashes_from_link
+ else:
+ self._hashes = {**hashes, **hashes_from_link}
+
+ self.comes_from = comes_from
+ self.requires_python = requires_python if requires_python else None
+ self.yanked_reason = yanked_reason
+ self.metadata_file_data = metadata_file_data
+
+ super().__init__(key=url, defining_class=Link)
+
+ self.cache_link_parsing = cache_link_parsing
+ self.egg_fragment = self._egg_fragment()
+
+ @classmethod
+ def from_json(
+ cls,
+ file_data: Dict[str, Any],
+ page_url: str,
+ ) -> Optional["Link"]:
+ """
+ Convert an pypi json document from a simple repository page into a Link.
+ """
+ file_url = file_data.get("url")
+ if file_url is None:
+ return None
+
+ url = _ensure_quoted_url(urllib.parse.urljoin(page_url, file_url))
+ pyrequire = file_data.get("requires-python")
+ yanked_reason = file_data.get("yanked")
+ hashes = file_data.get("hashes", {})
+
+ # PEP 714: Indexes must use the name core-metadata, but
+ # clients should support the old name as a fallback for compatibility.
+ metadata_info = file_data.get("core-metadata")
+ if metadata_info is None:
+ metadata_info = file_data.get("dist-info-metadata")
+
+ # The metadata info value may be a boolean, or a dict of hashes.
+ if isinstance(metadata_info, dict):
+ # The file exists, and hashes have been supplied
+ metadata_file_data = MetadataFile(supported_hashes(metadata_info))
+ elif metadata_info:
+ # The file exists, but there are no hashes
+ metadata_file_data = MetadataFile(None)
+ else:
+ # False or not present: the file does not exist
+ metadata_file_data = None
+
+ # The Link.yanked_reason expects an empty string instead of a boolean.
+ if yanked_reason and not isinstance(yanked_reason, str):
+ yanked_reason = ""
+ # The Link.yanked_reason expects None instead of False.
+ elif not yanked_reason:
+ yanked_reason = None
+
+ return cls(
+ url,
+ comes_from=page_url,
+ requires_python=pyrequire,
+ yanked_reason=yanked_reason,
+ hashes=hashes,
+ metadata_file_data=metadata_file_data,
+ )
+
+ @classmethod
+ def from_element(
+ cls,
+ anchor_attribs: Dict[str, Optional[str]],
+ page_url: str,
+ base_url: str,
+ ) -> Optional["Link"]:
+ """
+ Convert an anchor element's attributes in a simple repository page to a Link.
+ """
+ href = anchor_attribs.get("href")
+ if not href:
+ return None
+
+ url = _ensure_quoted_url(urllib.parse.urljoin(base_url, href))
+ pyrequire = anchor_attribs.get("data-requires-python")
+ yanked_reason = anchor_attribs.get("data-yanked")
+
+ # PEP 714: Indexes must use the name data-core-metadata, but
+ # clients should support the old name as a fallback for compatibility.
+ metadata_info = anchor_attribs.get("data-core-metadata")
+ if metadata_info is None:
+ metadata_info = anchor_attribs.get("data-dist-info-metadata")
+ # The metadata info value may be the string "true", or a string of
+ # the form "hashname=hashval"
+ if metadata_info == "true":
+ # The file exists, but there are no hashes
+ metadata_file_data = MetadataFile(None)
+ elif metadata_info is None:
+ # The file does not exist
+ metadata_file_data = None
+ else:
+ # The file exists, and hashes have been supplied
+ hashname, sep, hashval = metadata_info.partition("=")
+ if sep == "=":
+ metadata_file_data = MetadataFile(supported_hashes({hashname: hashval}))
+ else:
+ # Error - data is wrong. Treat as no hashes supplied.
+ logger.debug(
+ "Index returned invalid data-dist-info-metadata value: %s",
+ metadata_info,
+ )
+ metadata_file_data = MetadataFile(None)
+
+ return cls(
+ url,
+ comes_from=page_url,
+ requires_python=pyrequire,
+ yanked_reason=yanked_reason,
+ metadata_file_data=metadata_file_data,
+ )
+
+ def __str__(self) -> str:
+ if self.requires_python:
+ rp = f" (requires-python:{self.requires_python})"
+ else:
+ rp = ""
+ if self.comes_from:
+ return f"{redact_auth_from_url(self._url)} (from {self.comes_from}){rp}"
+ else:
+ return redact_auth_from_url(str(self._url))
+
+ def __repr__(self) -> str:
+ return f""
+
+ @property
+ def url(self) -> str:
+ return self._url
+
+ @property
+ def filename(self) -> str:
+ path = self.path.rstrip("/")
+ name = posixpath.basename(path)
+ if not name:
+ # Make sure we don't leak auth information if the netloc
+ # includes a username and password.
+ netloc, user_pass = split_auth_from_netloc(self.netloc)
+ return netloc
+
+ name = urllib.parse.unquote(name)
+ assert name, f"URL {self._url!r} produced no filename"
+ return name
+
+ @property
+ def file_path(self) -> str:
+ return url_to_path(self.url)
+
+ @property
+ def scheme(self) -> str:
+ return self._parsed_url.scheme
+
+ @property
+ def netloc(self) -> str:
+ """
+ This can contain auth information.
+ """
+ return self._parsed_url.netloc
+
+ @property
+ def path(self) -> str:
+ return urllib.parse.unquote(self._parsed_url.path)
+
+ def splitext(self) -> Tuple[str, str]:
+ return splitext(posixpath.basename(self.path.rstrip("/")))
+
+ @property
+ def ext(self) -> str:
+ return self.splitext()[1]
+
+ @property
+ def url_without_fragment(self) -> str:
+ scheme, netloc, path, query, fragment = self._parsed_url
+ return urllib.parse.urlunsplit((scheme, netloc, path, query, ""))
+
+ _egg_fragment_re = re.compile(r"[#&]egg=([^&]*)")
+
+ # Per PEP 508.
+ _project_name_re = re.compile(
+ r"^([A-Z0-9]|[A-Z0-9][A-Z0-9._-]*[A-Z0-9])$", re.IGNORECASE
+ )
+
+ def _egg_fragment(self) -> Optional[str]:
+ match = self._egg_fragment_re.search(self._url)
+ if not match:
+ return None
+
+ # An egg fragment looks like a PEP 508 project name, along with
+ # an optional extras specifier. Anything else is invalid.
+ project_name = match.group(1)
+ if not self._project_name_re.match(project_name):
+ deprecated(
+ reason=f"{self} contains an egg fragment with a non-PEP 508 name",
+ replacement="to use the req @ url syntax, and remove the egg fragment",
+ gone_in="25.0",
+ issue=11617,
+ )
+
+ return project_name
+
+ _subdirectory_fragment_re = re.compile(r"[#&]subdirectory=([^&]*)")
+
+ @property
+ def subdirectory_fragment(self) -> Optional[str]:
+ match = self._subdirectory_fragment_re.search(self._url)
+ if not match:
+ return None
+ return match.group(1)
+
+ def metadata_link(self) -> Optional["Link"]:
+ """Return a link to the associated core metadata file (if any)."""
+ if self.metadata_file_data is None:
+ return None
+ metadata_url = f"{self.url_without_fragment}.metadata"
+ if self.metadata_file_data.hashes is None:
+ return Link(metadata_url)
+ return Link(metadata_url, hashes=self.metadata_file_data.hashes)
+
+ def as_hashes(self) -> Hashes:
+ return Hashes({k: [v] for k, v in self._hashes.items()})
+
+ @property
+ def hash(self) -> Optional[str]:
+ return next(iter(self._hashes.values()), None)
+
+ @property
+ def hash_name(self) -> Optional[str]:
+ return next(iter(self._hashes), None)
+
+ @property
+ def show_url(self) -> str:
+ return posixpath.basename(self._url.split("#", 1)[0].split("?", 1)[0])
+
+ @property
+ def is_file(self) -> bool:
+ return self.scheme == "file"
+
+ def is_existing_dir(self) -> bool:
+ return self.is_file and os.path.isdir(self.file_path)
+
+ @property
+ def is_wheel(self) -> bool:
+ return self.ext == WHEEL_EXTENSION
+
+ @property
+ def is_vcs(self) -> bool:
+ from pip._internal.vcs import vcs
+
+ return self.scheme in vcs.all_schemes
+
+ @property
+ def is_yanked(self) -> bool:
+ return self.yanked_reason is not None
+
+ @property
+ def has_hash(self) -> bool:
+ return bool(self._hashes)
+
+ def is_hash_allowed(self, hashes: Optional[Hashes]) -> bool:
+ """
+ Return True if the link has a hash and it is allowed by `hashes`.
+ """
+ if hashes is None:
+ return False
+ return any(hashes.is_hash_allowed(k, v) for k, v in self._hashes.items())
+
+
+class _CleanResult(NamedTuple):
+ """Convert link for equivalency check.
+
+ This is used in the resolver to check whether two URL-specified requirements
+ likely point to the same distribution and can be considered equivalent. This
+ equivalency logic avoids comparing URLs literally, which can be too strict
+ (e.g. "a=1&b=2" vs "b=2&a=1") and produce conflicts unexpecting to users.
+
+ Currently this does three things:
+
+ 1. Drop the basic auth part. This is technically wrong since a server can
+ serve different content based on auth, but if it does that, it is even
+ impossible to guarantee two URLs without auth are equivalent, since
+ the user can input different auth information when prompted. So the
+ practical solution is to assume the auth doesn't affect the response.
+ 2. Parse the query to avoid the ordering issue. Note that ordering under the
+ same key in the query are NOT cleaned; i.e. "a=1&a=2" and "a=2&a=1" are
+ still considered different.
+ 3. Explicitly drop most of the fragment part, except ``subdirectory=`` and
+ hash values, since it should have no impact the downloaded content. Note
+ that this drops the "egg=" part historically used to denote the requested
+ project (and extras), which is wrong in the strictest sense, but too many
+ people are supplying it inconsistently to cause superfluous resolution
+ conflicts, so we choose to also ignore them.
+ """
+
+ parsed: urllib.parse.SplitResult
+ query: Dict[str, List[str]]
+ subdirectory: str
+ hashes: Dict[str, str]
+
+
+def _clean_link(link: Link) -> _CleanResult:
+ parsed = link._parsed_url
+ netloc = parsed.netloc.rsplit("@", 1)[-1]
+ # According to RFC 8089, an empty host in file: means localhost.
+ if parsed.scheme == "file" and not netloc:
+ netloc = "localhost"
+ fragment = urllib.parse.parse_qs(parsed.fragment)
+ if "egg" in fragment:
+ logger.debug("Ignoring egg= fragment in %s", link)
+ try:
+ # If there are multiple subdirectory values, use the first one.
+ # This matches the behavior of Link.subdirectory_fragment.
+ subdirectory = fragment["subdirectory"][0]
+ except (IndexError, KeyError):
+ subdirectory = ""
+ # If there are multiple hash values under the same algorithm, use the
+ # first one. This matches the behavior of Link.hash_value.
+ hashes = {k: fragment[k][0] for k in _SUPPORTED_HASHES if k in fragment}
+ return _CleanResult(
+ parsed=parsed._replace(netloc=netloc, query="", fragment=""),
+ query=urllib.parse.parse_qs(parsed.query),
+ subdirectory=subdirectory,
+ hashes=hashes,
+ )
+
+
+@functools.lru_cache(maxsize=None)
+def links_equivalent(link1: Link, link2: Link) -> bool:
+ return _clean_link(link1) == _clean_link(link2)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/search_scope.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/search_scope.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe61e8116b71e073351939ed7a499ee752398f1c
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/search_scope.py
@@ -0,0 +1,132 @@
+import itertools
+import logging
+import os
+import posixpath
+import urllib.parse
+from typing import List
+
+from pip._vendor.packaging.utils import canonicalize_name
+
+from pip._internal.models.index import PyPI
+from pip._internal.utils.compat import has_tls
+from pip._internal.utils.misc import normalize_path, redact_auth_from_url
+
+logger = logging.getLogger(__name__)
+
+
+class SearchScope:
+
+ """
+ Encapsulates the locations that pip is configured to search.
+ """
+
+ __slots__ = ["find_links", "index_urls", "no_index"]
+
+ @classmethod
+ def create(
+ cls,
+ find_links: List[str],
+ index_urls: List[str],
+ no_index: bool,
+ ) -> "SearchScope":
+ """
+ Create a SearchScope object after normalizing the `find_links`.
+ """
+ # Build find_links. If an argument starts with ~, it may be
+ # a local file relative to a home directory. So try normalizing
+ # it and if it exists, use the normalized version.
+ # This is deliberately conservative - it might be fine just to
+ # blindly normalize anything starting with a ~...
+ built_find_links: List[str] = []
+ for link in find_links:
+ if link.startswith("~"):
+ new_link = normalize_path(link)
+ if os.path.exists(new_link):
+ link = new_link
+ built_find_links.append(link)
+
+ # If we don't have TLS enabled, then WARN if anyplace we're looking
+ # relies on TLS.
+ if not has_tls():
+ for link in itertools.chain(index_urls, built_find_links):
+ parsed = urllib.parse.urlparse(link)
+ if parsed.scheme == "https":
+ logger.warning(
+ "pip is configured with locations that require "
+ "TLS/SSL, however the ssl module in Python is not "
+ "available."
+ )
+ break
+
+ return cls(
+ find_links=built_find_links,
+ index_urls=index_urls,
+ no_index=no_index,
+ )
+
+ def __init__(
+ self,
+ find_links: List[str],
+ index_urls: List[str],
+ no_index: bool,
+ ) -> None:
+ self.find_links = find_links
+ self.index_urls = index_urls
+ self.no_index = no_index
+
+ def get_formatted_locations(self) -> str:
+ lines = []
+ redacted_index_urls = []
+ if self.index_urls and self.index_urls != [PyPI.simple_url]:
+ for url in self.index_urls:
+ redacted_index_url = redact_auth_from_url(url)
+
+ # Parse the URL
+ purl = urllib.parse.urlsplit(redacted_index_url)
+
+ # URL is generally invalid if scheme and netloc is missing
+ # there are issues with Python and URL parsing, so this test
+ # is a bit crude. See bpo-20271, bpo-23505. Python doesn't
+ # always parse invalid URLs correctly - it should raise
+ # exceptions for malformed URLs
+ if not purl.scheme and not purl.netloc:
+ logger.warning(
+ 'The index url "%s" seems invalid, please provide a scheme.',
+ redacted_index_url,
+ )
+
+ redacted_index_urls.append(redacted_index_url)
+
+ lines.append(
+ "Looking in indexes: {}".format(", ".join(redacted_index_urls))
+ )
+
+ if self.find_links:
+ lines.append(
+ "Looking in links: {}".format(
+ ", ".join(redact_auth_from_url(url) for url in self.find_links)
+ )
+ )
+ return "\n".join(lines)
+
+ def get_index_urls_locations(self, project_name: str) -> List[str]:
+ """Returns the locations found via self.index_urls
+
+ Checks the url_name on the main (first in the list) index and
+ use this url_name to produce all locations
+ """
+
+ def mkurl_pypi_url(url: str) -> str:
+ loc = posixpath.join(
+ url, urllib.parse.quote(canonicalize_name(project_name))
+ )
+ # For maximum compatibility with easy_install, ensure the path
+ # ends in a trailing slash. Although this isn't in the spec
+ # (and PyPI can handle it without the slash) some other index
+ # implementations might break if they relied on easy_install's
+ # behavior.
+ if not loc.endswith("/"):
+ loc = loc + "/"
+ return loc
+
+ return [mkurl_pypi_url(url) for url in self.index_urls]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/selection_prefs.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/selection_prefs.py
new file mode 100644
index 0000000000000000000000000000000000000000..977bc4caa75c1e76156fa97e2841a01332f6fa47
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/selection_prefs.py
@@ -0,0 +1,51 @@
+from typing import Optional
+
+from pip._internal.models.format_control import FormatControl
+
+
+class SelectionPreferences:
+ """
+ Encapsulates the candidate selection preferences for downloading
+ and installing files.
+ """
+
+ __slots__ = [
+ "allow_yanked",
+ "allow_all_prereleases",
+ "format_control",
+ "prefer_binary",
+ "ignore_requires_python",
+ ]
+
+ # Don't include an allow_yanked default value to make sure each call
+ # site considers whether yanked releases are allowed. This also causes
+ # that decision to be made explicit in the calling code, which helps
+ # people when reading the code.
+ def __init__(
+ self,
+ allow_yanked: bool,
+ allow_all_prereleases: bool = False,
+ format_control: Optional[FormatControl] = None,
+ prefer_binary: bool = False,
+ ignore_requires_python: Optional[bool] = None,
+ ) -> None:
+ """Create a SelectionPreferences object.
+
+ :param allow_yanked: Whether files marked as yanked (in the sense
+ of PEP 592) are permitted to be candidates for install.
+ :param format_control: A FormatControl object or None. Used to control
+ the selection of source packages / binary packages when consulting
+ the index and links.
+ :param prefer_binary: Whether to prefer an old, but valid, binary
+ dist over a new source dist.
+ :param ignore_requires_python: Whether to ignore incompatible
+ "Requires-Python" values in links. Defaults to False.
+ """
+ if ignore_requires_python is None:
+ ignore_requires_python = False
+
+ self.allow_yanked = allow_yanked
+ self.allow_all_prereleases = allow_all_prereleases
+ self.format_control = format_control
+ self.prefer_binary = prefer_binary
+ self.ignore_requires_python = ignore_requires_python
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/wheel.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/wheel.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5dc12bdd63163c86f87ce4b5430cdb16d73769d
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35/lib/python3.12/site-packages/pip/_internal/models/wheel.py
@@ -0,0 +1,92 @@
+"""Represents a wheel file and provides access to the various parts of the
+name that have meaning.
+"""
+import re
+from typing import Dict, Iterable, List
+
+from pip._vendor.packaging.tags import Tag
+
+from pip._internal.exceptions import InvalidWheelFilename
+
+
+class Wheel:
+ """A wheel file"""
+
+ wheel_file_re = re.compile(
+ r"""^(?P(?P[^\s-]+?)-(?P[^\s-]*?))
+ ((-(?P\d[^-]*?))?-(?P[^\s-]+?)-(?P[^\s-]+?)-(?P[^\s-]+?)
+ \.whl|\.dist-info)$""",
+ re.VERBOSE,
+ )
+
+ def __init__(self, filename: str) -> None:
+ """
+ :raises InvalidWheelFilename: when the filename is invalid for a wheel
+ """
+ wheel_info = self.wheel_file_re.match(filename)
+ if not wheel_info:
+ raise InvalidWheelFilename(f"{filename} is not a valid wheel filename.")
+ self.filename = filename
+ self.name = wheel_info.group("name").replace("_", "-")
+ # we'll assume "_" means "-" due to wheel naming scheme
+ # (https://github.com/pypa/pip/issues/1150)
+ self.version = wheel_info.group("ver").replace("_", "-")
+ self.build_tag = wheel_info.group("build")
+ self.pyversions = wheel_info.group("pyver").split(".")
+ self.abis = wheel_info.group("abi").split(".")
+ self.plats = wheel_info.group("plat").split(".")
+
+ # All the tag combinations from this file
+ self.file_tags = {
+ Tag(x, y, z) for x in self.pyversions for y in self.abis for z in self.plats
+ }
+
+ def get_formatted_file_tags(self) -> List[str]:
+ """Return the wheel's tags as a sorted list of strings."""
+ return sorted(str(tag) for tag in self.file_tags)
+
+ def support_index_min(self, tags: List[Tag]) -> int:
+ """Return the lowest index that one of the wheel's file_tag combinations
+ achieves in the given list of supported tags.
+
+ For example, if there are 8 supported tags and one of the file tags
+ is first in the list, then return 0.
+
+ :param tags: the PEP 425 tags to check the wheel against, in order
+ with most preferred first.
+
+ :raises ValueError: If none of the wheel's file tags match one of
+ the supported tags.
+ """
+ try:
+ return next(i for i, t in enumerate(tags) if t in self.file_tags)
+ except StopIteration:
+ raise ValueError()
+
+ def find_most_preferred_tag(
+ self, tags: List[Tag], tag_to_priority: Dict[Tag, int]
+ ) -> int:
+ """Return the priority of the most preferred tag that one of the wheel's file
+ tag combinations achieves in the given list of supported tags using the given
+ tag_to_priority mapping, where lower priorities are more-preferred.
+
+ This is used in place of support_index_min in some cases in order to avoid
+ an expensive linear scan of a large list of tags.
+
+ :param tags: the PEP 425 tags to check the wheel against.
+ :param tag_to_priority: a mapping from tag to priority of that tag, where
+ lower is more preferred.
+
+ :raises ValueError: If none of the wheel's file tags match one of
+ the supported tags.
+ """
+ return min(
+ tag_to_priority[tag] for tag in self.file_tags if tag in tag_to_priority
+ )
+
+ def supported(self, tags: Iterable[Tag]) -> bool:
+ """Return whether the wheel is compatible with one of the given tags.
+
+ :param tags: the PEP 425 tags to check the wheel against.
+ """
+ return not self.file_tags.isdisjoint(tags)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c456a6b88378a5461e11dadfafb8820698b6b9d0
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 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 typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_m2m_100 import *
+ from .modeling_m2m_100 import *
+ from .tokenization_m2m_100 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/configuration_m2m_100.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/configuration_m2m_100.py
new file mode 100644
index 0000000000000000000000000000000000000000..41415686ffd9d60d978065608e9d0852cb55ee63
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/configuration_m2m_100.py
@@ -0,0 +1,75 @@
+# Copyright 2021 The Fairseq Authors and 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.
+"""M2M100 model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/m2m100_418M")
+@strict
+class M2M100Config(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import M2M100Config, M2M100Model
+
+ >>> # Initializing a M2M100 facebook/m2m100_418M style configuration
+ >>> configuration = M2M100Config()
+
+ >>> # Initializing a model (with random weights) from the facebook/m2m100_418M style configuration
+ >>> model = M2M100Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "m2m_100"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_attention_heads": "encoder_attention_heads",
+ "hidden_size": "d_model",
+ "num_hidden_layers": "encoder_layers",
+ }
+
+ vocab_size: int = 128112
+ max_position_embeddings: int = 1024
+ encoder_layers: int = 12
+ encoder_ffn_dim: int = 4096
+ encoder_attention_heads: int = 16
+ decoder_layers: int = 12
+ decoder_ffn_dim: int = 4096
+ decoder_attention_heads: int = 16
+ encoder_layerdrop: float | int = 0.05
+ decoder_layerdrop: float | int = 0.05
+ use_cache: bool = True
+ is_encoder_decoder: bool = True
+ activation_function: str = "relu"
+ d_model: int = 1024
+ dropout: float | int = 0.1
+ attention_dropout: float | int = 0.1
+ activation_dropout: float | int = 0.0
+ init_std: float = 0.02
+ decoder_start_token_id: int | None = 2
+ scale_embedding: bool = True
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["M2M100Config"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/modeling_m2m_100.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/modeling_m2m_100.py
new file mode 100644
index 0000000000000000000000000000000000000000..8985687435bc85680019fe957c653262bd9a7ca7
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/modeling_m2m_100.py
@@ -0,0 +1,923 @@
+# Copyright 2021 The Fairseq Authors and 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.
+"""PyTorch M2M100 model."""
+
+import math
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import (
+ FlashAttentionKwargs,
+)
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_m2m_100 import M2M100Config
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("self.model.config.pad_token_id has to be defined.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+# Copied from transformers.models.bart.modeling_bart.BartScaledWordEmbedding with Bart->M2M100
+class M2M100ScaledWordEmbedding(nn.Embedding):
+ """
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float | None = 1.0):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+class M2M100SinusoidalPositionalEmbedding(nn.Module):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None):
+ super().__init__()
+ self.offset = 2
+ self.num_positions = num_positions
+ self.embedding_dim = embedding_dim
+ self.padding_idx = padding_idx
+ self.make_weights(num_positions + self.offset, embedding_dim, padding_idx)
+
+ def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: int | None = None):
+ emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx)
+ if hasattr(self, "weights"):
+ # in forward put the weights on the correct dtype and device of the param
+ emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
+
+ self.register_buffer("weights", emb_weights, persistent=False)
+
+ @staticmethod
+ def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: int | None = None):
+ """
+ Build sinusoidal embeddings.
+
+ This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of
+ "Attention Is All You Need".
+ """
+ half_dim = embedding_dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
+ emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
+ if embedding_dim % 2 == 1:
+ # zero pad
+ emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+ if padding_idx is not None:
+ emb[padding_idx, :] = 0
+
+ return emb.to(torch.get_default_dtype())
+
+ @torch.no_grad()
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ past_key_values_length: int = 0,
+ ):
+ if input_ids is not None:
+ bsz, seq_len = input_ids.size()
+ # Create the position ids from the input token ids. Any padded tokens remain padded.
+ position_ids = self.create_position_ids_from_input_ids(
+ input_ids, self.padding_idx, past_key_values_length
+ ).to(input_ids.device)
+ else:
+ bsz, seq_len = inputs_embeds.size()[:-1]
+ position_ids = self.create_position_ids_from_inputs_embeds(
+ inputs_embeds, past_key_values_length, self.padding_idx
+ )
+
+ # expand embeddings if needed
+ max_pos = self.padding_idx + 1 + seq_len + past_key_values_length
+ if max_pos > self.weights.size(0):
+ self.make_weights(max_pos + self.offset, self.embedding_dim, self.padding_idx)
+
+ return self.weights.index_select(0, position_ids.view(-1)).view(bsz, seq_len, self.weights.shape[-1]).detach()
+
+ @staticmethod
+ def create_position_ids_from_inputs_embeds(inputs_embeds, past_key_values_length, padding_idx):
+ """
+ We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+ Args:
+ inputs_embeds: torch.Tensor
+
+ Returns: torch.Tensor
+ """
+ input_shape = inputs_embeds.size()[:-1]
+ sequence_length = input_shape[1]
+
+ position_ids = torch.arange(
+ padding_idx + 1, sequence_length + padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+ )
+ return position_ids.unsqueeze(0).expand(input_shape).contiguous() + past_key_values_length
+
+ @staticmethod
+ # Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings.create_position_ids_from_input_ids
+ def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
+ """
+ Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+ are ignored. This is modified from fairseq's `utils.make_positions`.
+
+ Args:
+ x: torch.Tensor x:
+
+ Returns: torch.Tensor
+ """
+ # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+ mask = input_ids.ne(padding_idx).int()
+ incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
+ return incremental_indices.long() + padding_idx
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->M2M100
+class M2M100Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ is_causal: bool = False,
+ config: M2M100Config | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if layer_idx is None and self.is_decoder:
+ logger.warning_once(
+ f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+ "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ input_shape = hidden_states.shape[:-1]
+
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ kv_shape = (*current_states.shape[:-1], -1, self.head_dim)
+ key_states = key_states.view(kv_shape).transpose(1, 2)
+ value_states = value_states.view(kv_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->M2M100, MBART->M2M100
+class M2M100EncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: M2M100Config):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = M2M100Attention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ config=config,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16:
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartDecoderLayer with MBart->M2M100, MBART->M2M100
+class M2M100DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: M2M100Config, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = M2M100Attention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = M2M100Attention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+ self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ encoder_hidden_states (`torch.FloatTensor`):
+ cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
+ encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class M2M100PreTrainedModel(PreTrainedModel):
+ config: M2M100Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["M2M100EncoderLayer", "M2M100DecoderLayer"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ # Doesn't support `compile` (dynamic control flow). Can be fixed but low usage model
+ _can_compile_fullgraph = False
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, M2M100SinusoidalPositionalEmbedding):
+ emb_weights = module.get_embedding(
+ module.num_positions + module.offset, module.embedding_dim, module.padding_idx
+ )
+ init.copy_(module.weights, emb_weights)
+
+
+class M2M100Encoder(M2M100PreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`M2M100EncoderLayer`].
+
+ Args:
+ config: M2M100Config
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": M2M100EncoderLayer,
+ "attentions": M2M100Attention,
+ }
+
+ def __init__(self, config: M2M100Config):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+
+ embed_dim = config.d_model
+ self.padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ self.embed_tokens = M2M100ScaledWordEmbedding(
+ config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = M2M100SinusoidalPositionalEmbedding(
+ config.max_position_embeddings,
+ embed_dim,
+ self.padding_idx,
+ )
+ self.layers = nn.ModuleList([M2M100EncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ embed_pos = self.embed_positions(input_ids, inputs_embeds)
+ embed_pos = embed_pos.to(inputs_embeds.device)
+
+ hidden_states = inputs_embeds + embed_pos
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ for idx, encoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class M2M100Decoder(M2M100PreTrainedModel):
+ """
+ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`M2M100DecoderLayer`]
+
+ Args:
+ config: M2M100Config
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": M2M100DecoderLayer,
+ "attentions": OutputRecorder(M2M100Attention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(M2M100Attention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: M2M100Config):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ self.padding_idx = config.pad_token_id
+ self.max_target_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+ self.embed_tokens = M2M100ScaledWordEmbedding(
+ config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = M2M100SinusoidalPositionalEmbedding(
+ config.max_position_embeddings,
+ config.d_model,
+ self.padding_idx,
+ )
+ self.layers = nn.ModuleList([M2M100DecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)])
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize `past_key_values`
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ if attention_mask is None and not is_torchdynamo_compiling():
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = (
+ past_key_values.self_attention_cache
+ if isinstance(past_key_values, EncoderDecoderCache)
+ else past_key_values
+ )
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ # embed positions
+ positions = self.embed_positions(input_ids, inputs_embeds, past_key_values_length)
+ positions = positions.to(inputs_embeds.device)
+
+ hidden_states = inputs_embeds + positions
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class M2M100Model(M2M100PreTrainedModel):
+ _tied_weights_keys = {
+ "decoder.embed_tokens.weight": "shared.weight",
+ "encoder.embed_tokens.weight": "shared.weight",
+ }
+
+ def __init__(self, config: M2M100Config):
+ super().__init__(config)
+
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+ self.shared = M2M100ScaledWordEmbedding(vocab_size, config.d_model, padding_idx, embed_scale=embed_scale)
+
+ self.encoder = M2M100Encoder(config)
+ self.decoder = M2M100Decoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.shared
+
+ def set_input_embeddings(self, value):
+ self.shared = value
+ self.encoder.embed_tokens = self.shared
+ self.decoder.embed_tokens = self.shared
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | Seq2SeqModelOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ M2M100 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+ """
+ if encoder_outputs is None:
+ encoder_outputs = self.encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn)
+ decoder_outputs = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs[0],
+ encoder_attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The M2M100 Model with a language modeling head. Can be used for summarization.
+ """
+)
+class M2M100ForConditionalGeneration(M2M100PreTrainedModel, GenerationMixin):
+ base_model_prefix = "model"
+ _tied_weights_keys = {"lm_head.weight": "model.shared.weight"}
+
+ def __init__(self, config: M2M100Config):
+ super().__init__(config)
+ self.model = M2M100Model(config)
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | Seq2SeqLMOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ M2M100 uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example Translation:
+
+ ```python
+ >>> from transformers import AutoTokenizer, M2M100ForConditionalGeneration
+
+ >>> model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
+
+ >>> text_to_translate = "Life is like a box of chocolates"
+ >>> model_inputs = tokenizer(text_to_translate, return_tensors="pt")
+
+ >>> # translate to French
+ >>> gen_tokens = model.generate(**model_inputs, forced_bos_token_id=tokenizer.get_lang_id("fr"))
+ >>> print(tokenizer.batch_decode(gen_tokens, skip_special_tokens=True))
+ ```
+ """
+ if labels is not None:
+ if decoder_input_ids is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ lm_logits = self.lm_head(outputs.last_hidden_state)
+
+ masked_lm_loss = None
+ if labels is not None:
+ # move labels to the correct device to enable PP
+ labels = labels.to(lm_logits.device)
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return Seq2SeqLMOutput(
+ loss=masked_lm_loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+__all__ = ["M2M100ForConditionalGeneration", "M2M100Model", "M2M100PreTrainedModel"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/tokenization_m2m_100.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/tokenization_m2m_100.py
new file mode 100644
index 0000000000000000000000000000000000000000..b5d7f26c126b5c79d7abdde9d09c77397bfb08c0
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/m2m_100/tokenization_m2m_100.py
@@ -0,0 +1,384 @@
+# Copyright 2021 The Fairseq Authors and 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.
+"""Tokenization classes for M2M100."""
+
+import json
+import os
+from pathlib import Path
+from shutil import copyfile
+from typing import Any
+
+import sentencepiece
+
+from ...tokenization_python import BatchEncoding, PreTrainedTokenizer
+from ...utils import logging
+from ...utils.import_utils import requires
+
+
+logger = logging.get_logger(__name__)
+
+SPIECE_UNDERLINE = "▁"
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.json",
+ "spm_file": "sentencepiece.bpe.model",
+ "tokenizer_config_file": "tokenizer_config.json",
+}
+
+
+# fmt: off
+FAIRSEQ_LANGUAGE_CODES = {
+ "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"],
+ "wmt21": ['en', 'ha', 'is', 'ja', 'cs', 'ru', 'zh', 'de']
+}
+# fmt: on
+
+
+@requires(backends=("sentencepiece",))
+class M2M100Tokenizer(PreTrainedTokenizer):
+ """
+ Construct an M2M100 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ spm_file (`str`):
+ Path to [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
+ contains the vocabulary.
+ src_lang (`str`, *optional*):
+ A string representing the source language.
+ tgt_lang (`str`, *optional*):
+ A string representing the target language.
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ language_codes (`str`, *optional*, defaults to `"m2m100"`):
+ What language codes to use. Should be one of `"m2m100"` or `"wmt21"`.
+ sp_model_kwargs (`dict`, *optional*):
+ Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+ SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+ to set:
+
+ - `enable_sampling`: Enable subword regularization.
+ - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+ - `nbest_size = {0,1}`: No sampling is performed.
+ - `nbest_size > 1`: samples from the nbest_size results.
+ - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+ using forward-filtering-and-backward-sampling algorithm.
+
+ - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+ BPE-dropout.
+
+ Examples:
+
+ ```python
+ >>> from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+
+ >>> model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+ >>> tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", src_lang="en", tgt_lang="ro")
+ >>> src_text = " UN Chief Says There Is No Military Solution in Syria"
+ >>> tgt_text = "Şeful ONU declară că nu există o soluţie militară în Siria"
+ >>> model_inputs = tokenizer(src_text, text_target=tgt_text, return_tensors="pt")
+ >>> outputs = model(**model_inputs) # should work
+ ```"""
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+
+ prefix_tokens: list[int] = []
+ suffix_tokens: list[int] = []
+
+ def __init__(
+ self,
+ vocab_file,
+ spm_file,
+ src_lang=None,
+ tgt_lang=None,
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ pad_token="",
+ unk_token="",
+ language_codes="m2m100",
+ sp_model_kwargs: dict[str, Any] | None = None,
+ num_madeup_words=8,
+ **kwargs,
+ ) -> None:
+ self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+
+ self.language_codes = language_codes
+ fairseq_language_code = FAIRSEQ_LANGUAGE_CODES[language_codes]
+ self.lang_code_to_token = {lang_code: f"__{lang_code}__" for lang_code in fairseq_language_code}
+
+ additional_special_tokens = kwargs.pop("additional_special_tokens", [])
+ for lang_code in fairseq_language_code:
+ token = self.get_lang_token(lang_code)
+ if token not in additional_special_tokens and lang_code not in str(token) not in self.added_tokens_encoder:
+ additional_special_tokens.append(token)
+
+ self.vocab_file = vocab_file
+ self.encoder = load_json(vocab_file)
+ self.decoder = {v: k for k, v in self.encoder.items()}
+ self.spm_file = spm_file
+ self.sp_model = load_spm(spm_file, self.sp_model_kwargs)
+
+ self.encoder_size = len(self.encoder)
+
+ self.lang_token_to_id = {
+ self.get_lang_token(lang_code): self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)
+ }
+ self.lang_code_to_id = {lang_code: self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)}
+ self.id_to_lang_token = {v: k for k, v in self.lang_token_to_id.items()}
+
+ self._src_lang = src_lang if src_lang is not None else "en"
+ self.tgt_lang = tgt_lang
+ self.cur_lang_id = self.get_lang_id(self._src_lang)
+
+ self.num_madeup_words = num_madeup_words
+
+ super().__init__(
+ src_lang=src_lang,
+ tgt_lang=tgt_lang,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ language_codes=language_codes,
+ sp_model_kwargs=self.sp_model_kwargs,
+ additional_special_tokens=additional_special_tokens,
+ num_madeup_words=num_madeup_words,
+ **kwargs,
+ )
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.encoder)
+
+ def get_vocab(self) -> dict:
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+ vocab.update(self.added_tokens_encoder)
+ return vocab
+
+ @property
+ def src_lang(self) -> str:
+ return self._src_lang
+
+ @src_lang.setter
+ def src_lang(self, new_src_lang: str) -> None:
+ self._src_lang = new_src_lang
+ self.set_src_lang_special_tokens(self._src_lang)
+
+ def _tokenize(self, text: str) -> list[str]:
+ return self.sp_model.encode(text, out_type=str)
+
+ def _convert_token_to_id(self, token):
+ if token in self.lang_token_to_id:
+ return self.lang_token_to_id[token]
+ return self.encoder.get(token, self.encoder[self.unk_token])
+
+ def _convert_id_to_token(self, index: int) -> str:
+ """Converts an index (integer) in a token (str) using the decoder."""
+ if index in self.id_to_lang_token:
+ return self.id_to_lang_token[index]
+ return self.decoder.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ current_sub_tokens = []
+ out_string = ""
+ for token in tokens:
+ # make sure that special tokens are not decoded using sentencepiece model
+ if token in self.all_special_tokens:
+ out_string += self.sp_model.decode(current_sub_tokens) + token
+ current_sub_tokens = []
+ else:
+ current_sub_tokens.append(token)
+ out_string += self.sp_model.decode(current_sub_tokens)
+ return out_string.strip()
+
+ def get_special_tokens_mask(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False
+ ) -> list[int]:
+ """
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+ special tokens using the tokenizer `prepare_for_model` method.
+
+ Args:
+ token_ids_0 (`list[int]`):
+ List of IDs.
+ token_ids_1 (`list[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not the token list is already formatted with special tokens for the model.
+
+ Returns:
+ `list[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+ """
+
+ if already_has_special_tokens:
+ return super().get_special_tokens_mask(
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+ )
+
+ prefix_ones = [1] * len(self.prefix_tokens)
+ suffix_ones = [1] * len(self.suffix_tokens)
+ if token_ids_1 is None:
+ return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
+ return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
+
+ def build_inputs_with_special_tokens(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+ ) -> list[int]:
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. An MBART sequence has the following format, where `X` represents the sequence:
+
+ - `input_ids` (for encoder) `X [eos, src_lang_code]`
+ - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
+
+ BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
+ separator.
+
+ Args:
+ token_ids_0 (`list[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`list[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `list[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return self.prefix_tokens + token_ids_0 + self.suffix_tokens
+ # We don't expect to process pairs, but leave the pair logic for API consistency
+ return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
+
+ def __getstate__(self) -> dict:
+ state = self.__dict__.copy()
+ state["sp_model"] = None
+ return state
+
+ def __setstate__(self, d: dict) -> None:
+ self.__dict__ = d
+
+ # for backward compatibility
+ if not hasattr(self, "sp_model_kwargs"):
+ self.sp_model_kwargs = {}
+
+ self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ save_dir = Path(save_directory)
+ if not save_dir.is_dir():
+ raise OSError(f"{save_directory} should be a directory")
+ vocab_save_path = save_dir / (
+ (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
+ )
+ spm_save_path = save_dir / (
+ (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
+ )
+
+ save_json(self.encoder, vocab_save_path)
+
+ if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file):
+ copyfile(self.spm_file, spm_save_path)
+ elif not os.path.isfile(self.spm_file):
+ with open(spm_save_path, "wb") as fi:
+ content_spiece_model = self.sp_model.serialized_model_proto()
+ fi.write(content_spiece_model)
+
+ return (str(vocab_save_path), str(spm_save_path))
+
+ def prepare_seq2seq_batch(
+ self,
+ src_texts: list[str],
+ src_lang: str = "en",
+ tgt_texts: list[str] | None = None,
+ tgt_lang: str = "ro",
+ **kwargs,
+ ) -> BatchEncoding:
+ self.src_lang = src_lang
+ self.tgt_lang = tgt_lang
+ self.set_src_lang_special_tokens(self.src_lang)
+ return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
+
+ def _build_translation_inputs(self, raw_inputs, src_lang: str | None, tgt_lang: str | None, **extra_kwargs):
+ """Used by translation pipeline, to prepare inputs for the generate function"""
+ if src_lang is None or tgt_lang is None:
+ raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model")
+ self.src_lang = src_lang
+ inputs = self(raw_inputs, add_special_tokens=True, **extra_kwargs)
+ tgt_lang_id = self.get_lang_id(tgt_lang)
+ inputs["forced_bos_token_id"] = tgt_lang_id
+ return inputs
+
+ def _switch_to_input_mode(self):
+ self.set_src_lang_special_tokens(self.src_lang)
+
+ def _switch_to_target_mode(self):
+ self.set_tgt_lang_special_tokens(self.tgt_lang)
+
+ def set_src_lang_special_tokens(self, src_lang: str) -> None:
+ """Reset the special tokens to the source lang setting. No prefix and suffix=[eos, src_lang_code]."""
+ lang_token = self.get_lang_token(src_lang)
+ self.cur_lang_id = self.lang_token_to_id[lang_token]
+ self.prefix_tokens = [self.cur_lang_id]
+ self.suffix_tokens = [self.eos_token_id]
+
+ def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None:
+ """Reset the special tokens to the target language setting. No prefix and suffix=[eos, tgt_lang_code]."""
+ lang_token = self.get_lang_token(tgt_lang)
+ self.cur_lang_id = self.lang_token_to_id[lang_token]
+ self.prefix_tokens = [self.cur_lang_id]
+ self.suffix_tokens = [self.eos_token_id]
+
+ def get_lang_token(self, lang: str) -> str:
+ return self.lang_code_to_token[lang]
+
+ def get_lang_id(self, lang: str) -> int:
+ lang_token = self.get_lang_token(lang)
+ return self.lang_token_to_id[lang_token]
+
+
+def load_spm(path: str, sp_model_kwargs: dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
+ spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
+ spm.Load(str(path))
+ return spm
+
+
+def load_json(path: str) -> dict | list:
+ with open(path, "r") as f:
+ return json.load(f)
+
+
+def save_json(data, path: str) -> None:
+ with open(path, "w") as f:
+ json.dump(data, f, indent=2)
+
+
+__all__ = ["M2M100Tokenizer"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6610a71fa701d0f8854c96ed5b63ef389666cc90
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/__init__.py
@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_nanochat import *
+ from .modeling_nanochat import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/configuration_nanochat.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/configuration_nanochat.py
new file mode 100644
index 0000000000000000000000000000000000000000..24a0ab7b6d097afbda06b1a0a9977cb43240d2ce
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/configuration_nanochat.py
@@ -0,0 +1,81 @@
+# 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 huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="karpathy/nanochat-d32")
+@strict
+class NanoChatConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import NanoChatModel, NanoChatConfig
+
+ >>> # Initializing a NanoChat style configuration
+ >>> configuration = NanoChatConfig()
+
+ >>> # Initializing a model from the NanoChat style configuration
+ >>> model = NanoChatModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "nanochat"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.fc1": "colwise",
+ "layers.*.mlp.fc2": "rowwise",
+ }
+
+ vocab_size: int = 50304
+ hidden_size: int = 768
+ intermediate_size: int = 8192
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 6
+ num_key_value_heads: int | None = None
+ max_position_embeddings: int = 2048
+ hidden_act: str = "relu2"
+ attention_dropout: float | int = 0.0
+ rms_norm_eps: float = 1e-6
+ initializer_range: float = 0.02
+ rope_parameters: RopeParameters | dict | None = None
+ use_cache: bool = True
+ final_logit_softcapping: float | None = 15.0
+ attention_bias: bool = False
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ pad_token_id: int | None = 1
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["NanoChatConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modeling_nanochat.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modeling_nanochat.py
new file mode 100644
index 0000000000000000000000000000000000000000..9205b89cd360e1363fcda18dc7826adf01dd2724
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modeling_nanochat.py
@@ -0,0 +1,518 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/nanochat/modular_nanochat.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_nanochat.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# 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 math
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_nanochat import NanoChatConfig
+
+
+class NanoChatRMSNorm(torch.nn.Module):
+ def __init__(self, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ return self._norm(x.float()).type_as(x)
+
+ def extra_repr(self):
+ return f"eps={self.eps}"
+
+
+class NanoChatRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: NanoChatConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: NanoChatConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input with flipped signs for NanoChat."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((x2, -x1), dim=-1)
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class NanoChatAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: NanoChatConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+
+ self.q_norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+ self.k_norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ # RoPE -> Norm (instead of usual Norm -> RoPE)
+ query_states = self.q_norm(query_states)
+ key_states = self.k_norm(key_states)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class NanoChatMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class NanoChatDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: NanoChatConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = NanoChatAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = NanoChatMLP(config)
+
+ self.input_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+ self.post_attention_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+@auto_docstring
+class NanoChatPreTrainedModel(PreTrainedModel):
+ config: NanoChatConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["NanoChatDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": NanoChatDecoderLayer,
+ "attentions": NanoChatAttention,
+ }
+
+ def _init_weights(self, module: nn.Module) -> None:
+ super()._init_weights(module)
+ if isinstance(module, NanoChatAttention):
+ init.normal_(
+ module.o_proj.weight,
+ mean=0.0,
+ std=self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers),
+ )
+
+
+@auto_docstring
+class NanoChatModel(NanoChatPreTrainedModel):
+ def __init__(self, config: NanoChatConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [NanoChatDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+ self.norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+ self.rotary_emb = NanoChatRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ hidden_states = self.norm(hidden_states) # Additional norm before the layers
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class NanoChatForCausalLM(NanoChatPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = NanoChatModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+ >>> model = AutoModelForCausalLM.from_pretrained("karpathy/nanochat-d32")
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("karpathy/nanochat-d32")
+
+ >>> conversation = [
+ {"role": "user", "content": "What is the capital of France?"},
+ ]
+
+ >>> inputs = tokenizer.apply_chat_template(
+ conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
+ ).to(device)
+
+ >>> with torch.no_grad():
+ >>> outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False)
+
+ >>> generated_tokens = outputs[0, inputs["input_ids"].shape[1] :]
+ >>> output = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+ ```"""
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ if self.config.final_logit_softcapping is not None:
+ logits = logits / self.config.final_logit_softcapping
+ logits = torch.tanh(logits)
+ logits = logits * self.config.final_logit_softcapping
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["NanoChatPreTrainedModel", "NanoChatModel", "NanoChatForCausalLM"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modular_nanochat.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modular_nanochat.py
new file mode 100644
index 0000000000000000000000000000000000000000..713cc29b81eb304e2d92f506294021858bb58c92
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/nanochat/modular_nanochat.py
@@ -0,0 +1,235 @@
+# 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 math
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...masking_utils import create_causal_mask
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ..clip.modeling_clip import CLIPMLP
+from ..gemma2.modeling_gemma2 import Gemma2ForCausalLM
+from ..llama.modeling_llama import (
+ LlamaDecoderLayer,
+ LlamaModel,
+ LlamaPreTrainedModel,
+ LlamaRotaryEmbedding,
+ apply_rotary_pos_emb,
+ eager_attention_forward,
+)
+from ..llama4.modeling_llama4 import Llama4TextL2Norm
+from ..qwen3.modeling_qwen3 import Qwen3Attention
+from .configuration_nanochat import NanoChatConfig
+
+
+class NanoChatRMSNorm(Llama4TextL2Norm):
+ pass
+
+
+class NanoChatRotaryEmbedding(LlamaRotaryEmbedding):
+ pass
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input with flipped signs for NanoChat."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((x2, -x1), dim=-1)
+
+
+class NanoChatAttention(Qwen3Attention):
+ def __init__(self, config: NanoChatConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+ del self.sliding_window
+ del self.layer_type
+
+ self.q_norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+ self.k_norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ # RoPE -> Norm (instead of usual Norm -> RoPE)
+ query_states = self.q_norm(query_states)
+ key_states = self.k_norm(key_states)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class NanoChatMLP(CLIPMLP):
+ def __init__(self, config):
+ super().__init__(config)
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+
+class NanoChatDecoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: NanoChatConfig, layer_idx: int):
+ super().__init__()
+
+ self.input_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+ self.post_attention_layernorm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+
+
+@auto_docstring
+class NanoChatPreTrainedModel(LlamaPreTrainedModel):
+ def _init_weights(self, module: nn.Module) -> None:
+ PreTrainedModel._init_weights(self, module)
+ if isinstance(module, NanoChatAttention):
+ init.normal_(
+ module.o_proj.weight,
+ mean=0.0,
+ std=self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers),
+ )
+
+
+@auto_docstring
+class NanoChatModel(LlamaModel):
+ def __init__(self, config: NanoChatConfig):
+ super().__init__(config)
+
+ self.norm = NanoChatRMSNorm(eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ hidden_states = self.norm(hidden_states) # Additional norm before the layers
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class NanoChatForCausalLM(Gemma2ForCausalLM):
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+
+ def forward(self, **super_kwargs) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+ >>> model = AutoModelForCausalLM.from_pretrained("karpathy/nanochat-d32")
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("karpathy/nanochat-d32")
+
+ >>> conversation = [
+ {"role": "user", "content": "What is the capital of France?"},
+ ]
+
+ >>> inputs = tokenizer.apply_chat_template(
+ conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
+ ).to(device)
+
+ >>> with torch.no_grad():
+ >>> outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False)
+
+ >>> generated_tokens = outputs[0, inputs["input_ids"].shape[1] :]
+ >>> output = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+ ```"""
+ super().forward(**super_kwargs)
+
+
+__all__ = [
+ "NanoChatPreTrainedModel",
+ "NanoChatModel",
+ "NanoChatForCausalLM",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fabf00e524e6b096b2ca33324c3f38b5b444a4e1
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 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 typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_qwen3_5_moe import *
+ from .modeling_qwen3_5_moe import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6f9594e0d73f7c088cd451c1b57910dfa10da84
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py
@@ -0,0 +1,197 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen3_5_moe.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The Qwen Team and 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 huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3.5-35B-A3B")
+@strict
+class Qwen3_5MoeTextConfig(PreTrainedConfig):
+ r"""
+ linear_conv_kernel_dim (`int`, *optional*, defaults to 4):
+ Kernel size of the convolution used in linear attention layers.
+ linear_key_head_dim (`int`, *optional*, defaults to 128):
+ Dimension of each key head in linear attention.
+ linear_value_head_dim (`int`, *optional*, defaults to 128):
+ Dimension of each value head in linear attention.
+ linear_num_key_heads (`int`, *optional*, defaults to 16):
+ Number of key heads used in linear attention layers.
+ linear_num_value_heads (`int`, *optional*, defaults to 32):
+ Number of value heads used in linear attention layers.
+
+ ```python
+ >>> from transformers import Qwen3_5MoeTextModel, Qwen3_5MoeTextConfig
+
+ >>> # Initializing a Qwen3.5-MoE style configuration
+ >>> configuration = Qwen3_5MoeTextConfig()
+
+ >>> # Initializing a model from the Qwen3.5-35B-A3B style configuration
+ >>> model = Qwen3_5MoeTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "qwen3_5_moe_text"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+ "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+ "layers.*.mlp.experts.down_proj": "rowwise",
+ "layers.*.mlp.experts": "moe_tp_experts",
+ "layers.*.mlp.shared_expert.gate_proj": "colwise",
+ "layers.*.mlp.shared_expert.up_proj": "colwise",
+ "layers.*.mlp.shared_expert.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ vocab_size: int = 248320
+ hidden_size: int = 2048
+ num_hidden_layers: int = 40
+ num_attention_heads: int = 16
+ num_key_value_heads: int = 2
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 32768
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+ head_dim: int = 256
+ linear_conv_kernel_dim: int = 4
+ linear_key_head_dim: int = 128
+ linear_value_head_dim: int = 128
+ linear_num_key_heads: int = 16
+ linear_num_value_heads: int = 32
+ moe_intermediate_size: int = 512
+ shared_expert_intermediate_size: int = 512
+ num_experts_per_tok: int = 8
+ num_experts: int = 256
+ output_router_logits: bool = False
+ router_aux_loss_coef: float = 0.001
+ layer_types: list[str] | None = None
+ pad_token_id: int | None = None
+ bos_token_id: int | None = None
+ eos_token_id: int | list[int] | None = None
+ base_config_key = "text_config"
+ ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
+
+ def __post_init__(self, **kwargs):
+ kwargs.setdefault("partial_rotary_factor", 0.25) # assign default for BC
+ if self.layer_types is None:
+ interval_pattern = kwargs.pop("full_attention_interval", 4)
+ self.layer_types = [
+ "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3.5-35B-A3B")
+@strict
+class Qwen3_5MoeVisionConfig(PreTrainedConfig):
+ r"""
+ out_hidden_size (`int`, *optional*, defaults to 3584):
+ The output hidden size of the vision model.
+ num_position_embeddings (`int`, *optional*, defaults to 2304):
+ The maximum sequence length that this model might ever be used with
+ """
+
+ model_type = "qwen3_5_moe_vision"
+ base_config_key = "vision_config"
+
+ depth: int = 27
+ hidden_size: int = 1152
+ hidden_act: str = "gelu_pytorch_tanh"
+ intermediate_size: int = 4304
+ num_heads: int = 16
+ in_channels: int = 3
+ patch_size: int | list[int] | tuple[int, int] = 16
+ spatial_merge_size: int = 2
+ temporal_patch_size: int | list[int] | tuple[int, int] = 2
+ out_hidden_size: int = 3584
+ num_position_embeddings: int = 2304
+ initializer_range: float = 0.02
+
+
+@auto_docstring(checkpoint="Qwen/Qwen3.5-35B-A3B")
+@strict
+class Qwen3_5MoeConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Qwen3_5MoeForConditionalGeneration, Qwen3_5MoeConfig
+
+ >>> # Initializing a Qwen3.5-MoE style configuration
+ >>> configuration = Qwen3_5MoeConfig()
+
+ >>> # Initializing a model from the Qwen3.5-35B-A3B style configuration
+ >>> model = Qwen3_5MoeForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen3_5_moe"
+ sub_configs = {"vision_config": Qwen3_5MoeVisionConfig, "text_config": Qwen3_5MoeTextConfig}
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+
+ image_token_id: int = 248056
+ video_token_id: int = 248057
+ vision_start_token_id: int = 248053
+ vision_end_token_id: int = 248054
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"]()
+
+ if isinstance(self.text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = self.sub_configs["text_config"]()
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig", "Qwen3_5MoeVisionConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbf459287d0e4c464c98c7c1e69e3aa138cd3e65
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py
@@ -0,0 +1,2325 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen3_5_moe/modular_qwen3_5_moe.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen3_5_moe.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The Qwen Team and 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 itertools
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_experts_implementation, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPast,
+ BaseModelOutputWithPooling,
+ ModelOutput,
+ MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check
+from ...utils.generic import (
+ accepts_precomputed_kwargs,
+ is_flash_attention_requested,
+ maybe_autocast,
+ merge_with_config_defaults,
+)
+from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ...vision_utils import get_vision_bilinear_indices_and_weights, get_vision_cu_seqlens, get_vision_position_ids
+from .configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5MoeVisionConfig
+
+
+if is_causal_conv1d_available():
+ from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+else:
+ causal_conv1d_update, causal_conv1d_fn = None, None
+
+if is_flash_linear_attention_available():
+ from fla.modules import FusedRMSNormGated
+ from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
+else:
+ chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None
+ FusedRMSNormGated = None
+
+logger = logging.get_logger(__name__)
+
+
+class Qwen3_5MoeVisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+ return (position_ids.unsqueeze(-1) * self.inv_freq).flatten(1)
+
+
+class Qwen3_5MoeTextRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen3_5MoeTextConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+ self.mrope_section = config.rope_parameters.get("mrope_section", [11, 11, 10])
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Qwen3_5MoeTextConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, Qwen3_5Moe has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+ inv_freq_expanded = (
+ self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device)
+ )
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+ def apply_interleaved_mrope(self, freqs, mrope_section):
+ """Apply interleaved MRoPE to 3D rotary embeddings.
+ Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
+ interleaved [THWTHWTHW...TT], preserving frequency continuity.
+ args:
+ x: (3, bs, seq_len, head_dim // 2)
+ mrope_section: (3,)
+ returns:
+ x_t: (bs, seq_len, head_dim // 2)
+ """
+ freqs_t = freqs[0] # just overwrite the first dimension T
+ for dim, offset in enumerate((1, 2), start=1): # H, W
+ length = mrope_section[dim] * 3
+ idx = slice(offset, length, 3)
+ freqs_t[..., idx] = freqs[dim, ..., idx]
+ return freqs_t
+
+
+class Qwen3_5MoeRMSNormGated(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6, **kwargs):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states, gate=None):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ # Norm before gate
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ hidden_states = self.weight * hidden_states.to(input_dtype)
+ hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+
+ return hidden_states.to(input_dtype)
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+ """
+ Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+ """
+ # NOTE: attention mask is a 2D boolean tensor
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ dtype = hidden_states.dtype
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+ return hidden_states
+
+
+is_fast_path_available = all(
+ (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)
+)
+
+
+def torch_causal_conv1d_update(
+ hidden_states,
+ conv_state,
+ weight,
+ bias=None,
+ activation=None,
+):
+ _, hidden_size, seq_len = hidden_states.shape
+ state_len = conv_state.shape[-1]
+
+ hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
+ conv_state.copy_(hidden_states_new[:, :, -state_len:])
+ out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
+ out = F.silu(out[:, :, -seq_len:])
+ out = out.to(hidden_states.dtype)
+ return out
+
+
+def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
+ """This function is intended to align with the l2norm implementation in the FLA library."""
+ inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+ return x * inv_norm
+
+
+def torch_chunk_gated_delta_rule(
+ query,
+ key,
+ value,
+ g,
+ beta,
+ chunk_size=64,
+ initial_state=None,
+ output_final_state=False,
+ use_qk_l2norm_in_kernel=False,
+ **kwargs,
+):
+ initial_dtype = query.dtype
+ if use_qk_l2norm_in_kernel:
+ query = l2norm(query, dim=-1, eps=1e-6)
+ key = l2norm(key, dim=-1, eps=1e-6)
+ query, key, value, beta, g = [
+ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
+ ]
+
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
+ v_head_dim = value.shape[-1]
+ pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
+ query = F.pad(query, (0, 0, 0, pad_size))
+ key = F.pad(key, (0, 0, 0, pad_size))
+ value = F.pad(value, (0, 0, 0, pad_size))
+ beta = F.pad(beta, (0, pad_size))
+ g = F.pad(g, (0, pad_size))
+ total_sequence_length = sequence_length + pad_size
+ scale = 1 / (query.shape[-1] ** 0.5)
+ query = query * scale
+
+ v_beta = value * beta.unsqueeze(-1)
+ k_beta = key * beta.unsqueeze(-1)
+ # reshape to chunks
+ query, key, value, k_beta, v_beta = [
+ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta)
+ ]
+ g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
+ mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
+
+ # chunk decay
+ g = g.cumsum(dim=-1)
+ decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
+ attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
+ for i in range(1, chunk_size):
+ row = attn[..., i, :i].clone()
+ sub = attn[..., :i, :i].clone()
+ attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
+ attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
+ value = attn @ v_beta
+ k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
+ last_recurrent_state = (
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim, dtype=value.dtype, device=value.device)
+ if initial_state is None
+ else initial_state.to(value)
+ )
+ core_attn_out = torch.zeros_like(value)
+ mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
+
+ # for each chunk
+ for i in range(0, total_sequence_length // chunk_size):
+ q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
+ attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]
+ v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
+ v_new = v_i - v_prime
+ attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
+ core_attn_out[:, :, i] = attn_inter + attn @ v_new
+ last_recurrent_state = (
+ last_recurrent_state * g[:, :, i, -1, None, None].exp()
+ + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
+ )
+
+ if not output_final_state:
+ last_recurrent_state = None
+ core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
+ core_attn_out = core_attn_out[:, :, :sequence_length]
+ core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+ return core_attn_out, last_recurrent_state
+
+
+def torch_recurrent_gated_delta_rule(
+ query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
+):
+ initial_dtype = query.dtype
+ if use_qk_l2norm_in_kernel:
+ query = l2norm(query, dim=-1, eps=1e-6)
+ key = l2norm(key, dim=-1, eps=1e-6)
+ query, key, value, beta, g = [
+ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
+ ]
+
+ batch_size, num_heads, sequence_length, k_head_dim = key.shape
+ v_head_dim = value.shape[-1]
+ scale = 1 / (query.shape[-1] ** 0.5)
+ query = query * scale
+
+ core_attn_out = torch.zeros(
+ batch_size, num_heads, sequence_length, v_head_dim, dtype=value.dtype, device=value.device
+ )
+ last_recurrent_state = (
+ torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim, dtype=value.dtype, device=value.device)
+ if initial_state is None
+ else initial_state.to(value)
+ )
+
+ for i in range(sequence_length):
+ q_t = query[:, :, i]
+ k_t = key[:, :, i]
+ v_t = value[:, :, i]
+ g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
+ beta_t = beta[:, :, i].unsqueeze(-1)
+
+ last_recurrent_state = last_recurrent_state * g_t
+ kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)
+ delta = (v_t - kv_mem) * beta_t
+ last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
+ core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)
+
+ if not output_final_state:
+ last_recurrent_state = None
+ core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+ return core_attn_out, last_recurrent_state
+
+
+class Qwen3_5MoeGatedDeltaNet(nn.Module):
+ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.num_v_heads = config.linear_num_value_heads
+ self.num_k_heads = config.linear_num_key_heads
+ self.head_k_dim = config.linear_key_head_dim
+ self.head_v_dim = config.linear_value_head_dim
+ self.key_dim = self.head_k_dim * self.num_k_heads
+ self.value_dim = self.head_v_dim * self.num_v_heads
+
+ self.conv_kernel_size = config.linear_conv_kernel_dim
+ self.layer_idx = layer_idx
+ self.activation = config.hidden_act
+ self.act = ACT2FN[config.hidden_act]
+ self.layer_norm_epsilon = config.rms_norm_eps
+
+ # QKV
+ self.conv_dim = self.key_dim * 2 + self.value_dim
+ self.conv1d = nn.Conv1d(
+ in_channels=self.conv_dim,
+ out_channels=self.conv_dim,
+ bias=False,
+ kernel_size=self.conv_kernel_size,
+ groups=self.conv_dim,
+ padding=self.conv_kernel_size - 1,
+ )
+
+ # time step projection (discretization)
+ # instantiate once and copy inv_dt in init_weights of PretrainedModel
+ self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
+
+ A = torch.empty(self.num_v_heads).uniform_(0, 16)
+ self.A_log = nn.Parameter(torch.log(A))
+
+ self.norm = (
+ Qwen3_5MoeRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)
+ if FusedRMSNormGated is None
+ else FusedRMSNormGated(
+ self.head_v_dim,
+ eps=self.layer_norm_epsilon,
+ activation=self.activation,
+ device=torch.cuda.current_device(),
+ dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
+ )
+ )
+
+ self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+ self.causal_conv1d_fn = causal_conv1d_fn
+ self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update
+ self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule
+ self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule
+
+ if not is_fast_path_available:
+ logger.warning_once(
+ "The fast path is not available because one of the required library is not installed. Falling back to "
+ "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"
+ " https://github.com/Dao-AILab/causal-conv1d"
+ )
+
+ self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)
+ self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)
+ self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+ self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+
+ # Set up dimensions for reshapes later
+ batch_size, seq_len, _ = hidden_states.shape
+
+ # We have cached `conv_state` / `recurrent_state` to continue from. The two cached modes
+ # (single-token decode and chunk-tokens continuation) share the state read here; they only
+ # diverge in how the conv input is assembled and which kernel consumes the states below,
+ # which we gate locally on `seq_len`.
+ use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx)
+
+ # getting projected states from cache if it exists
+ if use_precomputed_states:
+ conv_state = cache_params.layers[self.layer_idx].conv_states
+ recurrent_state = cache_params.layers[self.layer_idx].recurrent_states
+
+ mixed_qkv = self.in_proj_qkv(hidden_states)
+ mixed_qkv = mixed_qkv.transpose(1, 2)
+
+ z = self.in_proj_z(hidden_states)
+ z = z.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+ b = self.in_proj_b(hidden_states)
+ a = self.in_proj_a(hidden_states)
+
+ if use_precomputed_states and seq_len == 1:
+ # Single-token cached decode: the fused per-step kernel updates the conv state in-place.
+ mixed_qkv = self.causal_conv1d_update(
+ mixed_qkv,
+ conv_state,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.activation,
+ )
+ else:
+ # Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state).
+ if use_precomputed_states:
+ # Cached chunked-tokens decode: prepend the cached conv context so the causal conv
+ # sees the correct left-context rather than zero-padding. Dropped from the output
+ # at the end of this branch.
+ mixed_qkv = torch.cat([conv_state, mixed_qkv], dim=-1)
+ if cache_params is not None:
+ new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
+ cache_params.update_conv_state(new_conv_state, self.layer_idx)
+ if self.causal_conv1d_fn is not None:
+ mixed_qkv = self.causal_conv1d_fn(
+ x=mixed_qkv,
+ weight=self.conv1d.weight.squeeze(1),
+ bias=self.conv1d.bias,
+ activation=self.activation,
+ seq_idx=kwargs.get("seq_idx"),
+ )
+ else:
+ mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
+ if use_precomputed_states:
+ mixed_qkv = mixed_qkv[:, :, -seq_len:]
+
+ mixed_qkv = mixed_qkv.transpose(1, 2)
+ query, key, value = torch.split(
+ mixed_qkv,
+ [
+ self.key_dim,
+ self.key_dim,
+ self.value_dim,
+ ],
+ dim=-1,
+ )
+
+ query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
+ key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
+ value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)
+
+ beta = b.sigmoid()
+ # If the model is loaded in fp16, without the .float() here, A might be -inf
+ g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+ if self.num_v_heads // self.num_k_heads > 1:
+ query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+ key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+
+ if use_precomputed_states and seq_len == 1:
+ core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+ query,
+ key,
+ value,
+ g=g,
+ beta=beta,
+ initial_state=recurrent_state,
+ output_final_state=cache_params is not None,
+ use_qk_l2norm_in_kernel=True,
+ )
+ else:
+ core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+ query,
+ key,
+ value,
+ g=g,
+ beta=beta,
+ initial_state=recurrent_state if use_precomputed_states else None,
+ output_final_state=cache_params is not None,
+ use_qk_l2norm_in_kernel=True,
+ # The chunked FLA kernel takes a single `cu_seqlens` arg; for packed self-attention this matches q-side lengths.
+ cu_seqlens=kwargs.get("cu_seq_lens_q"),
+ )
+
+ # Update cache
+ if cache_params is not None:
+ cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx)
+
+ # reshape input data into 2D tensor
+ core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
+ z = z.reshape(-1, self.head_v_dim)
+ core_attn_out = self.norm(core_attn_out, z)
+ core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)
+
+ output = self.out_proj(core_attn_out)
+ return output
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Removes the interleaving of cos and sin from GLM
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Qwen3_5MoeAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim * 2, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.q_norm = Qwen3_5MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
+ self.k_norm = Qwen3_5MoeRMSNorm(
+ self.head_dim, eps=config.rms_norm_eps
+ ) # thus post q_norm does not need reshape
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states, gate = torch.chunk(
+ self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
+ )
+ gate = gate.reshape(*input_shape, -1)
+
+ query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = attn_output * torch.sigmoid(gate)
+
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class Qwen3_5MoeMLP(nn.Module):
+ def __init__(self, config: Qwen3_5MoeConfig, intermediate_size: int):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+@use_experts_implementation
+class Qwen3_5MoeExperts(nn.Module):
+ """Collection of expert weights stored as 3D tensors."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.num_experts = config.num_experts
+ self.hidden_dim = config.hidden_size
+ self.intermediate_dim = config.moe_intermediate_size
+ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+ self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ top_k_index: torch.Tensor,
+ top_k_weights: torch.Tensor,
+ ) -> torch.Tensor:
+ final_hidden_states = torch.zeros_like(hidden_states)
+ with torch.no_grad():
+ expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+ expert_mask = expert_mask.permute(2, 1, 0)
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+ for expert_idx in expert_hit:
+ expert_idx = expert_idx[0]
+ if expert_idx == self.num_experts:
+ continue
+ top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+ current_state = hidden_states[token_idx]
+ gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+ current_hidden_states = self.act_fn(gate) * up
+ current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+ current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+ final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+ return final_hidden_states
+
+
+class Qwen3_5MoeTopKRouter(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.top_k = config.num_experts_per_tok
+ self.num_experts = config.num_experts
+ self.hidden_dim = config.hidden_size
+ self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+
+ def forward(self, hidden_states):
+ hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+ router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts)
+ router_probs = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)
+ router_top_value, router_indices = torch.topk(router_probs, self.top_k, dim=-1) # (seq_len, top_k)
+ router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
+ router_top_value = router_top_value.to(router_logits.dtype)
+ router_scores = router_top_value
+ return router_logits, router_scores, router_indices
+
+
+class Qwen3_5MoeSparseMoeBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.gate = Qwen3_5MoeTopKRouter(config)
+ self.experts = Qwen3_5MoeExperts(config)
+ self.shared_expert = Qwen3_5MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size)
+ self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
+ hidden_states_reshaped = hidden_states.view(-1, hidden_dim)
+ shared_expert_output = self.shared_expert(hidden_states_reshaped)
+ _, routing_weights, selected_experts = self.gate(hidden_states_reshaped)
+ expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights)
+
+ shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output
+
+ expert_output = expert_output + shared_expert_output
+ expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim)
+ return expert_output
+
+
+class Qwen3_5MoeRMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.zeros(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float())
+ # Llama does x.to(float16) * w whilst Qwen3_5Moe is (x * w).to(float16)
+ # See https://github.com/huggingface/transformers/pull/29402
+ output = output * (1.0 + self.weight.float())
+ return output.type_as(x)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+class Qwen3_5MoeDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen3_5MoeTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.layer_type = config.layer_types[layer_idx]
+ if self.layer_type == "linear_attention":
+ self.linear_attn = Qwen3_5MoeGatedDeltaNet(config, layer_idx)
+ elif self.layer_type == "full_attention":
+ self.self_attn = Qwen3_5MoeAttention(config, layer_idx)
+ self.mlp = Qwen3_5MoeSparseMoeBlock(config)
+ self.input_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Token Mixer
+ if self.layer_type == "linear_attention":
+ hidden_states = self.linear_attn(
+ hidden_states=hidden_states,
+ cache_params=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ elif self.layer_type == "full_attention":
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ # For the MoE layers, we need to unpack
+ if isinstance(hidden_states, tuple):
+ hidden_states, _ = hidden_states
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class Qwen3_5MoePreTrainedModel(PreTrainedModel):
+ config: Qwen3_5MoeConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen3_5MoeDecoderLayer", "Qwen3_5MoeVisionBlock"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _keys_to_ignore_on_load_unexpected = [r"^mtp.*"]
+ _can_record_outputs = {
+ "router_logits": OutputRecorder(Qwen3_5MoeTopKRouter, index=0),
+ "hidden_states": Qwen3_5MoeDecoderLayer,
+ "attentions": Qwen3_5MoeAttention,
+ }
+ _is_stateful = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, Qwen3_5MoeGatedDeltaNet):
+ init.ones_(module.dt_bias)
+ init.copy_(module.A_log, torch.empty_like(module.A_log).uniform_(0, 16).log_())
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
+ elif isinstance(module, Qwen3_5MoeRMSNorm):
+ init.zeros_(module.weight)
+ elif isinstance(module, Qwen3_5MoeExperts):
+ init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
+ init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, Qwen3_5MoeSparseMoeBlock):
+ init.normal_(module.gate.weight, mean=0.0, std=self.config.initializer_range)
+ elif isinstance(module, Qwen3_5MoeVisionRotaryEmbedding):
+ inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+
+
+class Qwen3_5MoeVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3_5MoeVisionPatchEmbed(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+
+ kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+ self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen3_5MoeVisionPatchMerger(nn.Module):
+ def __init__(self, config: Qwen3_5MoeVisionConfig, use_postshuffle_norm=False) -> None:
+ super().__init__()
+ self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)
+ self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
+ self.act_fn = nn.GELU()
+ self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)
+ x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+ return x
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+class Qwen3_5MoeVisionAttention(nn.Module):
+ def __init__(self, config: Qwen3_5MoeVisionConfig) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen3_5MoeVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config, attn_implementation: str = "sdpa") -> None:
+ super().__init__()
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+ self.attn = Qwen3_5MoeVisionAttention(config=config)
+ self.mlp = Qwen3_5MoeVisionMLP(config=config)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor`):
+ Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels.
+ rotary_pos_emb (`torch.Tensor`, *optional*):
+ Precomputed rotary positional embeddings applied to the vision attention query/key states.
+ """
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen3_5MoeVisionModel(Qwen3_5MoePreTrainedModel):
+ config: Qwen3_5MoeVisionConfig
+ input_modalities = ("image", "video")
+ _no_split_modules = ["Qwen3_5MoeVisionBlock"]
+ _can_record_outputs = {
+ "hidden_states": Qwen3_5MoeVisionBlock,
+ "attentions": Qwen3_5MoeVisionAttention,
+ }
+
+ def __init__(self, config, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen3_5MoeVisionPatchEmbed(
+ config=config,
+ )
+
+ self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
+ self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen3_5MoeVisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList([Qwen3_5MoeVisionBlock(config) for _ in range(config.depth)])
+ self.merger = Qwen3_5MoeVisionPatchMerger(
+ config=config,
+ use_postshuffle_norm=False,
+ )
+
+ self.gradient_checkpointing = False
+
+ self.post_init()
+
+ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+ warnings.warn(
+ f"`{self.__class__.__name__}.rot_pos_emb` is deprecated and will be removed in v5.11. Use `get_vision_position_ids` from `transformers.vision_utils` and apply the rotary embedding module.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size)
+ rotary_pos_emb = self.rotary_pos_emb(position_ids)
+ return rotary_pos_emb
+
+ def fast_pos_embed_interpolate(self, grid_thw):
+ warnings.warn(
+ f"`{self.__class__.__name__}.fast_pos_embed_interpolate` is deprecated and will be removed in v5.11. Use `get_vision_bilinear_indices_and_weights` from `transformers.vision_utils` and apply `self.pos_embed`.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ bilinear_indices, bilinear_weights = get_vision_bilinear_indices_and_weights(
+ grid_thw,
+ num_grid_per_side=self.num_grid_per_side,
+ spatial_merge_size=self.config.spatial_merge_size,
+ )
+ return (self.pos_embed(bilinear_indices) * bilinear_weights[:, :, None]).sum(0)
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ bilinear_indices, bilinear_weights = get_vision_bilinear_indices_and_weights(
+ grid_thw,
+ num_grid_per_side=self.num_grid_per_side,
+ spatial_merge_size=self.config.spatial_merge_size,
+ kwargs=kwargs,
+ )
+ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+ cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs)
+
+ hidden_states = self.patch_embed(hidden_states)
+ pos_embeds = (self.pos_embed(bilinear_indices) * bilinear_weights[:, :, None]).sum(0)
+ hidden_states = hidden_states + pos_embeds.to(hidden_states.dtype)
+ rotary_pos_emb = self.rotary_pos_emb(position_ids)
+
+ seq_len, _ = hidden_states.size()
+ hidden_states = hidden_states.reshape(seq_len, -1)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ for blk in self.blocks:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ merged_hidden_states = self.merger(hidden_states)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=hidden_states,
+ pooler_output=merged_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Llava outputs, with hidden states and attentions.
+ """
+)
+@dataclass
+class Qwen3_5MoeModelOutputWithPast(ModelOutput):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+ router_logits: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen3_5Moe causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class Qwen3_5MoeCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+ router_logits: tuple[torch.FloatTensor] | None = None
+ aux_loss: torch.FloatTensor | None = None
+
+
+class Qwen3_5MoeTextModel(Qwen3_5MoePreTrainedModel):
+ config: Qwen3_5MoeTextConfig
+
+ def __init__(self, config: Qwen3_5MoeTextConfig):
+ super().__init__(config)
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+ self.layers = nn.ModuleList(
+ [Qwen3_5MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Qwen3_5MoeTextRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ # the hard coded `4` is for text, temporal, height and width.
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(4, position_ids.shape[0], -1)
+
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ text_position_ids = None
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=text_position_ids,
+ )
+ linear_attn_mask = self._update_linear_attn_mask(attention_mask, past_key_values)
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+ layer_mask = linear_attn_mask if self.config.layer_types[i] == "linear_attention" else causal_mask
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=layer_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return Qwen3_5MoeModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def _update_linear_attn_mask(self, attention_mask, past_key_values):
+ """
+ NOTE: Left-padding is used for linear attention mask.
+ No need for zeroing states when
+ 1. Cached forward
+ 2. Attending to all inputs
+ """
+ linear_attn_mask = attention_mask
+ if (past_key_values is not None and past_key_values.has_previous_state()) or (
+ attention_mask is not None and torch.all(attention_mask == 1)
+ ):
+ linear_attn_mask = None
+ return linear_attn_mask
+
+
+@auto_docstring
+class Qwen3_5MoeModel(Qwen3_5MoePreTrainedModel):
+ base_model_prefix = "model"
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ config: Qwen3_5MoeConfig
+ _no_split_modules = ["Qwen3_5MoeDecoderLayer", "Qwen3_5MoeVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = Qwen3_5MoeVisionModel._from_config(config.vision_config)
+ self.language_model = Qwen3_5MoeTextModel._from_config(config.text_config)
+ self.rope_deltas = None # cache rope_deltas here
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_vision_position_ids(
+ self,
+ start_position: int,
+ grid_thw: list[int, int, int] | torch.Tensor,
+ temp_merge_size: int = 1,
+ spatial_merge_size: int = 1,
+ time_interval: int = 1,
+ device: str | torch.device | None = None,
+ ):
+ """
+ Compute 3D positional indices for vision tokens derived from a single image or video input.
+
+ The positions are generated from the input grid defined by temporal (T), height (H), and
+ width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the
+ merge sizes used in the vision backbone. The resulting positions are offset by `start_position`.
+
+ Args:
+ start_position (`int`):
+ Offset added to all computed positional indices.
+ grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`):
+ The (T, H, W) grid representing the feature layout of the current image or video after patch embedding.
+ temp_merge_size (`int`, *optional*):
+ Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided
+ by this value. Defaults to 1.
+ spatial_merge_size (`int`, *optional*):
+ Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided
+ by this value. Defaults to 1.
+ time_interval (`int`, *optional*):
+ Spacing factor applied between consecutive temporal position indices.Defaults to 1.
+ device (`str` or `torch.device`, *optional*):
+ Device on which the resulting tensor is allocated. If `None`, uses the current default device.
+
+ Returns:
+ torch.LongTensor of shape (3, sequence_length):
+ Positional indices for temporal, height, and width dimensions,
+ flattened into sequence form and offset by `start_position`.
+ """
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ grid_thw[0].item() // temp_merge_size,
+ grid_thw[1].item() // spatial_merge_size,
+ grid_thw[2].item() // spatial_merge_size,
+ )
+
+ # Add `start_position` after arange for compile
+ position_temporal = torch.arange(llm_grid_t, device=device) * time_interval
+ position_width = torch.arange(llm_grid_w, device=device) + start_position
+ position_height = torch.arange(llm_grid_h, device=device) + start_position
+
+ # Repeat the positions per each grid and per video frame. Repeat patterns are important
+ # do not modify without checking values!
+ position_width = position_width.repeat(llm_grid_h * llm_grid_t)
+ position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t)
+ # Important: add `start_positions` after applying `time_interval`, order matters
+ position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position
+ vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
+
+ return vision_position_ids
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor,
+ mm_token_type_ids: torch.IntTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Difference from Qwen2VL/Qwen2.5VL's get_rope_index:
+ - Since Qwen3.5 use timestamps to separate videos, like , the video_grid_thw should also be split too.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
+ Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2).
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+
+ # Separate video grid thw into multiple grids because timestamps are used to separate videos.
+ if video_grid_thw is not None:
+ video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+ video_grid_thw[:, 0] = 1
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+
+ mrope_position_deltas = []
+ position_ids = torch.zeros(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ grid_iters = {
+ 1: iter(image_grid_thw) if image_grid_thw is not None else None,
+ 2: iter(video_grid_thw) if video_grid_thw is not None else None,
+ }
+
+ for batch_idx, current_input_ids in enumerate(input_ids):
+ input_token_type = mm_token_type_ids[batch_idx]
+ if attention_mask is not None:
+ current_input_ids = current_input_ids[attention_mask[batch_idx].bool()]
+ input_token_type = input_token_type[attention_mask[batch_idx].bool()]
+
+ input_type_group = []
+ for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]):
+ group = list(group)
+ start_index = group[0][0]
+ end_index = group[-1][0] + 1
+ input_type_group.append((key, start_index, end_index))
+
+ current_pos = 0
+ llm_pos_ids_list = []
+ for modality_type, start_idx, end_idx in input_type_group:
+ # text == 0
+ if modality_type == 0:
+ text_len = end_idx - start_idx
+ llm_pos_ids_list.append(
+ torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos
+ )
+ current_pos += text_len
+ # image == 1, video == 2
+ else:
+ grid_thw = next(grid_iters[modality_type])
+ vision_position_ids = self.get_vision_position_ids(
+ current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device
+ )
+ llm_pos_ids_list.append(vision_position_ids)
+ current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ if attention_mask is not None:
+ position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device)
+ else:
+ position_ids[:, batch_idx] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+
+ @accepts_precomputed_kwargs(modality="video")
+ @can_return_tuple
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ # Same implementation as for images
+ return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+ @accepts_precomputed_kwargs(modality="image")
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ vision_output: BaseModelOutputWithPooling = self.visual(
+ pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs
+ )
+ image_embeds = vision_output.pooler_output
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(image_embeds, split_sizes)
+ vision_output.pooler_output = image_embeds
+
+ return vision_output
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: torch.FloatTensor | None = None,
+ video_features: torch.FloatTensor | None = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_video_mask].numel() == video_features.numel(),
+ f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+ )
+ return special_image_mask, special_video_mask
+
+ def compute_3d_position_ids(
+ self,
+ input_ids: torch.Tensor | None,
+ inputs_embeds: torch.Tensor | None,
+ image_grid_thw: torch.Tensor | None = None,
+ video_grid_thw: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: torch.Tensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ ) -> torch.Tensor | None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ has_multimodal = image_grid_thw is not None or video_grid_thw is not None
+ if has_multimodal and mm_token_type_ids is None and input_ids is not None:
+ raise ValueError(
+ "Multimodal data was passed (via `image_grid_thw` or `video_grid_thw`) but `mm_token_type_ids` is "
+ "missing. Please pass `mm_token_type_ids` to the model so that multimodal RoPE (M-RoPE) can be "
+ "computed correctly. `mm_token_type_ids` is returned by the processor alongside `input_ids`."
+ )
+ can_compute_mrope = input_ids is not None and mm_token_type_ids is not None and has_multimodal
+
+ if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ attention_mask=attention_mask,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+ self.rope_deltas = rope_deltas
+ # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+ # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+ # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+ # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+ elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+ position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device)
+ else:
+ position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length)
+ position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device)
+ delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0)
+ position_ids = position_ids + delta.to(device=inputs_embeds.device)
+ else:
+ # Can't build correct 3D positions. Let the model infer it
+ position_ids = None
+ return position_ids
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Qwen3_5MoeModelOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_outputs: BaseModelOutputWithPooling = self.get_image_features(
+ pixel_values, image_grid_thw, return_dict=True, **kwargs
+ )
+ image_embeds = image_outputs.pooler_output
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_outputs: BaseModelOutputWithPooling = self.get_video_features(
+ pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+ )
+ video_embeds = video_outputs.pooler_output
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if position_ids is None:
+ position_ids = self.compute_3d_position_ids(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return Qwen3_5MoeModelOutputWithPast(
+ **outputs,
+ rope_deltas=self.rope_deltas,
+ )
+
+
+def load_balancing_loss_func(
+ gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+ num_experts: int | None = None,
+ top_k=2,
+ attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+ r"""
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+ See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+ experts is too unbalanced.
+
+ Args:
+ gate_logits:
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+ shape [batch_size X sequence_length, num_experts].
+ num_experts:
+ Number of experts
+ top_k:
+ The number of experts to route per-token, can be also interpreted as the `top-k` routing
+ parameter.
+ attention_mask (`torch.Tensor`, *optional*):
+ The attention_mask used in forward function
+ shape [batch_size X sequence_length] if not None.
+
+ Returns:
+ The auxiliary loss.
+ """
+ if gate_logits is None or not isinstance(gate_logits, tuple):
+ return 0
+
+ if isinstance(gate_logits, tuple):
+ compute_device = gate_logits[0].device
+ concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+ if attention_mask is None:
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
+ else:
+ batch_size, sequence_length = attention_mask.shape
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+ expert_attention_mask = (
+ attention_mask[None, :, :, None, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+ .reshape(-1, top_k, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+ expert_attention_mask, dim=0
+ )
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+ router_per_expert_attention_mask = (
+ attention_mask[None, :, :, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+ .reshape(-1, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+ router_per_expert_attention_mask, dim=0
+ )
+
+ overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+ return overall_loss * num_experts
+
+
+@auto_docstring
+class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+ config: Qwen3_5MoeTextConfig
+ _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Qwen3_5MoeTextModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+ self.router_aux_loss_coef = config.router_aux_loss_coef
+ self.num_experts = config.num_experts
+ self.num_experts_per_tok = config.num_experts_per_tok
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_router_logits: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> MoeCausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, Qwen3_5MoeForCausalLM
+
+ >>> model = Qwen3_5MoeForCausalLM.from_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct")
+ >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+
+ output_router_logits = (
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs: MoeModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_router_logits=output_router_logits,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+ aux_loss = None
+ if output_router_logits:
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.num_experts,
+ self.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
+
+ return MoeCausalLMOutputWithPast(
+ loss=loss,
+ aux_loss=aux_loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ router_logits=outputs.router_logits,
+ )
+
+
+class Qwen3_5MoeForConditionalGeneration(Qwen3_5MoePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ config: Qwen3_5MoeConfig
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Qwen3_5MoeModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Qwen3_5MoeCausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+
+ Example:
+ ```python
+ >>> from transformers import AutoProcessor, Qwen3_5MoeForConditionalGeneration
+
+ >>> model = Qwen3_5MoeForConditionalGeneration.from_pretrained("Qwen/Qwen3.5-35B-A3B-Instruct", dtype="auto", device_map="auto")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3.5-35B-A3B-Instruct")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "image",
+ "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+ },
+ {"type": "text", "text": "Describe this image in short."},
+ ],
+ }
+ ]
+
+ >>> # Preparation for inference
+ >>> inputs = processor.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ return_dict=True,
+ return_tensors="pt"
+ )
+ >>> inputs = inputs.to(model.device)
+
+ >>> # Generate
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=128)
+ >>> generated_ids_trimmed = [
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+ ]
+ >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background."
+ ```"""
+
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ mm_token_type_ids=mm_token_type_ids,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ aux_loss = None
+ if kwargs.get("output_router_logits", False):
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.config.text_config.num_experts,
+ self.config.text_config.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.config.text_config.router_aux_loss_coef * aux_loss.to(
+ loss.device
+ ) # make sure to reside in the same device
+
+ return Qwen3_5MoeCausalLMOutputWithPast(
+ loss=loss,
+ aux_loss=aux_loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=outputs.rope_deltas,
+ router_logits=outputs.router_logits,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs):
+ # Overwritten -- requires 3D position ids
+
+ text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)
+
+ # Early exit in case we are continuing generation from past kv
+ past_length = 0
+ if (cache := model_kwargs.get("past_key_values")) is not None:
+ past_length = cache.get_seq_length()
+ if past_length != 0 and self.model.rope_deltas is not None:
+ position_ids = text_positions[None, ...] + self.model.rope_deltas
+ return position_ids
+
+ # Otherwise compute 3d position ids for vision tokens and concat with text position ids
+ if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+ inputs_tensor = model_kwargs["input_ids"]
+
+ is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+ if (
+ is_input_ids
+ and model_kwargs.get("mm_token_type_ids") is not None
+ and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None)
+ ):
+ model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"}
+ vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs)
+ self.model.rope_deltas = rope_deltas
+ else:
+ vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1)
+ self.model.rope_deltas = torch.zeros(
+ inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device
+ )
+
+ # Concatenate "text + vision" positions into [4, bs, seq-len]
+ text_positions = text_positions[None, ...]
+ position_ids = torch.cat([text_positions, vision_positions], dim=0)
+
+ return position_ids
+
+ def _get_image_nums_and_video_nums(
+ self,
+ input_ids: torch.LongTensor | None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+ These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Returns:
+ image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+ video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+ """
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+
+ if inputs_embeds is not None:
+ vision_start_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(vision_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ image_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ video_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ else:
+ vision_start_mask = input_ids == vision_start_token_id
+ image_mask = input_ids == image_token_id
+ video_mask = input_ids == video_token_id
+
+ vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
+ image_nums = torch.sum(vision_first_mask & image_mask, dim=1)
+ video_nums = torch.sum(vision_first_mask & video_mask, dim=1)
+
+ return image_nums, video_nums
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: torch.LongTensor | None = None,
+ **model_kwargs,
+ ) -> tuple[torch.LongTensor, dict[str, Any]]:
+ # Overwritten -- Qwen3_5Moe use timestamps and remove second_per_grid_ts
+ # Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ video_grid_thw = model_kwargs.get("video_grid_thw", None)
+ image_nums, video_nums = self._get_image_nums_and_video_nums(
+ input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
+ )
+
+ # video_nums: (batch_size,)
+ # since video_nums is the number of videos in the input dependent on the input_ids(vision_start),
+ # but Qwen3_5Moe append vision_start to each frame of each video, so we need to recover the real video_nums according to video_grid_thw
+ if video_grid_thw is not None:
+ cumulative_frame_counts = torch.cumsum(video_grid_thw[:, 0], dim=0)
+ cumulative_token_video_counts = torch.cumsum(video_nums, dim=0)
+ # Find video boundaries in cumulative_frame_counts
+ video_boundary_indices = torch.searchsorted(cumulative_frame_counts, cumulative_token_video_counts)
+ # example: video_boundary_indices = [3, 5] means video_nums = [4, 2]
+ video_nums = torch.diff(torch.cat([-video_boundary_indices.new_ones(1), video_boundary_indices]))
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # split images into samples
+ samples = torch.split(image_grid_thw, list(image_nums))
+ # compute the sequence length of images for each sample
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # get the num of images for each sample
+ lengths = list(image_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "pixel_values_videos":
+ samples = torch.split(video_grid_thw, list(video_nums))
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "video_grid_thw":
+ lengths = list(video_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if key == "position_ids" and dict_to_expand[key].ndim == 3:
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=1)
+ elif (
+ dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
+
+
+__all__ = [
+ "Qwen3_5MoeVisionModel",
+ "Qwen3_5MoeTextModel",
+ "Qwen3_5MoeModel",
+ "Qwen3_5MoeForCausalLM",
+ "Qwen3_5MoeForConditionalGeneration",
+ "Qwen3_5MoePreTrainedModel",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/t5gemma2/configuration_t5gemma2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/t5gemma2/configuration_t5gemma2.py
new file mode 100644
index 0000000000000000000000000000000000000000..d9a9a3f5769f8114ca7211089bad7315cea00e3b
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/t5gemma2/configuration_t5gemma2.py
@@ -0,0 +1,403 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/t5gemma2/modular_t5gemma2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_t5gemma2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Google Inc. 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 typing import Any
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+from ..siglip import SiglipVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="google/t5gemma-2-270m-270m")
+@strict
+class T5Gemma2TextConfig(PreTrainedConfig):
+ r"""
+ query_pre_attn_scalar (`float`, *optional*, defaults to 256):
+ Scaling factor used on the attention scores
+ final_logit_softcapping (`float`, *optional*):
+ Scaling factor when applying tanh softcapping on the logits.
+ attn_logit_softcapping (`float`, *optional*):
+ Scaling factor when applying tanh softcapping on the attention scores.
+ """
+
+ model_type = "t5gemma2_text"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+ "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ vocab_size: int = 262_208
+ hidden_size: int = 2304
+ intermediate_size: int = 9216
+ num_hidden_layers: int = 26
+ num_attention_heads: int = 8
+ num_key_value_heads: int = 4
+ head_dim: int = 256
+ hidden_activation: str = "gelu_pytorch_tanh"
+ max_position_embeddings: int = 131_072
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ pad_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ bos_token_id: int | None = 2
+ tie_word_embeddings: bool = True
+ rope_parameters: dict | None = None
+ attention_bias: bool = False
+ attention_dropout: int | float | None = 0.0
+ query_pre_attn_scalar: int = 256
+ sliding_window: int | None = 4096
+ layer_types: list[str] | None = None
+ final_logit_softcapping: float | None = None
+ attn_logit_softcapping: float | None = None
+ default_theta = {"global": 1_000_000.0, "local": 10_000.0}
+
+ def __post_init__(self, **kwargs):
+ # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub
+ _sliding_window_pattern = kwargs.pop("sliding_window_pattern", 6)
+ if self.layer_types is None:
+ self.layer_types = [
+ "sliding_attention" if bool((i + 1) % _sliding_window_pattern) else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.hidden_size % self.num_attention_heads != 0:
+ raise ValueError(
+ f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({self.num_attention_heads})."
+ )
+
+ def convert_rope_params_to_dict(self, **kwargs):
+ rope_scaling = kwargs.pop("rope_scaling", None)
+
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`. If we find `rope_parameters`
+ # as arg in the inputs, we can safely assume that it is in the new format. New naming used -> new format
+ default_rope_params = {
+ "sliding_attention": {"rope_type": "default"},
+ "full_attention": {"rope_type": "default"},
+ }
+ self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else default_rope_params
+ if rope_scaling is not None:
+ self.rope_parameters["full_attention"].update(rope_scaling)
+
+ # Set default values if not present
+ if self.rope_parameters.get("full_attention") is None:
+ self.rope_parameters["full_attention"] = {"rope_type": "default"}
+ self.rope_parameters["full_attention"].setdefault(
+ "rope_theta", kwargs.pop("rope_theta", self.default_theta["global"])
+ )
+ if self.rope_parameters.get("sliding_attention") is None:
+ self.rope_parameters["sliding_attention"] = {"rope_type": "default"}
+ self.rope_parameters["sliding_attention"].setdefault(
+ "rope_theta", kwargs.pop("rope_local_base_freq", self.default_theta["local"])
+ )
+
+ # Standardize and validate the correctness of rotary position embeddings parameters
+ self.standardize_rope_params()
+ return kwargs
+
+
+@auto_docstring(checkpoint="google/t5gemma-2-270m-270m")
+@strict
+class T5Gemma2EncoderConfig(PreTrainedConfig):
+ r"""
+ mm_tokens_per_image (`int`, *optional*, defaults to 256):
+ The number of tokens per image embedding.
+ boi_token_index (`int`, *optional*, defaults to 255999):
+ The begin-of-image token index to wrap the image prompt.
+ eoi_token_index (`int`, *optional*, defaults to 256000):
+ The end-of-image token index to wrap the image prompt.
+
+ Example:
+
+ ```python
+ >>> from transformers import T5Gemma2EncoderForConditionalGeneration, T5Gemma2EncoderConfig, SiglipVisionConfig, T5Gemma2EncoderTextConfig
+
+ >>> # Initializing a Siglip-like vision config
+ >>> vision_config = SiglipVisionConfig()
+
+ >>> # Initializing a T5Gemma2Encoder Text config
+ >>> text_config = T5Gemma2EncoderTextConfig()
+
+ >>> # Initializing a T5Gemma2Encoder gemma-3-4b style configuration
+ >>> configuration = T5Gemma2EncoderConfig(vision_config, text_config)
+
+ >>> # Initializing a model from the gemma-3-4b style configuration
+ >>> model = T5Gemma2EncoderTextConfig(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "t5gemma2_encoder"
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ "boi_token_id": "boi_token_index",
+ "eoi_token_id": "eoi_token_index",
+ }
+
+ sub_configs = {
+ "text_config": T5Gemma2TextConfig,
+ "vision_config": SiglipVisionConfig,
+ }
+
+ text_config: T5Gemma2TextConfig | dict[str, Any] | None = None
+ vision_config: SiglipVisionConfig | dict[str, Any] | None = None
+ mm_tokens_per_image: int | None = 256
+ boi_token_index: int | None = 255_999
+ eoi_token_index: int | None = 256_000
+ image_token_index: int | None = 262_144
+ initializer_range: float | None = 0.02
+ tie_word_embeddings: bool | None = True
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ self.text_config = T5Gemma2TextConfig()
+ logger.info("text_config is None, using default T5Gemma2EncoderTextConfig text config.")
+ elif isinstance(self.text_config, dict):
+ self.text_config = T5Gemma2TextConfig(**self.text_config)
+
+ if isinstance(self.vision_config, dict):
+ self.vision_config = SiglipVisionConfig(**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = SiglipVisionConfig()
+ logger.info("vision_config is None, using default SiglipVisionConfig vision config.")
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="google/t5gemma-2-270m-270m")
+@strict
+class T5Gemma2DecoderConfig(PreTrainedConfig):
+ r"""
+ query_pre_attn_scalar (`float`, *optional*, defaults to 256):
+ Scaling factor used on the attention scores
+ final_logit_softcapping (`float`, *optional*):
+ Scaling factor when applying tanh softcapping on the logits.
+ attn_logit_softcapping (`float`, *optional*):
+ Scaling factor when applying tanh softcapping on the attention scores.
+ """
+
+ model_type = "t5gemma2_decoder"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+ "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ vocab_size: int = 262_208
+ hidden_size: int = 2304
+ intermediate_size: int = 9216
+ num_hidden_layers: int = 26
+ num_attention_heads: int = 8
+ num_key_value_heads: int = 4
+ head_dim: int = 256
+ hidden_activation: str = "gelu_pytorch_tanh"
+ max_position_embeddings: int = 131_072
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ pad_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ bos_token_id: int | None = 2
+ tie_word_embeddings: bool = True
+ rope_parameters: dict | None = None
+ attention_bias: bool = False
+ attention_dropout: int | float | None = 0.0
+ query_pre_attn_scalar: int = 256
+ sliding_window: int | None = 4096
+ layer_types: list[str] | None = None
+ final_logit_softcapping: float | None = None
+ attn_logit_softcapping: float | None = None
+ default_theta = {"global": 1_000_000.0, "local": 10_000.0}
+
+ def __post_init__(self, **kwargs):
+ # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub
+ _sliding_window_pattern = kwargs.pop("sliding_window_pattern", 6)
+ if self.layer_types is None:
+ self.layer_types = [
+ "sliding_attention" if bool((i + 1) % _sliding_window_pattern) else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.hidden_size % self.num_attention_heads != 0:
+ raise ValueError(
+ f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({self.num_attention_heads})."
+ )
+
+ def convert_rope_params_to_dict(self, **kwargs):
+ rope_scaling = kwargs.pop("rope_scaling", None)
+
+ # Try to set `rope_scaling` if available, otherwise use `rope_parameters`. If we find `rope_parameters`
+ # as arg in the inputs, we can safely assume that it is in the new format. New naming used -> new format
+ default_rope_params = {
+ "sliding_attention": {"rope_type": "default"},
+ "full_attention": {"rope_type": "default"},
+ }
+ self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else default_rope_params
+ if rope_scaling is not None:
+ self.rope_parameters["full_attention"].update(rope_scaling)
+
+ # Set default values if not present
+ if self.rope_parameters.get("full_attention") is None:
+ self.rope_parameters["full_attention"] = {"rope_type": "default"}
+ self.rope_parameters["full_attention"].setdefault(
+ "rope_theta", kwargs.pop("rope_theta", self.default_theta["global"])
+ )
+ if self.rope_parameters.get("sliding_attention") is None:
+ self.rope_parameters["sliding_attention"] = {"rope_type": "default"}
+ self.rope_parameters["sliding_attention"].setdefault(
+ "rope_theta", kwargs.pop("rope_local_base_freq", self.default_theta["local"])
+ )
+
+ # Standardize and validate the correctness of rotary position embeddings parameters
+ self.standardize_rope_params()
+ return kwargs
+
+
+@auto_docstring(checkpoint="google/t5gemma-2-270m-270m")
+@strict
+class T5Gemma2Config(PreTrainedConfig):
+ r"""
+ encoder (`Union[T5Gemma2EncoderConfig, dict]`, optional, *optional*):
+ Configuration for the encoder.
+ decoder (`Union[T5Gemma2DecoderConfig, dict]`, optional, *optional*):
+ Configuration for the decoder.
+ eoi_token_index (`int`, *optional*):
+ The end-of-image token index to wrap the image prompt. Will be same as
+ `self.encoder.eoi_token_index`
+
+ ```python
+ >>> from transformers import T5Gemma2Config, T5Gemma2Model
+ >>> t5gemma2_config = T5Gemma2Config.from_pretrained("google/t5gemma-270m-270m")
+ >>> model = T5Gemma2Model(t5gemma2_config)
+ ```
+ """
+
+ model_type = "t5gemma2"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ sub_configs = {
+ "encoder": T5Gemma2EncoderConfig,
+ "decoder": T5Gemma2DecoderConfig,
+ }
+
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ "eoi_token_id": "eoi_token_index",
+ }
+
+ encoder: T5Gemma2EncoderConfig | dict[str, Any] | None = None
+ decoder: T5Gemma2DecoderConfig | dict[str, Any] | None = None
+ is_encoder_decoder: bool = True
+ dropout_rate: float | int = 0.0
+ attention_dropout: float | int = 0.0
+ classifier_dropout_rate: float | int = 0.0
+ initializer_range: float = 0.02
+ image_token_index: int = 256_001
+ eoi_token_index: int | None = None
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.encoder, dict):
+ self.encoder = T5Gemma2EncoderConfig(**self.encoder)
+ elif self.encoder is None:
+ self.encoder = T5Gemma2EncoderConfig()
+ logger.info("encoder is None, using default T5Gemma2EncoderConfig encoder config.")
+
+ if isinstance(self.decoder, dict):
+ self.decoder = T5Gemma2DecoderConfig(**self.decoder)
+ elif self.decoder is None:
+ self.decoder = T5Gemma2DecoderConfig()
+ logger.info("decoder is None, using default T5Gemma2DecoderConfig decoder config.")
+
+ self.encoder.text_config.dropout_rate = self.dropout_rate
+ self.encoder.text_config.attention_dropout = self.attention_dropout
+ self.encoder.vision_config.attention_dropout = self.attention_dropout
+ self.encoder.image_token_index = self.image_token_index
+
+ self.decoder.dropout_rate = self.dropout_rate
+ self.decoder.attention_dropout = self.attention_dropout
+ self.eoi_token_index = self.encoder.eoi_token_index
+
+ for special_token_key in ["bos_token_id", "pad_token_id", "eos_token_id", "vocab_size"]:
+ if special_token_key not in kwargs:
+ kwargs[special_token_key] = getattr(self.decoder, special_token_key)
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.encoder.text_config.hidden_size != self.decoder.hidden_size:
+ raise ValueError(
+ "Imbalanced encoder-decoder is not supported in T5Gemma2: "
+ f"encoder ({self.encoder.text_config.hidden_size}) vs decoder ({self.decoder.hidden_size})."
+ )
+
+ if not self.is_encoder_decoder:
+ raise ValueError("T5Gemma2Model only support encoder-decoder modeling.")
+
+ if self.encoder.text_config.vocab_size != self.decoder.vocab_size:
+ raise ValueError(
+ "Imbalanced encoder-decoder vocabulary size is not supported in T5Gemma2: "
+ f"encoder ({self.encoder.text_config.vocab_size}) vs decoder ({self.decoder.vocab_size})."
+ )
+
+
+__all__ = ["T5Gemma2Config", "T5Gemma2TextConfig", "T5Gemma2EncoderConfig", "T5Gemma2DecoderConfig"]