diff --git a/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/INSTALLER b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/LICENSE b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..f26bcf4d2de6eb136e31006ca3ab447d5e488adf --- /dev/null +++ b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/LICENSE @@ -0,0 +1,279 @@ +A. HISTORY OF THE SOFTWARE +========================== + +Python was created in the early 1990s by Guido van Rossum at Stichting +Mathematisch Centrum (CWI, see https://www.cwi.nl) in the Netherlands +as a successor of a language called ABC. Guido remains Python's +principal author, although it includes many contributions from others. + +In 1995, Guido continued his work on Python at the Corporation for +National Research Initiatives (CNRI, see https://www.cnri.reston.va.us) +in Reston, Virginia where he released several versions of the +software. + +In May 2000, Guido and the Python core development team moved to +BeOpen.com to form the BeOpen PythonLabs team. In October of the same +year, the PythonLabs team moved to Digital Creations, which became +Zope Corporation. In 2001, the Python Software Foundation (PSF, see +https://www.python.org/psf/) was formed, a non-profit organization +created specifically to own Python-related Intellectual Property. +Zope Corporation was a sponsoring member of the PSF. + +All Python releases are Open Source (see https://opensource.org for +the Open Source Definition). Historically, most, but not all, Python +releases have also been GPL-compatible; the table below summarizes +the various releases. + + Release Derived Year Owner GPL- + from compatible? (1) + + 0.9.0 thru 1.2 1991-1995 CWI yes + 1.3 thru 1.5.2 1.2 1995-1999 CNRI yes + 1.6 1.5.2 2000 CNRI no + 2.0 1.6 2000 BeOpen.com no + 1.6.1 1.6 2001 CNRI yes (2) + 2.1 2.0+1.6.1 2001 PSF no + 2.0.1 2.0+1.6.1 2001 PSF yes + 2.1.1 2.1+2.0.1 2001 PSF yes + 2.1.2 2.1.1 2002 PSF yes + 2.1.3 2.1.2 2002 PSF yes + 2.2 and above 2.1.1 2001-now PSF yes + +Footnotes: + +(1) GPL-compatible doesn't mean that we're distributing Python under + the GPL. All Python licenses, unlike the GPL, let you distribute + a modified version without making your changes open source. The + GPL-compatible licenses make it possible to combine Python with + other software that is released under the GPL; the others don't. + +(2) According to Richard Stallman, 1.6.1 is not GPL-compatible, + because its license has a choice of law clause. According to + CNRI, however, Stallman's lawyer has told CNRI's lawyer that 1.6.1 + is "not incompatible" with the GPL. + +Thanks to the many outside volunteers who have worked under Guido's +direction to make these releases possible. + + +B. TERMS AND CONDITIONS FOR ACCESSING OR OTHERWISE USING PYTHON +=============================================================== + +Python software and documentation are licensed under the +Python Software Foundation License Version 2. + +Starting with Python 3.8.6, examples, recipes, and other code in +the documentation are dual licensed under the PSF License Version 2 +and the Zero-Clause BSD license. + +Some software incorporated into Python is under different licenses. +The licenses are listed with code falling under that license. + + +PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +-------------------------------------------- + +1. This LICENSE AGREEMENT is between the Python Software Foundation +("PSF"), and the Individual or Organization ("Licensee") accessing and +otherwise using this software ("Python") in source or binary form and +its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, PSF hereby +grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce, +analyze, test, perform and/or display publicly, prepare derivative works, +distribute, and otherwise use Python alone or in any derivative version, +provided, however, that PSF's License Agreement and PSF's notice of copyright, +i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, +2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022, 2023 Python Software Foundation; +All Rights Reserved" are retained in Python alone or in any derivative version +prepared by Licensee. + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python. + +4. PSF is making Python available to Licensee on an "AS IS" +basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between PSF and +Licensee. This License Agreement does not grant permission to use PSF +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using Python, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +BEOPEN.COM LICENSE AGREEMENT FOR PYTHON 2.0 +------------------------------------------- + +BEOPEN PYTHON OPEN SOURCE LICENSE AGREEMENT VERSION 1 + +1. This LICENSE AGREEMENT is between BeOpen.com ("BeOpen"), having an +office at 160 Saratoga Avenue, Santa Clara, CA 95051, and the +Individual or Organization ("Licensee") accessing and otherwise using +this software in source or binary form and its associated +documentation ("the Software"). + +2. Subject to the terms and conditions of this BeOpen Python License +Agreement, BeOpen hereby grants Licensee a non-exclusive, +royalty-free, world-wide license to reproduce, analyze, test, perform +and/or display publicly, prepare derivative works, distribute, and +otherwise use the Software alone or in any derivative version, +provided, however, that the BeOpen Python License is retained in the +Software, alone or in any derivative version prepared by Licensee. + +3. BeOpen is making the Software available to Licensee on an "AS IS" +basis. BEOPEN MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, BEOPEN MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF THE SOFTWARE WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +4. BEOPEN SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF THE +SOFTWARE FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS +AS A RESULT OF USING, MODIFYING OR DISTRIBUTING THE SOFTWARE, OR ANY +DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +5. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +6. This License Agreement shall be governed by and interpreted in all +respects by the law of the State of California, excluding conflict of +law provisions. Nothing in this License Agreement shall be deemed to +create any relationship of agency, partnership, or joint venture +between BeOpen and Licensee. This License Agreement does not grant +permission to use BeOpen trademarks or trade names in a trademark +sense to endorse or promote products or services of Licensee, or any +third party. As an exception, the "BeOpen Python" logos available at +http://www.pythonlabs.com/logos.html may be used according to the +permissions granted on that web page. + +7. By copying, installing or otherwise using the software, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +CNRI LICENSE AGREEMENT FOR PYTHON 1.6.1 +--------------------------------------- + +1. This LICENSE AGREEMENT is between the Corporation for National +Research Initiatives, having an office at 1895 Preston White Drive, +Reston, VA 20191 ("CNRI"), and the Individual or Organization +("Licensee") accessing and otherwise using Python 1.6.1 software in +source or binary form and its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, CNRI +hereby grants Licensee a nonexclusive, royalty-free, world-wide +license to reproduce, analyze, test, perform and/or display publicly, +prepare derivative works, distribute, and otherwise use Python 1.6.1 +alone or in any derivative version, provided, however, that CNRI's +License Agreement and CNRI's notice of copyright, i.e., "Copyright (c) +1995-2001 Corporation for National Research Initiatives; All Rights +Reserved" are retained in Python 1.6.1 alone or in any derivative +version prepared by Licensee. Alternately, in lieu of CNRI's License +Agreement, Licensee may substitute the following text (omitting the +quotes): "Python 1.6.1 is made available subject to the terms and +conditions in CNRI's License Agreement. This Agreement together with +Python 1.6.1 may be located on the internet using the following +unique, persistent identifier (known as a handle): 1895.22/1013. This +Agreement may also be obtained from a proxy server on the internet +using the following URL: http://hdl.handle.net/1895.22/1013". + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python 1.6.1 or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python 1.6.1. + +4. CNRI is making Python 1.6.1 available to Licensee on an "AS IS" +basis. CNRI MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, CNRI MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON 1.6.1 WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. CNRI SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +1.6.1 FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON 1.6.1, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. This License Agreement shall be governed by the federal +intellectual property law of the United States, including without +limitation the federal copyright law, and, to the extent such +U.S. federal law does not apply, by the law of the Commonwealth of +Virginia, excluding Virginia's conflict of law provisions. +Notwithstanding the foregoing, with regard to derivative works based +on Python 1.6.1 that incorporate non-separable material that was +previously distributed under the GNU General Public License (GPL), the +law of the Commonwealth of Virginia shall govern this License +Agreement only as to issues arising under or with respect to +Paragraphs 4, 5, and 7 of this License Agreement. Nothing in this +License Agreement shall be deemed to create any relationship of +agency, partnership, or joint venture between CNRI and Licensee. This +License Agreement does not grant permission to use CNRI trademarks or +trade name in a trademark sense to endorse or promote products or +services of Licensee, or any third party. + +8. By clicking on the "ACCEPT" button where indicated, or by copying, +installing or otherwise using Python 1.6.1, Licensee agrees to be +bound by the terms and conditions of this License Agreement. + + ACCEPT + + +CWI LICENSE AGREEMENT FOR PYTHON 0.9.0 THROUGH 1.2 +-------------------------------------------------- + +Copyright (c) 1991 - 1995, Stichting Mathematisch Centrum Amsterdam, +The Netherlands. All rights reserved. + +Permission to use, copy, modify, and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appear in all copies and that +both that copyright notice and this permission notice appear in +supporting documentation, and that the name of Stichting Mathematisch +Centrum or CWI not be used in advertising or publicity pertaining to +distribution of the software without specific, written prior +permission. + +STICHTING MATHEMATISCH CENTRUM DISCLAIMS ALL WARRANTIES WITH REGARD TO +THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS, IN NO EVENT SHALL STICHTING MATHEMATISCH CENTRUM BE LIABLE +FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT +OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +ZERO-CLAUSE BSD LICENSE FOR CODE IN THE PYTHON DOCUMENTATION +---------------------------------------------------------------------- + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/METADATA b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c632040d66bf120a377fc3785940934361273a66 --- /dev/null +++ b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/METADATA @@ -0,0 +1,123 @@ +Metadata-Version: 2.3 +Name: aiohappyeyeballs +Version: 2.6.1 +Summary: Happy Eyeballs for asyncio +License: PSF-2.0 +Author: J. Nick Koston +Author-email: nick@koston.org +Requires-Python: >=3.9 +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Natural Language :: English +Classifier: Operating System :: OS Independent +Classifier: Topic :: Software Development :: Libraries +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: License :: OSI Approved :: Python Software Foundation License +Project-URL: Bug Tracker, https://github.com/aio-libs/aiohappyeyeballs/issues +Project-URL: Changelog, https://github.com/aio-libs/aiohappyeyeballs/blob/main/CHANGELOG.md +Project-URL: Documentation, https://aiohappyeyeballs.readthedocs.io +Project-URL: Repository, https://github.com/aio-libs/aiohappyeyeballs +Description-Content-Type: text/markdown + +# aiohappyeyeballs + +

+ + CI Status + + + Documentation Status + + + Test coverage percentage + +

+

+ + Poetry + + + Ruff + + + pre-commit + +

+

+ + PyPI Version + + Supported Python versions + License +

+ +--- + +**Documentation**: https://aiohappyeyeballs.readthedocs.io + +**Source Code**: https://github.com/aio-libs/aiohappyeyeballs + +--- + +[Happy Eyeballs](https://en.wikipedia.org/wiki/Happy_Eyeballs) +([RFC 8305](https://www.rfc-editor.org/rfc/rfc8305.html)) + +## Use case + +This library exists to allow connecting with +[Happy Eyeballs](https://en.wikipedia.org/wiki/Happy_Eyeballs) +([RFC 8305](https://www.rfc-editor.org/rfc/rfc8305.html)) +when you +already have a list of addrinfo and not a DNS name. + +The stdlib version of `loop.create_connection()` +will only work when you pass in an unresolved name which +is not a good fit when using DNS caching or resolving +names via another method such as `zeroconf`. + +## Installation + +Install this via pip (or your favourite package manager): + +`pip install aiohappyeyeballs` + +## License + +[aiohappyeyeballs is licensed under the same terms as cpython itself.](https://github.com/python/cpython/blob/main/LICENSE) + +## Example usage + +```python + +addr_infos = await loop.getaddrinfo("example.org", 80) + +socket = await start_connection(addr_infos) +socket = await start_connection(addr_infos, local_addr_infos=local_addr_infos, happy_eyeballs_delay=0.2) + +transport, protocol = await loop.create_connection( + MyProtocol, sock=socket, ...) + +# Remove the first address for each family from addr_info +pop_addr_infos_interleave(addr_info, 1) + +# Remove all matching address from addr_info +remove_addr_infos(addr_info, "dead::beef::") + +# Convert a local_addr to local_addr_infos +local_addr_infos = addr_to_addr_infos(("127.0.0.1",0)) +``` + +## Credits + +This package contains code from cpython and is licensed under the same terms as cpython itself. + +This package was created with +[Copier](https://copier.readthedocs.io/) and the +[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template) +project template. + diff --git a/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/RECORD b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..9e0e2023f21e2df03b9b2d75aae5b116887feede --- /dev/null +++ b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/RECORD @@ -0,0 +1,16 @@ +aiohappyeyeballs-2.6.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +aiohappyeyeballs-2.6.1.dist-info/LICENSE,sha256=Oy-B_iHRgcSZxZolbI4ZaEVdZonSaaqFNzv7avQdo78,13936 +aiohappyeyeballs-2.6.1.dist-info/METADATA,sha256=NSXlhJwAfi380eEjAo7BQ4P_TVal9xi0qkyZWibMsVM,5915 +aiohappyeyeballs-2.6.1.dist-info/RECORD,, +aiohappyeyeballs-2.6.1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88 +aiohappyeyeballs/__init__.py,sha256=x7kktHEtaD9quBcWDJPuLeKyjuVAI-Jj14S9B_5hcTs,361 +aiohappyeyeballs/__pycache__/__init__.cpython-312.pyc,, +aiohappyeyeballs/__pycache__/_staggered.cpython-312.pyc,, +aiohappyeyeballs/__pycache__/impl.cpython-312.pyc,, +aiohappyeyeballs/__pycache__/types.cpython-312.pyc,, +aiohappyeyeballs/__pycache__/utils.cpython-312.pyc,, +aiohappyeyeballs/_staggered.py,sha256=edfVowFx-P-ywJjIEF3MdPtEMVODujV6CeMYr65otac,6900 +aiohappyeyeballs/impl.py,sha256=Dlcm2mTJ28ucrGnxkb_fo9CZzLAkOOBizOt7dreBbXE,9681 +aiohappyeyeballs/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +aiohappyeyeballs/types.py,sha256=YZJIAnyoV4Dz0WFtlaf_OyE4EW7Xus1z7aIfNI6tDDQ,425 +aiohappyeyeballs/utils.py,sha256=on9GxIR0LhEfZu8P6Twi9hepX9zDanuZM20MWsb3xlQ,3028 diff --git a/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/WHEEL b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..0582547b15f02d3a51659106262832565d5dc5ea --- /dev/null +++ b/lib/python3.12/site-packages/aiohappyeyeballs-2.6.1.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: poetry-core 2.1.1 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/RECORD b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..5c99d84368090819b30a1af0cfb3e570c16119ad --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/RECORD @@ -0,0 +1,689 @@ +../../../bin/fonttools,sha256=NKIcaBcdPgzS3nnfEpl3yn4gvv1xUv8GdqAn2pnRLaw,253 +../../../bin/pyftmerge,sha256=WvFMdnNLpg3c-hj-PIDoyRUDiulZoS5-eVzr0WuuN6I,250 +../../../bin/pyftsubset,sha256=odvcECPVzBH4Qe03IPDPW2V00cI2MpQHJUpbIHeJ3y4,251 +../../../bin/ttx,sha256=qce32mvSYKqwLyqFH3adWZCHWv5y_2jIw43AhMKJphM,248 +../../../share/man/man1/ttx.1,sha256=cLbm_pOOj1C76T2QXvDxzwDj9gk-GTd5RztvTMsouFw,5377 +fontTools/__init__.py,sha256=w-yY9D2nPQeBbS3lLIyNrc-4KGqgArD4IhQXdd8qioQ,183 +fontTools/__main__.py,sha256=VjkGh1UD-i1zTDA1dXo1uecSs6PxHdGQ5vlCk_mCCYs,925 +fontTools/__pycache__/__init__.cpython-312.pyc,, +fontTools/__pycache__/__main__.cpython-312.pyc,, +fontTools/__pycache__/afmLib.cpython-312.pyc,, +fontTools/__pycache__/agl.cpython-312.pyc,, +fontTools/__pycache__/annotations.cpython-312.pyc,, +fontTools/__pycache__/fontBuilder.cpython-312.pyc,, +fontTools/__pycache__/help.cpython-312.pyc,, +fontTools/__pycache__/tfmLib.cpython-312.pyc,, +fontTools/__pycache__/ttx.cpython-312.pyc,, +fontTools/__pycache__/unicode.cpython-312.pyc,, +fontTools/afmLib.py,sha256=1MagIItOzRV4vV5kKPxeDZbPJsfxLB3wdHLFkQvl0uk,13164 +fontTools/agl.py,sha256=05bm8Uq45uVWW8nPbP6xbNgmFyxQr8sWhYAiP0VSjnI,112975 +fontTools/annotations.py,sha256=BdIIriNYDzBfgniwWFg_u71qLZnV0sCcn4-VAkXkYNM,1225 +fontTools/cffLib/CFF2ToCFF.py,sha256=7h0fK53ji4JlA4mAxlVb54eQTM4N2e3LAd4wwW3TLyE,8201 +fontTools/cffLib/CFFToCFF2.py,sha256=Qnk7lYlsTRHnlZQ6NXNdr_f4MJwZQ21kcS08KFbsyY8,10119 +fontTools/cffLib/__init__.py,sha256=62vpcR7u8cE407kXduAwnFttHnsoCpDQ7IBK-qOYFQ8,107886 +fontTools/cffLib/__pycache__/CFF2ToCFF.cpython-312.pyc,, +fontTools/cffLib/__pycache__/CFFToCFF2.cpython-312.pyc,, +fontTools/cffLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/cffLib/__pycache__/specializer.cpython-312.pyc,, +fontTools/cffLib/__pycache__/transforms.cpython-312.pyc,, +fontTools/cffLib/__pycache__/width.cpython-312.pyc,, +fontTools/cffLib/specializer.py,sha256=vsOPkR_jHNe6tESQEjmm0i76y7sWI5MKo3bsTmI3sNM,32609 +fontTools/cffLib/transforms.py,sha256=SEIZc1XxWYiVXVBsoNm6LTvM9SUN7Z76QOaSAlR1ZCo,17455 +fontTools/cffLib/width.py,sha256=IqGL0CLyCZqi_hvsHySG08qpYxS3kaqW-tsAT-bjHV4,6074 +fontTools/colorLib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +fontTools/colorLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/colorLib/__pycache__/builder.cpython-312.pyc,, +fontTools/colorLib/__pycache__/errors.cpython-312.pyc,, +fontTools/colorLib/__pycache__/geometry.cpython-312.pyc,, +fontTools/colorLib/__pycache__/table_builder.cpython-312.pyc,, +fontTools/colorLib/__pycache__/unbuilder.cpython-312.pyc,, +fontTools/colorLib/builder.py,sha256=kmO7OuudQQb3fEOS7aLzgTDVjqS9i2xIQmk9p1uBe8A,23008 +fontTools/colorLib/errors.py,sha256=CsaviiRxxrpgVX4blm7KCyK8553ljwL44xkJOeC5U7U,41 +fontTools/colorLib/geometry.py,sha256=3ScySrR2YDJa7d5K5_xM5Yt1-3NCV-ry8ikYA5VwVbI,5518 +fontTools/colorLib/table_builder.py,sha256=ZeltWY6n-YPiJv_hQ1iBXoEFAG70EKxZyScgsMKUFGU,7469 +fontTools/colorLib/unbuilder.py,sha256=iW-E5I39WsV82K3NgCO4Cjzwm1WqzGrtypHt8epwbHM,2142 +fontTools/config/__init__.py,sha256=JICOHIz06KuHCiBmrxj-ga19P6ZTuLXh0lHmPh-Ra1w,3154 +fontTools/config/__pycache__/__init__.cpython-312.pyc,, +fontTools/cu2qu/__init__.py,sha256=Cuc7Uglb0nSgaraTxXY5J8bReznH5wApW0uakN7MycY,618 +fontTools/cu2qu/__main__.py,sha256=kTUI-jczsHeelULLlory74QEeFjZWp9zigCc7PrdVQY,92 +fontTools/cu2qu/__pycache__/__init__.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/__main__.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/benchmark.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/cli.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/cu2qu.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/errors.cpython-312.pyc,, +fontTools/cu2qu/__pycache__/ufo.cpython-312.pyc,, +fontTools/cu2qu/benchmark.py,sha256=wasPJmf8q9k9UHjpHChC3WQAGbBAyHN9PvJzXvWC0Fw,1296 +fontTools/cu2qu/cli.py,sha256=MbAQnOpZwrUFe_tjAP3Tgf6uLdOgHlONUcPNeTXwH0Y,6076 +fontTools/cu2qu/cu2qu.c,sha256=Hacw9XncC5wYjSd-uT4MxoOQGDBQFnG8mLZrohGHtXQ,648759 +fontTools/cu2qu/cu2qu.cpython-312-x86_64-linux-gnu.so,sha256=IqlIidvItzEAjZ5Lko1FBJUsMj-oWloXxop80oBShRw,1082272 +fontTools/cu2qu/cu2qu.py,sha256=6LTe1ZI-jxW8y79s_UFjbkeFoFleIekTLm2jAE-uqGQ,17986 +fontTools/cu2qu/errors.py,sha256=PyJNMy8lHDtKpfFkc0nkM8F4jNLZAC4lPQCN1Km4bpg,2441 +fontTools/cu2qu/ufo.py,sha256=g2NkcMqdS-t3tBtvcqtt4E02QVyQr3pTC1qXDRtI31U,12307 +fontTools/designspaceLib/__init__.py,sha256=whzi46QY0FK77k9iE646rGP4Ivlq9BEXKO0y7UU9Plo,129502 +fontTools/designspaceLib/__main__.py,sha256=xhtYXo1T1tsykhQDD0tcconSNYgWL5hoTBORpVDUYrc,103 +fontTools/designspaceLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/designspaceLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/designspaceLib/__pycache__/split.cpython-312.pyc,, +fontTools/designspaceLib/__pycache__/statNames.cpython-312.pyc,, +fontTools/designspaceLib/__pycache__/types.cpython-312.pyc,, +fontTools/designspaceLib/split.py,sha256=FB1NuvhUO453UXveQZi9oyrW_caoCPM3RADp1rYWkDs,19239 +fontTools/designspaceLib/statNames.py,sha256=gXGKWVr1ju2_oL-R_DkyoZ3GlI7mfLORovpk1Ebgmvc,9237 +fontTools/designspaceLib/types.py,sha256=ofK65qXNADqcpl7zI72Pa5s07-cm7G41iEmLVV44-Es,5320 +fontTools/encodings/MacRoman.py,sha256=4vEooUDm2gLCG8KIIDhRxm5-A64w7XrhP9cjDRr2Eo0,3576 +fontTools/encodings/StandardEncoding.py,sha256=Eo3AGE8FE_p-IVYYuV097KouSsF3UrXoRRN0XyvYbrs,3581 +fontTools/encodings/__init__.py,sha256=DJBWmoX_Haau7qlgmvWyfbhSzrX2qL636Rns7CG01pk,75 +fontTools/encodings/__pycache__/MacRoman.cpython-312.pyc,, +fontTools/encodings/__pycache__/StandardEncoding.cpython-312.pyc,, +fontTools/encodings/__pycache__/__init__.cpython-312.pyc,, +fontTools/encodings/__pycache__/codecs.cpython-312.pyc,, +fontTools/encodings/codecs.py,sha256=u50ruwz9fcRsrUrRGpR17Cr55Ovn1fvCHCKrElVumDE,4721 +fontTools/feaLib/__init__.py,sha256=jlIru2ghxvb1HhC5Je2BCXjFJmFQlYKpruorPoz3BvQ,213 +fontTools/feaLib/__main__.py,sha256=Df2PA6LXwna98lSXiL7R4as_ZEdWCIk3egSM5w7GpvM,2240 +fontTools/feaLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/feaLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/feaLib/__pycache__/ast.cpython-312.pyc,, +fontTools/feaLib/__pycache__/builder.cpython-312.pyc,, +fontTools/feaLib/__pycache__/error.cpython-312.pyc,, +fontTools/feaLib/__pycache__/lexer.cpython-312.pyc,, +fontTools/feaLib/__pycache__/location.cpython-312.pyc,, +fontTools/feaLib/__pycache__/lookupDebugInfo.cpython-312.pyc,, +fontTools/feaLib/__pycache__/parser.cpython-312.pyc,, +fontTools/feaLib/__pycache__/variableScalar.cpython-312.pyc,, +fontTools/feaLib/ast.py,sha256=48Y6vpSD_wYfucWyh_bQtzf2AQFX-pOwBvsxdcpVDz0,74158 +fontTools/feaLib/builder.py,sha256=OI3slR3tTlH1LyWGHJMDw3jol5PzvRd8_T5j1603xZM,73968 +fontTools/feaLib/error.py,sha256=Bz_5tNcNVcY7_nrAmFlQNhQldtqZWd8WUGQ2E3PWhZo,648 +fontTools/feaLib/lexer.c,sha256=LWJfBpHqWqXW1IAGjiFT7pcXWmJWX55FjrQPvN3d9uo,747744 +fontTools/feaLib/lexer.cpython-312-x86_64-linux-gnu.so,sha256=C1G2mAF0TH-ZfPsd5PYxiTQTeuEcJpSy-0m34L75XEE,1388240 +fontTools/feaLib/lexer.py,sha256=emyMPmRoqNZkzxnJyI6JRCCtXrbCOFofwa9O6ABGLiw,11121 +fontTools/feaLib/location.py,sha256=JXzHqGV56EHdcq823AwA5oaK05hf_1ySWpScbo3zGC0,234 +fontTools/feaLib/lookupDebugInfo.py,sha256=gVRr5-APWfT_a5-25hRuawSVX8fEvXVsOSLWkH91T2w,304 +fontTools/feaLib/parser.py,sha256=eM4ph-_6ab8EpAWqg9Ax7is5We-KwkzJYlui31mJSBc,99911 +fontTools/feaLib/variableScalar.py,sha256=f6sOg9cfFJRI3fw04uRohDeFux0xnZanaPT_lcxAVOw,4200 +fontTools/fontBuilder.py,sha256=d9bBaId9K6KPpR5nP0l5Ou_8Pp7hr4gRuLeQV0NmwUY,34280 +fontTools/help.py,sha256=bAjatvIhV7TJyXI7WhsxdYO4YVlhScZXu_kRtHANEPo,1125 +fontTools/merge/__init__.py,sha256=8i6ownyQTAOBKWnTEHvvCYFw64Mv7Z1HPBgJI-ZiuKo,8256 +fontTools/merge/__main__.py,sha256=hDx3gfbUBO83AJKumSEhiV-xqNTJNNgK2uFjazOGTmw,94 +fontTools/merge/__pycache__/__init__.cpython-312.pyc,, +fontTools/merge/__pycache__/__main__.cpython-312.pyc,, +fontTools/merge/__pycache__/base.cpython-312.pyc,, +fontTools/merge/__pycache__/cmap.cpython-312.pyc,, +fontTools/merge/__pycache__/layout.cpython-312.pyc,, +fontTools/merge/__pycache__/options.cpython-312.pyc,, +fontTools/merge/__pycache__/tables.cpython-312.pyc,, +fontTools/merge/__pycache__/unicode.cpython-312.pyc,, +fontTools/merge/__pycache__/util.cpython-312.pyc,, +fontTools/merge/base.py,sha256=l0G1Px98E9ZdVuFLMUBKWdtr7Jb8JX8vxcjeaDUUnzY,2389 +fontTools/merge/cmap.py,sha256=HpthxVH5lA7VegJ8yHoBjd9vrFBV7UB5OknKGYpxWY8,6728 +fontTools/merge/layout.py,sha256=fkMPGPLxEdxohS3scVM4W7LmNthSz-UPyocsffe2KqE,16075 +fontTools/merge/options.py,sha256=xko_1-WErcNQkirECzIOOYxSJR_bRtdQYQYOtmgccYI,2501 +fontTools/merge/tables.py,sha256=7SzXYL04awDEDhvU2-9T_8A2gAjvgGyYAHUICUJOpZg,10958 +fontTools/merge/unicode.py,sha256=kb1Jrfuoq1KUcVhhSKnflAED_wMZxXDjVwB-CI9k05Y,4273 +fontTools/merge/util.py,sha256=BH3bZWNFy-Tsj1cth7aSpGVJ18YXKXqDakPn6Wzku6U,3378 +fontTools/misc/__init__.py,sha256=DJBWmoX_Haau7qlgmvWyfbhSzrX2qL636Rns7CG01pk,75 +fontTools/misc/__pycache__/__init__.cpython-312.pyc,, +fontTools/misc/__pycache__/arrayTools.cpython-312.pyc,, +fontTools/misc/__pycache__/bezierTools.cpython-312.pyc,, +fontTools/misc/__pycache__/classifyTools.cpython-312.pyc,, +fontTools/misc/__pycache__/cliTools.cpython-312.pyc,, +fontTools/misc/__pycache__/configTools.cpython-312.pyc,, +fontTools/misc/__pycache__/cython.cpython-312.pyc,, +fontTools/misc/__pycache__/dictTools.cpython-312.pyc,, +fontTools/misc/__pycache__/eexec.cpython-312.pyc,, +fontTools/misc/__pycache__/encodingTools.cpython-312.pyc,, +fontTools/misc/__pycache__/enumTools.cpython-312.pyc,, +fontTools/misc/__pycache__/etree.cpython-312.pyc,, +fontTools/misc/__pycache__/filenames.cpython-312.pyc,, +fontTools/misc/__pycache__/fixedTools.cpython-312.pyc,, +fontTools/misc/__pycache__/intTools.cpython-312.pyc,, +fontTools/misc/__pycache__/iterTools.cpython-312.pyc,, +fontTools/misc/__pycache__/lazyTools.cpython-312.pyc,, +fontTools/misc/__pycache__/loggingTools.cpython-312.pyc,, +fontTools/misc/__pycache__/macCreatorType.cpython-312.pyc,, +fontTools/misc/__pycache__/macRes.cpython-312.pyc,, +fontTools/misc/__pycache__/psCharStrings.cpython-312.pyc,, +fontTools/misc/__pycache__/psLib.cpython-312.pyc,, +fontTools/misc/__pycache__/psOperators.cpython-312.pyc,, +fontTools/misc/__pycache__/py23.cpython-312.pyc,, +fontTools/misc/__pycache__/roundTools.cpython-312.pyc,, +fontTools/misc/__pycache__/sstruct.cpython-312.pyc,, +fontTools/misc/__pycache__/symfont.cpython-312.pyc,, +fontTools/misc/__pycache__/testTools.cpython-312.pyc,, +fontTools/misc/__pycache__/textTools.cpython-312.pyc,, +fontTools/misc/__pycache__/timeTools.cpython-312.pyc,, +fontTools/misc/__pycache__/transform.cpython-312.pyc,, +fontTools/misc/__pycache__/treeTools.cpython-312.pyc,, +fontTools/misc/__pycache__/vector.cpython-312.pyc,, +fontTools/misc/__pycache__/visitor.cpython-312.pyc,, +fontTools/misc/__pycache__/xmlReader.cpython-312.pyc,, +fontTools/misc/__pycache__/xmlWriter.cpython-312.pyc,, +fontTools/misc/arrayTools.py,sha256=jZk__GE-K9VViZE_H-LPPj0smWbKng-yfPE8BfGp8HI,11483 +fontTools/misc/bezierTools.c,sha256=7ymo3Ld3xnDe1fJBXiqMVMFzoYRc5hY9Qxby9g3r5XI,1836341 +fontTools/misc/bezierTools.cpython-312-x86_64-linux-gnu.so,sha256=TiZqyjMhcxmhLrG4DzMkKDcx0sR_6dyo7uqXbq6C9NA,4786264 +fontTools/misc/bezierTools.py,sha256=2xvnsKboUZ-Ln1bbiJk2392foh3vRhulmpqaWl1a-lQ,45240 +fontTools/misc/classifyTools.py,sha256=zcg3EM4GOerBW9c063ljaLllgeeZ772EpFZjp9CdgLI,5613 +fontTools/misc/cliTools.py,sha256=qCznJMLCQu3ZHQD_4ctUnr3TkfAUdkGl-UuxZUrppy0,1862 +fontTools/misc/configTools.py,sha256=tDede_q8h81w6mDh132_F5WzfWkwmDGAF8nwj3dHVus,11229 +fontTools/misc/cython.py,sha256=eyLcL2Bw-SSToYro8f44dkkYRlQfiFbhcza0afS-qHE,682 +fontTools/misc/dictTools.py,sha256=VxjarsGJuk_wa3z29FSCtKZNCFfXtMBiNEu0RPAlpDk,2417 +fontTools/misc/eexec.py,sha256=GNn2OCRvO1HbbIeDPxk9i0glO7cux_AQaoVMXhBR8y8,3331 +fontTools/misc/encodingTools.py,sha256=hCv5PFfnXQJVCZA8Wyn1vr3vzLBbUuEPtGk5CzWM9RY,2073 +fontTools/misc/enumTools.py,sha256=YQZW-d2ES9KFFkAXOUMIBbRUk6v_3BT6Q7lXE1ufhxA,502 +fontTools/misc/etree.py,sha256=ZzJc6TvAS579deAgZLVDvTY_HeTm-ZsKJ5s3LYhZSSY,16304 +fontTools/misc/filenames.py,sha256=MMCO3xjk1pcDc-baobcKd8IdoFPt-bcGqu8t8HUGAkI,8223 +fontTools/misc/filesystem/__init__.py,sha256=iwoOj6DpXKk8q-NRRHqOfRxFF6lcXIhsIA46j-cZswU,2011 +fontTools/misc/filesystem/__pycache__/__init__.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_base.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_copy.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_errors.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_info.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_osfs.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_path.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_subfs.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_tempfs.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_tools.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_walk.cpython-312.pyc,, +fontTools/misc/filesystem/__pycache__/_zipfs.cpython-312.pyc,, +fontTools/misc/filesystem/_base.py,sha256=p74O7xREadfPQgzPJ9mP3ehu0ZHDgsmXlpsL9CnTRso,4010 +fontTools/misc/filesystem/_copy.py,sha256=ifMSs-A_bz1Aa4tIQrlUd9HtdJQ5fp5M3B6mbYuDtXI,1361 +fontTools/misc/filesystem/_errors.py,sha256=-YziRB1BT1I80ypmufvCR-M_4XoerCHyBVqX-cRnIzU,641 +fontTools/misc/filesystem/_info.py,sha256=pbV7bDTJ5F8ms6alK34J0FZYWzmRO7FT0NM3yRA3czo,2013 +fontTools/misc/filesystem/_osfs.py,sha256=RkKCE2IxcRaj7gyqFW10LhEm_-VJYtXxsS5s0DCXihM,5737 +fontTools/misc/filesystem/_path.py,sha256=frP6ZLmMeP9E3NiwoCbbgBPWpQbLBRh7T-0vOE-EuPo,1745 +fontTools/misc/filesystem/_subfs.py,sha256=vRotQwyLVfINbR88xIBQUbq9j4Kmg1_mJvEhnpvK_t4,3028 +fontTools/misc/filesystem/_tempfs.py,sha256=9FUXdCBTwFtZMFx8ghYuZVYoQdDb0tDB-jXNu3D-Qy0,924 +fontTools/misc/filesystem/_tools.py,sha256=r75dpadp7C9EdQ6r7pJQKZlCZDUJzVq2ikb_LXN-wCI,972 +fontTools/misc/filesystem/_walk.py,sha256=KMQ-GavWYr4SsA5V8ohLPmz3boilvY2P0JKrLoxW6NU,1655 +fontTools/misc/filesystem/_zipfs.py,sha256=i3qolbkDRntB_oL3v79KuEgfVlVojecPBnBA0X04PWc,6301 +fontTools/misc/fixedTools.py,sha256=XycD5QpaejJE-GFaFMlLlku6RG0ipJowuamMXjFrMjQ,7668 +fontTools/misc/intTools.py,sha256=l6pjk4UYlXcyLtfC0DdOC5RL6UJ8ihRR0zRiYow5xA8,586 +fontTools/misc/iterTools.py,sha256=17H6LPZszp32bTKoNorp6uZF1PKj47BAbe5QG8irUjo,390 +fontTools/misc/lazyTools.py,sha256=BC6MmF-OzJ3GrBD8TYDZ-VCSN4UOx0pN0r3oF4GSoiw,1020 +fontTools/misc/loggingTools.py,sha256=NOYROsLK5TzONK5967OGdVonNyXC6kP_CmPr7M2PW_c,19933 +fontTools/misc/macCreatorType.py,sha256=Je9jtqUr7EPbpH3QxlVl3pizoQ-1AOPMBIctHIMTM3k,1593 +fontTools/misc/macRes.py,sha256=GT_pnfPw2NCvvOF86nHLAnOtZ6SMHqEuLntaplXzvHM,8579 +fontTools/misc/plistlib/__init__.py,sha256=1HfhHPt3As6u2eRSlFfl6XdnXv_ypQImeQdWIw6wK7Y,21113 +fontTools/misc/plistlib/__pycache__/__init__.cpython-312.pyc,, +fontTools/misc/plistlib/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +fontTools/misc/psCharStrings.py,sha256=Zn8mr7NRTfEJbnaXbOnSarrFBPcvPM0AupAxF2C80LY,43468 +fontTools/misc/psLib.py,sha256=ioIPm5x3MHkBXF2vzNkC4iVZYobrkWcyvFhmYsjOrPY,12099 +fontTools/misc/psOperators.py,sha256=9SLl5PPBulLo0Xxg_dqlJMitNIBdiGKdkXhOWsNSYZE,15700 +fontTools/misc/py23.py,sha256=aPVCEUz_deggwLBCeTSsccX6QgJavZqvdVtuhpzrPvA,2238 +fontTools/misc/roundTools.py,sha256=1RSXZ0gyi1qW42tz6WSBMJD1FlPdtgqKfWixVN9bd78,3173 +fontTools/misc/sstruct.py,sha256=vUODd2CKvHLtjr7yn1K94Hui_yxOPWKmlgAmBMm3KDQ,7009 +fontTools/misc/symfont.py,sha256=x5ZwqK9Ik9orG6qSftgVGygBFE1wTSngrMK2We1Z5AM,6977 +fontTools/misc/testTools.py,sha256=3vj_KllUQVEiVFbS0SzTmeuKv44-L-disI1dZ4XhOfw,7052 +fontTools/misc/textTools.py,sha256=wNjH5zl1v9qNfmTl4BL52IO2IG1H5xY3o_pslqPPRjc,3483 +fontTools/misc/timeTools.py,sha256=e9h5pgzL04tBDXmCv_8eRGB4boFV8GKXlS6dq3ggEpw,2234 +fontTools/misc/transform.py,sha256=OR8dPsAw87z77gkZQMq00iUkDWLIxYv-12XiKH1-erk,15798 +fontTools/misc/treeTools.py,sha256=tLWkwyDHeZUPVOGNnJeD4Pn7x2bQeZetwJKaEAW2J2M,1269 +fontTools/misc/vector.py,sha256=6lqZcDjAgHJFQgjzD-ULQ_PrigAMfeZKaBZmAfcC0ig,4062 +fontTools/misc/visitor.py,sha256=zwBAVfZ3MTsrbhNFj03pSSjNRyT6oGkare-kfWkN5ns,5754 +fontTools/misc/xmlReader.py,sha256=igut4_d13RT4WarliqVvuuPybO1uSXVeoBOeW4j0_e4,6580 +fontTools/misc/xmlWriter.py,sha256=EIdcS_0y9gvpc82QvZC9VcerAqEsP8eeC0gqMoPK4jc,7374 +fontTools/mtiLib/__init__.py,sha256=EzYwNaENLf906h1THBeq6nSRHUKpOAYxuzO9x9PHzh8,46602 +fontTools/mtiLib/__main__.py,sha256=gd8X89jnZOe-752k7uaR1lWoiju-2zIT5Yx35Kl0Xek,94 +fontTools/mtiLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/mtiLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/otlLib/__init__.py,sha256=D2leUW-3gsUTOFcJYGC18edBYjIJ804ut4qitJYWsaQ,45 +fontTools/otlLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/otlLib/__pycache__/builder.cpython-312.pyc,, +fontTools/otlLib/__pycache__/error.cpython-312.pyc,, +fontTools/otlLib/__pycache__/maxContextCalc.cpython-312.pyc,, +fontTools/otlLib/builder.py,sha256=DEhmbyVzsPFwCDrxgIjfOaZSIcKcVmCwrg9KMxkksEs,129088 +fontTools/otlLib/error.py,sha256=cthuhBuOwZYpkTLi5gFPupUxkXkCHe-L_YgkE7N1wCI,335 +fontTools/otlLib/maxContextCalc.py,sha256=3es4Kt84TaZ49sA2ev1zrlwPJikJCAECx5KavwhyB-I,3175 +fontTools/otlLib/optimize/__init__.py,sha256=UUQRpNkHU2RczCRt-Gz7sEiYE9AQq9BHLXZEOyvsnX4,1530 +fontTools/otlLib/optimize/__main__.py,sha256=BvP472kA9KxBb9RMyyehPNevAfpmgW9MfdazkUiAO3M,104 +fontTools/otlLib/optimize/__pycache__/__init__.cpython-312.pyc,, +fontTools/otlLib/optimize/__pycache__/__main__.cpython-312.pyc,, +fontTools/otlLib/optimize/__pycache__/gpos.cpython-312.pyc,, +fontTools/otlLib/optimize/gpos.py,sha256=htOgSP743DZDUKF3eWAeJ-kdqNYOnpGXdlV-rbEXJ1A,17668 +fontTools/pens/__init__.py,sha256=DJBWmoX_Haau7qlgmvWyfbhSzrX2qL636Rns7CG01pk,75 +fontTools/pens/__pycache__/__init__.cpython-312.pyc,, +fontTools/pens/__pycache__/areaPen.cpython-312.pyc,, +fontTools/pens/__pycache__/basePen.cpython-312.pyc,, +fontTools/pens/__pycache__/boundsPen.cpython-312.pyc,, +fontTools/pens/__pycache__/cairoPen.cpython-312.pyc,, +fontTools/pens/__pycache__/cocoaPen.cpython-312.pyc,, +fontTools/pens/__pycache__/cu2quPen.cpython-312.pyc,, +fontTools/pens/__pycache__/explicitClosingLinePen.cpython-312.pyc,, +fontTools/pens/__pycache__/filterPen.cpython-312.pyc,, +fontTools/pens/__pycache__/freetypePen.cpython-312.pyc,, +fontTools/pens/__pycache__/hashPointPen.cpython-312.pyc,, +fontTools/pens/__pycache__/momentsPen.cpython-312.pyc,, +fontTools/pens/__pycache__/perimeterPen.cpython-312.pyc,, +fontTools/pens/__pycache__/pointInsidePen.cpython-312.pyc,, +fontTools/pens/__pycache__/pointPen.cpython-312.pyc,, +fontTools/pens/__pycache__/qtPen.cpython-312.pyc,, +fontTools/pens/__pycache__/qu2cuPen.cpython-312.pyc,, +fontTools/pens/__pycache__/quartzPen.cpython-312.pyc,, +fontTools/pens/__pycache__/recordingPen.cpython-312.pyc,, +fontTools/pens/__pycache__/reportLabPen.cpython-312.pyc,, +fontTools/pens/__pycache__/reverseContourPen.cpython-312.pyc,, +fontTools/pens/__pycache__/roundingPen.cpython-312.pyc,, +fontTools/pens/__pycache__/statisticsPen.cpython-312.pyc,, +fontTools/pens/__pycache__/svgPathPen.cpython-312.pyc,, +fontTools/pens/__pycache__/t2CharStringPen.cpython-312.pyc,, +fontTools/pens/__pycache__/teePen.cpython-312.pyc,, +fontTools/pens/__pycache__/transformPen.cpython-312.pyc,, +fontTools/pens/__pycache__/ttGlyphPen.cpython-312.pyc,, +fontTools/pens/__pycache__/wxPen.cpython-312.pyc,, +fontTools/pens/areaPen.py,sha256=Y1WkmqzcC4z_bpGAR0IZUKrtHFtxKUQBmr5-64_zCOk,1472 +fontTools/pens/basePen.py,sha256=eIGSKrKm6w4LLHuG6XJoQZ3eObtoKV5P6aF4gT4sk7U,17073 +fontTools/pens/boundsPen.py,sha256=wE3owOQA8DfhH-zBGC3lJvnVwp-oyIt0KZrEqXbmS9I,3129 +fontTools/pens/cairoPen.py,sha256=wuuOJ1qQDSt_K3zscM2nukRyHZTZMwMzzCXCirfq_qQ,592 +fontTools/pens/cocoaPen.py,sha256=IJRQcAxRuVOTQ90bB_Bgjnmz7px_ST5uLF9CW-Y0KPY,612 +fontTools/pens/cu2quPen.py,sha256=gMUwFUsm_-WzBlDjTMQiNnEuI2heomGeOJBX81zYXPo,13007 +fontTools/pens/explicitClosingLinePen.py,sha256=kKKtdZiwaf8Cj4_ytrIDdGB2GMpPPDXm5Nwbw5WDgwU,3219 +fontTools/pens/filterPen.py,sha256=REgspXaaSvF3XUwqe40abe3X_E7-WbBP13IqLUUBLCw,14703 +fontTools/pens/freetypePen.py,sha256=HD-gXJSbgImJdBc8sIBk0HWBdjv3WKFofs6PgCCsGOY,19908 +fontTools/pens/hashPointPen.py,sha256=gElrFyQoOQp3ZbpKHRWPwC61A9OgT2Js8crVUD8BQAY,3573 +fontTools/pens/momentsPen.c,sha256=YVgVJ8eTQIscGOc4zNztWjIdC4RGkiqLLpm8dDBamPY,567692 +fontTools/pens/momentsPen.cpython-312-x86_64-linux-gnu.so,sha256=_E9DVON4gfy3F8IDWnpbbpvJXk62xPyFfD6xmvxLjQM,986720 +fontTools/pens/momentsPen.py,sha256=kjLVXhGe55Abl__Yr1gob0bl0dHe7fPSwyr7TRJnbug,25658 +fontTools/pens/perimeterPen.py,sha256=lr6NzrIWxi4TXBJPbcJsKzqABWfQeil2Bgm9BgUD3N4,2153 +fontTools/pens/pointInsidePen.py,sha256=noEUvBQIeAheDMJwzvvfnEiKhmwbS1i0RQE9jik6Gl4,6355 +fontTools/pens/pointPen.py,sha256=oeE_uabVCNJ1Lpk5Hn3eBmafaao3QqKMjK6FAy0hKBo,24197 +fontTools/pens/qtPen.py,sha256=QRNLIry2rQl4E_7ct2tu10-qLHneQp0XV7FfaZ-tcL8,634 +fontTools/pens/qu2cuPen.py,sha256=pRST43-rUpzlOP83Z_Rr0IvIQBCx6RWI6nnNaitQcLk,3985 +fontTools/pens/quartzPen.py,sha256=EH482Kz_xsqYhVRovv6N_T1CXaSvOzUKPLxTaN956tU,1287 +fontTools/pens/recordingPen.py,sha256=VgFZ4NMhnZt1qSTzFEU0cma-gw3kBe47bfSxPYH73rs,12489 +fontTools/pens/reportLabPen.py,sha256=kpfMfOLXt2vOQ5smPsU82ft80FpCPWJzQLl7ENOH8Ew,2066 +fontTools/pens/reverseContourPen.py,sha256=oz64ZRhLAvT7DYMAwGKoLzZXQK8l81jRiYnTZkW6a-Y,4022 +fontTools/pens/roundingPen.py,sha256=vh_FjikRd82-S4I8glgGMGEuGrj5IkCjRT_wmZ8jfqY,4620 +fontTools/pens/statisticsPen.py,sha256=piWK6NjjWqk9MLROjeE2-4EsxVYMyNU7UQFGD_trE9g,9808 +fontTools/pens/svgPathPen.py,sha256=T3b6SZS9B9sVWMK9mSFDtjHeviQs_yOJOZKq5Sg5Zdg,8572 +fontTools/pens/t2CharStringPen.py,sha256=GgGklb5XsCer0w37ujgRLRXx-EuzdFsyCYuzCx4n-Qs,2931 +fontTools/pens/teePen.py,sha256=P1ARJOCMJ6MxK-PB1yZ-ips3CUfnadWYnQ_do6VIasQ,1290 +fontTools/pens/transformPen.py,sha256=s0kUyQdnemUwHvYr2SFboFmh4WY1S9OHBL8L4PJKRwE,4056 +fontTools/pens/ttGlyphPen.py,sha256=yLtB-E5pTQR59OKVYySttWBu1xC2vR8ezSaRhIMtVwg,11870 +fontTools/pens/wxPen.py,sha256=W9RRHlBWHp-CVC4Exvk3ytBmRaB4-LgJPP5Bv7o9BA0,680 +fontTools/qu2cu/__init__.py,sha256=Jfm1JljXbt91w4gyvZn6jzEmVnhRx50sh2fDongrOsE,618 +fontTools/qu2cu/__main__.py,sha256=9FWf6SIZaRaC8SiL0LhjAWC2yIdY9N_9wlRko8m1l2Q,93 +fontTools/qu2cu/__pycache__/__init__.cpython-312.pyc,, +fontTools/qu2cu/__pycache__/__main__.cpython-312.pyc,, +fontTools/qu2cu/__pycache__/benchmark.cpython-312.pyc,, +fontTools/qu2cu/__pycache__/cli.cpython-312.pyc,, +fontTools/qu2cu/__pycache__/qu2cu.cpython-312.pyc,, +fontTools/qu2cu/benchmark.py,sha256=GMcr_4r7L6K9SmJ13itt-_XKhnKqSVUDPlXUG6IZmmM,1400 +fontTools/qu2cu/cli.py,sha256=U2rooYnVVEalGRAWGFHk-Kp6Okys8wtzdaWLjw1bngY,3714 +fontTools/qu2cu/qu2cu.c,sha256=1F7GKT21uskZbU01eryh1sg-aTbB_KUXDZnL1JGOahY,692747 +fontTools/qu2cu/qu2cu.cpython-312-x86_64-linux-gnu.so,sha256=BlQPoIhu2uIqyuUaTTfSqDJby0p3yl993im96bJ_YsE,1234352 +fontTools/qu2cu/qu2cu.py,sha256=IYtpkwHdfKOXJr65Y_pJ9Lrt_MgJaISAKGMAs5ilFSM,12288 +fontTools/subset/__init__.py,sha256=R9VoZ2QWhqENHC5Ct1wyhLhEU-xo4mUXwCWl6EZGgwQ,143263 +fontTools/subset/__main__.py,sha256=bhtfP2SqP4k799pxtksFgnC-XGNQDr3LcO4lc8T5e5g,95 +fontTools/subset/__pycache__/__init__.cpython-312.pyc,, +fontTools/subset/__pycache__/__main__.cpython-312.pyc,, +fontTools/subset/__pycache__/cff.cpython-312.pyc,, +fontTools/subset/__pycache__/svg.cpython-312.pyc,, +fontTools/subset/__pycache__/util.cpython-312.pyc,, +fontTools/subset/cff.py,sha256=rqMRJOlX5FacV1LW8aDlVOglgEM87TkMA9bdsYenask,6145 +fontTools/subset/svg.py,sha256=sbuR0UHVB_9vYIb6dWCbZ5yCqushkpp5lg-f9NeatBM,9297 +fontTools/subset/util.py,sha256=9SXFYb5Ef9Z58uXmYPCQil8B2i3Q7aFB_1fFDFSppdU,754 +fontTools/svgLib/__init__.py,sha256=IGCLwSbU8jLhq6HI2vSdPQgNs6zDUi5774TgX5MCXPY,75 +fontTools/svgLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/svgLib/path/__init__.py,sha256=C82fh7xH6ZHsSFVnV848-xeDezpokx1EwTmayJCouFU,1996 +fontTools/svgLib/path/__pycache__/__init__.cpython-312.pyc,, +fontTools/svgLib/path/__pycache__/arc.cpython-312.pyc,, +fontTools/svgLib/path/__pycache__/parser.cpython-312.pyc,, +fontTools/svgLib/path/__pycache__/shapes.cpython-312.pyc,, +fontTools/svgLib/path/arc.py,sha256=-f5Ym6q4tDWQ76sMNSTUTWgL_7AfgXojvBhtBS7bWwQ,5812 +fontTools/svgLib/path/parser.py,sha256=8T6okMstvgM9ufb2zBcwSzsuuoYbqfnUjNYgb6kjznU,10788 +fontTools/svgLib/path/shapes.py,sha256=xvBUIckKyT9JLy7q_ZP50r6TjvZANyHdZP7wFDzErcI,5322 +fontTools/t1Lib/__init__.py,sha256=p42y70wEIbuX0IIxZG7-b_I-gHto1VLy0gLsDvxCfkw,20865 +fontTools/t1Lib/__pycache__/__init__.cpython-312.pyc,, +fontTools/tfmLib.py,sha256=UMbkM73JXRJVS9t2B-BJc13rSjImaWBuzCoehLwHFhs,14270 +fontTools/ttLib/__init__.py,sha256=1k7qp9z04gA3m6GvxDaINjqrKbzOkdTA_4RnqW_-LrA,661 +fontTools/ttLib/__main__.py,sha256=lHMPWsnzjKPuMFavf6i1gpk9KexiAk4qzgDd50Mbby0,4733 +fontTools/ttLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/ttLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/ttLib/__pycache__/macUtils.cpython-312.pyc,, +fontTools/ttLib/__pycache__/removeOverlaps.cpython-312.pyc,, +fontTools/ttLib/__pycache__/reorderGlyphs.cpython-312.pyc,, +fontTools/ttLib/__pycache__/scaleUpem.cpython-312.pyc,, +fontTools/ttLib/__pycache__/sfnt.cpython-312.pyc,, +fontTools/ttLib/__pycache__/standardGlyphOrder.cpython-312.pyc,, +fontTools/ttLib/__pycache__/ttCollection.cpython-312.pyc,, +fontTools/ttLib/__pycache__/ttFont.cpython-312.pyc,, +fontTools/ttLib/__pycache__/ttGlyphSet.cpython-312.pyc,, +fontTools/ttLib/__pycache__/ttVisitor.cpython-312.pyc,, +fontTools/ttLib/__pycache__/woff2.cpython-312.pyc,, +fontTools/ttLib/macUtils.py,sha256=lj3oeFpyjV7ko_JqnluneITmAtlc119J-vwTTg2s73A,1737 +fontTools/ttLib/removeOverlaps.py,sha256=ny3XhFw-RrriJ6n0sSh4SaVaKc3GhmCVZMw8g4oI0TY,12744 +fontTools/ttLib/reorderGlyphs.py,sha256=TbxLxqPTUGiKRX3ulGFCwVm2lEisFYlX6caONJr_4oY,10371 +fontTools/ttLib/scaleUpem.py,sha256=U_-NGkwfS9GRIackdEXjGYZ-wSomcUPXQahDneLeArI,14618 +fontTools/ttLib/sfnt.py,sha256=D5f7kvZX7rXDcmDqVBXqPPki4k27QgiBFYi8ziH2OiU,22981 +fontTools/ttLib/standardGlyphOrder.py,sha256=7AY_fVWdtwZ4iv5uWdyKAUcbEQiSDt1lN4sqx9xXwE0,5785 +fontTools/ttLib/tables/B_A_S_E_.py,sha256=H71A9pJ850mvjbrWHqy8iFI2Dxg7102YRtAkfdCooig,369 +fontTools/ttLib/tables/BitmapGlyphMetrics.py,sha256=9gcGPVzsxEYnVBO7YLWfeOuht9PaCl09GmbAqDYqKi0,1769 +fontTools/ttLib/tables/C_B_D_T_.py,sha256=5LNdc8FMir1kC5fvp5iHwWfeuE-RuqdxAArFXaqPjQ0,3646 +fontTools/ttLib/tables/C_B_L_C_.py,sha256=YXlwovoCHYx8THLQD9iBU_VGoaB9LFObEKtL6ddD320,520 +fontTools/ttLib/tables/C_F_F_.py,sha256=yg3mUtYBudgmpG7Bz475j_DNnCelsgrTsM8DH1uR4ek,1978 +fontTools/ttLib/tables/C_F_F__2.py,sha256=YoHfJQdF-ezx4OwRQ2Y2O7rRJEPjOkf3Hx5Y11Xq0AM,807 +fontTools/ttLib/tables/C_O_L_R_.py,sha256=SHwFVNVmoUQR2e87KuTSe-J9LfeegS4f2hEpee29_2o,5993 +fontTools/ttLib/tables/C_P_A_L_.py,sha256=odFjqM4GnjXyQYGEC-e0Gvqms1jQ5zHHG3SDg7y-BI0,11942 +fontTools/ttLib/tables/D_S_I_G_.py,sha256=AgQPM9Cdro1P-ehJjTfsC9mRTTtSc16At0nnpb1XOGI,5517 +fontTools/ttLib/tables/D__e_b_g.py,sha256=KDnfkNOUnm3F13wD_j3YNBOvYadZ40Gf_0170hFkJp0,1134 +fontTools/ttLib/tables/DefaultTable.py,sha256=cOtgkLWPY9qmOH2BSPt4c4IUSdANWTKx2rK1CTxQ4h0,1487 +fontTools/ttLib/tables/E_B_D_T_.py,sha256=uOpmt25gOJQeO1u1IGAyPWgVmh-4vSZqrQEHvOYwbwg,32534 +fontTools/ttLib/tables/E_B_L_C_.py,sha256=LfEVzBg_yWr9dhChzS0U2G-7wNOwzwB0LWoXIUYNKKM,30054 +fontTools/ttLib/tables/F_F_T_M_.py,sha256=_450vdbEH7Y-0_rOwb3Q0hg-Qq2W8C_sHljy7rZtqqM,1683 +fontTools/ttLib/tables/F__e_a_t.py,sha256=ct79Gf__5ALlqfSBn6wvw6fazb31Od71R6vIp6o9XF4,5483 +fontTools/ttLib/tables/G_D_E_F_.py,sha256=QXiILFCRnPNZcwpub6ojN5S9WP6y56LsXi25pUWLgp4,299 +fontTools/ttLib/tables/G_M_A_P_.py,sha256=fvIQumokOCLa8DFeq_xi069F9RROsXSVmDvWtxgyacQ,4720 +fontTools/ttLib/tables/G_P_K_G_.py,sha256=Xi4Hj2OxZ2IZgVyBQ-Qyiie0hPZjpXZkrao-E5EdTWM,4646 +fontTools/ttLib/tables/G_P_O_S_.py,sha256=UkP3mlnyvQg-jj6ZBOh6j-OieVg_goJQ31nlLvoLGSI,397 +fontTools/ttLib/tables/G_S_U_B_.py,sha256=cwFMKO-pgwsn1H8Q9Jb58Z6ZrBrCoN0sqJB0YunBfSk,294 +fontTools/ttLib/tables/G_V_A_R_.py,sha256=13oO2dD-L4yfkrBuR-KN2rc40wh5lLIlx_khwMz5GH4,94 +fontTools/ttLib/tables/G__l_a_t.py,sha256=Xh3IzFgYlvNjrAOn7Ja73DrWrQTJgJxmDFSUKS6yHdM,8645 +fontTools/ttLib/tables/G__l_o_c.py,sha256=5DsxGzaG7HyJVvLlKQeff1lXt-XPWaHNNaf-EYwsKh4,2685 +fontTools/ttLib/tables/H_V_A_R_.py,sha256=6kPLDUGT8EussA3y9gKr_mrgY5PNv7YaK1V0keMXD9w,313 +fontTools/ttLib/tables/J_S_T_F_.py,sha256=Q9TEf3OuyDIxZlmoz9a3c-mDMlJK6YBQ9KcYmiwFRbU,315 +fontTools/ttLib/tables/L_T_S_H_.py,sha256=Iu6syJFuhJj0_7Aan2NPlDuQDIq-AzLwsOQbXVTnlL0,2189 +fontTools/ttLib/tables/M_A_T_H_.py,sha256=-TVu9Nlcs-1shkElbIk-CWtUwXUMdycHFkjvPE8C_fs,342 +fontTools/ttLib/tables/M_E_T_A_.py,sha256=sA6ookcjchw8UYVEuS8QEXc62I9_Rms9cu_jKA6MkNI,11989 +fontTools/ttLib/tables/M_V_A_R_.py,sha256=67cEuiTw5y5W1Zk98L_S_SmJINIfy_mzWCkyHcujz94,308 +fontTools/ttLib/tables/O_S_2f_2.py,sha256=1Pq2Xu4oYOJePTHC_hTKg3RIfKely3j6T1u_lMTEpD8,28030 +fontTools/ttLib/tables/S_I_N_G_.py,sha256=CFDy8R2fDeYn7ocfrZr7Ui7U9D0h4G55CdPfY55g-Bk,3317 +fontTools/ttLib/tables/S_T_A_T_.py,sha256=y9NiWCtnlZtMjw4K9_SdA84Xa-dJk7G5eb2dSe6ciWc,498 +fontTools/ttLib/tables/S_V_G_.py,sha256=vT6QTW5ArtskVUxnPEH_ZxKz4DF4v1pKbylN6DG0R3o,7676 +fontTools/ttLib/tables/S__i_l_f.py,sha256=lPQV2RdhcJRgfDzHp_dkgSxVUUdkcAnY1Bz7V18Gt9U,34985 +fontTools/ttLib/tables/S__i_l_l.py,sha256=Vjtn7SI83vaLGIuQf2e-jhZSFOXb9vXB4jwqznjqnMc,3224 +fontTools/ttLib/tables/T_S_I_B_.py,sha256=3WhEtyNnjYumcowD0GpjubrgnS-RzouZxCxEe4yLDo8,341 +fontTools/ttLib/tables/T_S_I_C_.py,sha256=hAV9Hq_ALsWaducDpw1tDRREvFL7hx7onnUF0sXTelU,381 +fontTools/ttLib/tables/T_S_I_D_.py,sha256=TsdX-G2xxVQO9sSE1wE_xDRx-gor5YiXTHeUthMwCPY,341 +fontTools/ttLib/tables/T_S_I_J_.py,sha256=x8Tlvi6aTxoQcI12UL7muoWF1Q61iBDseAS1mRdOYrg,341 +fontTools/ttLib/tables/T_S_I_P_.py,sha256=-il2ucTBOghVBY7cmleHdLZc3W3CKh7-iPPT0A3KBzk,341 +fontTools/ttLib/tables/T_S_I_S_.py,sha256=tVBnl63vyZUIq93oM6dEjHCXvPn9vt5vvL3jG59b0Lg,341 +fontTools/ttLib/tables/T_S_I_V_.py,sha256=iUWxz2MSrtw7mzuQZj30QAJrCPnyJ4GincFfySFUNAg,855 +fontTools/ttLib/tables/T_S_I__0.py,sha256=O-2oI0eBgt4mP15-UwH0_0r7YWi3EEEhG-4etqDueGI,2505 +fontTools/ttLib/tables/T_S_I__1.py,sha256=nSUhni-fvYmeKXW4zLfP3FG_3LQU2QKPKS1_gKY5lYg,6971 +fontTools/ttLib/tables/T_S_I__2.py,sha256=q2rub-d77iWWiBM6awO0-TCl-Xq7kalPobHYC2QEOfc,496 +fontTools/ttLib/tables/T_S_I__3.py,sha256=0LcvvCzVZJzyz7i4zjIkUuYXEqXwOCs9WeCsgDFqKJ8,543 +fontTools/ttLib/tables/T_S_I__5.py,sha256=hhvJn6jiXs8kuBtun8krNUTXTljH-eKxaxXM1T-7SXM,1905 +fontTools/ttLib/tables/T_T_F_A_.py,sha256=LuT0w__AMtawnsBMobhEMW9gp2yk0mA5ZRzwF45c0UI,392 +fontTools/ttLib/tables/TupleVariation.py,sha256=4XTDTRPZWPg9_1K5SVgdNoxtgQvahtiO4LNO7fk1cK4,32235 +fontTools/ttLib/tables/V_A_R_C_.py,sha256=3jFX50J6X-Cc4dwwiztKKsDTRXVHTXlVdQH328UN1-k,289 +fontTools/ttLib/tables/V_D_M_X_.py,sha256=RbHl7vvO9pcjT_kKvcCmcByQj39n4PmVeq55wD5C14g,10437 +fontTools/ttLib/tables/V_O_R_G_.py,sha256=Cn3OxjVtcO-Uvp61P5c2336V9iEbuGr6vWAXnSIaihk,5965 +fontTools/ttLib/tables/V_V_A_R_.py,sha256=Cstw6tc_U4-EmTriRItBSpvTJODAjMFQjfyTaxLzsbI,319 +fontTools/ttLib/tables/__init__.py,sha256=eQPcuHCfRuGtt6nOa0KwV6vtUNKHnwuQyA7xSN8SPoc,2651 +fontTools/ttLib/tables/__pycache__/B_A_S_E_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/BitmapGlyphMetrics.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_B_D_T_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_B_L_C_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_F_F_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_F_F__2.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_O_L_R_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/C_P_A_L_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/D_S_I_G_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/D__e_b_g.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/DefaultTable.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/E_B_D_T_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/E_B_L_C_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/F_F_T_M_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/F__e_a_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_D_E_F_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_M_A_P_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_P_K_G_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_P_O_S_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_S_U_B_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G_V_A_R_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G__l_a_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/G__l_o_c.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/H_V_A_R_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/J_S_T_F_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/L_T_S_H_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/M_A_T_H_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/M_E_T_A_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/M_V_A_R_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/O_S_2f_2.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/S_I_N_G_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/S_T_A_T_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/S_V_G_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/S__i_l_f.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/S__i_l_l.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_B_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_C_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_D_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_J_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_P_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_S_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I_V_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I__0.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I__1.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I__2.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I__3.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_S_I__5.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/T_T_F_A_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/TupleVariation.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/V_A_R_C_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/V_D_M_X_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/V_O_R_G_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/V_V_A_R_.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/__init__.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_a_n_k_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_a_v_a_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_b_s_l_n.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_c_i_d_g.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_c_m_a_p.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_c_v_a_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_c_v_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_f_e_a_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_f_p_g_m.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_f_v_a_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_g_a_s_p.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_g_c_i_d.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_g_l_y_f.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_g_v_a_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_h_d_m_x.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_h_e_a_d.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_h_h_e_a.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_h_m_t_x.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_k_e_r_n.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_l_c_a_r.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_l_o_c_a.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_l_t_a_g.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_m_a_x_p.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_m_e_t_a.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_m_o_r_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_m_o_r_x.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_n_a_m_e.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_o_p_b_d.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_p_o_s_t.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_p_r_e_p.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_p_r_o_p.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_s_b_i_x.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_t_r_a_k.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_v_h_e_a.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/_v_m_t_x.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/asciiTable.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/grUtils.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/otBase.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/otConverters.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/otData.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/otTables.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/otTraverse.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/sbixGlyph.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/sbixStrike.cpython-312.pyc,, +fontTools/ttLib/tables/__pycache__/ttProgram.cpython-312.pyc,, +fontTools/ttLib/tables/_a_n_k_r.py,sha256=MpAzIifmIi_3gx2oP6PC3R2lu36Ewsr2-W1rXjsz2Ug,483 +fontTools/ttLib/tables/_a_v_a_r.py,sha256=_qplWEwFos0VwRTdZa7LhfLIWw-CF1G76NbiRM_Rgvs,7374 +fontTools/ttLib/tables/_b_s_l_n.py,sha256=_848o7SQqztzBDfHYei-80u9ltxIHVBzXu1dYHLV57M,465 +fontTools/ttLib/tables/_c_i_d_g.py,sha256=yt8rVIadpJSDUCoVH4dZetNiy0Azm5ESAxHjB2BX_eA,913 +fontTools/ttLib/tables/_c_m_a_p.py,sha256=r8-bB_E0EQh5h4TGX5nTnDnwTUtXuRB3iuqEDoN_IOM,62202 +fontTools/ttLib/tables/_c_v_a_r.py,sha256=35ayk2kX1pcLGwyx0y4I1l-r7LHgdKv0ulVx8oBPteI,3527 +fontTools/ttLib/tables/_c_v_t.py,sha256=1_RhEcTmhWQWQp7Hsj8UsByKmXCIppZyIbIArGywEEM,1618 +fontTools/ttLib/tables/_f_e_a_t.py,sha256=Fi1XnjhkCG0tp43AcvpIaivD-YRFpufo6feGIrenQDo,469 +fontTools/ttLib/tables/_f_p_g_m.py,sha256=uZHZzqL6OdLn_Hxskv-xf3XuE4fyaSv_jbALEjwXYug,1633 +fontTools/ttLib/tables/_f_v_a_r.py,sha256=rV33H2BgHUl3Wuydsou1G-Hi4uASBppWaLj3FMmiLjs,8837 +fontTools/ttLib/tables/_g_a_s_p.py,sha256=YvhAVDvdssN2fjPMTfSrO4WBCfTuh9T2cU5zquDVnSw,2203 +fontTools/ttLib/tables/_g_c_i_d.py,sha256=AJ4uV7PTHbnsw4Tfw8c2Ezh0VMox3oAH0qhhq7y8hdM,362 +fontTools/ttLib/tables/_g_l_y_f.py,sha256=dDV65llsEDI9fKcVKC5TOiaXpXSyMHNEytuYOGt7adM,85584 +fontTools/ttLib/tables/_g_v_a_r.py,sha256=E9WCKjeITUfd5hcJLQ0rjQFBtZdxw1eswFlWp1U6bD4,12196 +fontTools/ttLib/tables/_h_d_m_x.py,sha256=wMrO4D04QNT8u30p8AV-aG3bndXCq4wlPNvtbd8ip7c,4252 +fontTools/ttLib/tables/_h_e_a_d.py,sha256=yY2GTFq6Mn6nN8EegbMVJRMUWIqDYFln3FhTk3ziw6s,4926 +fontTools/ttLib/tables/_h_h_e_a.py,sha256=X4t1aF1MZMuz3phCVSFwKcNTeoZdx-042wFtHc-nK9w,4767 +fontTools/ttLib/tables/_h_m_t_x.py,sha256=rbxr3cy9-9Jm0HCGIWQiX6fGH5iu6yojp9kfgWrW2Ks,6192 +fontTools/ttLib/tables/_k_e_r_n.py,sha256=DQNLmD_HEdDKPfp4tamOd9W3T5a1lXFM5tDaWrKl164,10794 +fontTools/ttLib/tables/_l_c_a_r.py,sha256=8W6xFOj-sm003MCXX4bIHxs9ntfVvT0FXYllPxa3z4I,390 +fontTools/ttLib/tables/_l_o_c_a.py,sha256=yxiwLKXLZjNju5XYmLb6EhNLec1d7ezEDDe1dszceHo,2180 +fontTools/ttLib/tables/_l_t_a_g.py,sha256=9YpApjI-rZ4e3HeT8Pj-osiHl3uALD9JXg5O7pqk9L0,2552 +fontTools/ttLib/tables/_m_a_x_p.py,sha256=cIDIZWse9czwwsnlxIh3qwgwaXbt7PQAjXKAcmMDspY,5264 +fontTools/ttLib/tables/_m_e_t_a.py,sha256=A0CZPEAVxYrpytjXUGQJCTddwG8KrvUVbtBe3A1MqgI,3913 +fontTools/ttLib/tables/_m_o_r_t.py,sha256=u35tYqn3cjzKxeCF0FUFeLtaf36mjDDSN08uuk0Kme8,487 +fontTools/ttLib/tables/_m_o_r_x.py,sha256=OwamVpIO7REDnFr95HuFPoY_0U6i9zQPb11K1sFTvDY,548 +fontTools/ttLib/tables/_n_a_m_e.py,sha256=86_0fUeA5_c-GY5ZnkqUI0jyWwMh1mn6yVOf6KKqIlU,41266 +fontTools/ttLib/tables/_o_p_b_d.py,sha256=TNZv_2YTrj4dGzd6wA9Jb-KGZ99un177s5p3LlfxQ74,448 +fontTools/ttLib/tables/_p_o_s_t.py,sha256=9siVXSisWGdTfj_mC1E9dUDz9Jdm1i3QzI-l3i3VWME,11671 +fontTools/ttLib/tables/_p_r_e_p.py,sha256=CcKr4HrswkupLmbJdrJLTM-z9XgLefQyv8467j9V0zs,427 +fontTools/ttLib/tables/_p_r_o_p.py,sha256=Eg8x5qWyXDzPezMafFu0s0qyPDHj-sPsFxGtE6h29qo,427 +fontTools/ttLib/tables/_s_b_i_x.py,sha256=tkkKbNKNYkUhZJuN0kl7q37x5KK5OovB06y28obPV6A,4865 +fontTools/ttLib/tables/_t_r_a_k.py,sha256=rrrPZLELFYA5F8PERoafIS9cb_d_i6xtpAzHEbsFHSw,11379 +fontTools/ttLib/tables/_v_h_e_a.py,sha256=FuULIBl4OQyUeLPOFEY8buB0pAnQhGa1-5a6kN9i5Sc,4459 +fontTools/ttLib/tables/_v_m_t_x.py,sha256=AUuxtyQvMWrTBNbOIaL6uKcB_DNpNb0YX28JIuTHw_Y,500 +fontTools/ttLib/tables/asciiTable.py,sha256=4c69jsAirUnDEpylf9CYBoCKTzwbmfbtUAOrtPnpHjY,637 +fontTools/ttLib/tables/grUtils.py,sha256=hcOJ5oJPOd2uJWnWA7qwR7AfL37YZ5zUT7g8o5BBV80,2270 +fontTools/ttLib/tables/otBase.py,sha256=k1Mt5sLd2EL6ufX3e0ZaBDgTvSLCUKOW6qy5UDeKwxI,52986 +fontTools/ttLib/tables/otConverters.py,sha256=ihE_WMSKAKSaBbMvnFYDj2eMxf7PvRMMa8zGwfoYuYc,74202 +fontTools/ttLib/tables/otData.py,sha256=-XXRwdVfP-Wz7oBjMPpku0A0QH9lw_fFGNzZlt9N0mo,197262 +fontTools/ttLib/tables/otTables.py,sha256=2U04ot_2ITlBZx2QtpnIOtBGftPFs9ZX2FWfz4vz1G0,96987 +fontTools/ttLib/tables/otTraverse.py,sha256=HznEVAlVf_8eyqjsO2edgELtMlXnjnUqccK3PytvVUE,5518 +fontTools/ttLib/tables/sbixGlyph.py,sha256=tjEUPVRfx6gr5yme8UytGTtVrimKN5qmbzT1GZPjXiM,5796 +fontTools/ttLib/tables/sbixStrike.py,sha256=dL8O9K8R4S6RVQDP-PVjIPBrvbqbE9zwra0uRL0nLq0,6651 +fontTools/ttLib/tables/table_API_readme.txt,sha256=eZlRTLUkLzc_9Ot3pdfhyMb3ahU0_Iipx0vSbzOVGy8,2748 +fontTools/ttLib/tables/ttProgram.py,sha256=tgtxgd-EnOq-2PUlYEihp-6NHu_7HnE5rxeSAtmXOtU,35888 +fontTools/ttLib/ttCollection.py,sha256=aRph2MkBK3kd9-JCLqhJ1EN9pffN_lVX6WWmOTTewc8,3963 +fontTools/ttLib/ttFont.py,sha256=TKcK7qSvxHum0iXpX0zS3YiDjt6yW9wlK_YsoP2torE,63707 +fontTools/ttLib/ttGlyphSet.py,sha256=cUBhMGa5hszeVqOm2KpOdeJh-LsiqE7RNdyIUPZ2vO8,17476 +fontTools/ttLib/ttVisitor.py,sha256=_tah4C42Tv6Pm9QeLNQwwVCxqI4VNEAqYCbmThp6cvY,1025 +fontTools/ttLib/woff2.py,sha256=6LPISeBQ1dubzKjWrUcYm_vgETC46BTLY4XkG52qvSA,60921 +fontTools/ttx.py,sha256=FxuGubujWCGJWSTrJEjoNH--25fVIPy-ZRtYy9H6iTk,17277 +fontTools/ufoLib/__init__.py,sha256=nKG8gu6NEvqGJoZ781IARoQ7ii4LoWfMMvX3Yf5TsVw,98981 +fontTools/ufoLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/converters.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/errors.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/etree.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/filenames.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/glifLib.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/kerning.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/plistlib.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/pointPen.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/utils.cpython-312.pyc,, +fontTools/ufoLib/__pycache__/validators.cpython-312.pyc,, +fontTools/ufoLib/converters.py,sha256=YnBKr8kmyjwLcq8LdD46ubGOgyL9Pxt9avlvZn9anKI,13444 +fontTools/ufoLib/errors.py,sha256=9f8l5NaFAj3BZPa6Bbqt06FL4afffLuMzy4nPf-eOlE,845 +fontTools/ufoLib/etree.py,sha256=T3sjLTgjMAq6VyYRicWPaMIVBJ2YSuwZxV6Vc5yZtQI,231 +fontTools/ufoLib/filenames.py,sha256=hoyUhzzQMDaeckT7UdreISANq4-gLR2jGyk5yAyYtOA,10654 +fontTools/ufoLib/glifLib.py,sha256=Y-xzf4qbTIOl3-dVLXvu3iFCIDtAEu_klId2_UNngWs,77170 +fontTools/ufoLib/kerning.py,sha256=o1BeJDVZ_CZZPzmOPwRKTqglYmhA_JZPjwq2JLgdQIk,4836 +fontTools/ufoLib/plistlib.py,sha256=jzMGOGvHO6XvS-IO8hS04ur7r8-v2dnVq-vKMoJZvqQ,1510 +fontTools/ufoLib/pointPen.py,sha256=CuREcm3IYteZNBDAd_ZRAV4XqBsy0s07jdWc4en9r-8,244 +fontTools/ufoLib/utils.py,sha256=nZoJJqHXQSL-LXYE58_WHA97XlbTkEbYkdH3GL32SmQ,3192 +fontTools/ufoLib/validators.py,sha256=MWBqcLThGyYpst61QothA_BSlc6jGVhPvFiay-pobCY,32387 +fontTools/unicode.py,sha256=ZZ7OMmWvIyV1IL1k6ioTzaRAh3tUvm6gvK7QgFbOIHY,1237 +fontTools/unicodedata/Blocks.py,sha256=frShN07WqD1uLt-V0TNlUZPnn0-cKc0Nuhtd2uNDoSw,33139 +fontTools/unicodedata/Mirrored.py,sha256=kdhwCWOWaArmfNkDah0Thv-67M9wWz45R5IMPhqyzFM,9242 +fontTools/unicodedata/OTTags.py,sha256=wOPpbMsNcp_gdvPFeITtgVMnTN8TJSNAsVEdu_nuPXE,1196 +fontTools/unicodedata/ScriptExtensions.py,sha256=udL53Wdd5n-ZFEVHeIC0LpKKBqPHD2WqRiBbzXs2U7U,28647 +fontTools/unicodedata/Scripts.py,sha256=wBD6SElfM3w-u6JcgyNsyofXPvqW1e4YltgeHE4XaOc,131148 +fontTools/unicodedata/__init__.py,sha256=4eR0Luk4QJHC9YDG0_wi6UrIDawED36e6Xbf3S5hDPg,9085 +fontTools/unicodedata/__pycache__/Blocks.cpython-312.pyc,, +fontTools/unicodedata/__pycache__/Mirrored.cpython-312.pyc,, +fontTools/unicodedata/__pycache__/OTTags.cpython-312.pyc,, +fontTools/unicodedata/__pycache__/ScriptExtensions.cpython-312.pyc,, +fontTools/unicodedata/__pycache__/Scripts.cpython-312.pyc,, +fontTools/unicodedata/__pycache__/__init__.cpython-312.pyc,, +fontTools/varLib/__init__.py,sha256=nzk6FSRozAMz9CDVqhMVYhh4Ho_hjIKrYtkMIJLn0y0,57467 +fontTools/varLib/__main__.py,sha256=wbdYC5bPjWCxA0I4SKcLO88gl-UMtsYS8MxdW9ySTkY,95 +fontTools/varLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/varLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/varLib/__pycache__/avarPlanner.cpython-312.pyc,, +fontTools/varLib/__pycache__/builder.cpython-312.pyc,, +fontTools/varLib/__pycache__/cff.cpython-312.pyc,, +fontTools/varLib/__pycache__/errors.cpython-312.pyc,, +fontTools/varLib/__pycache__/featureVars.cpython-312.pyc,, +fontTools/varLib/__pycache__/hvar.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolatable.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolatableHelpers.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolatablePlot.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolatableTestContourOrder.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolatableTestStartingPoint.cpython-312.pyc,, +fontTools/varLib/__pycache__/interpolate_layout.cpython-312.pyc,, +fontTools/varLib/__pycache__/iup.cpython-312.pyc,, +fontTools/varLib/__pycache__/merger.cpython-312.pyc,, +fontTools/varLib/__pycache__/models.cpython-312.pyc,, +fontTools/varLib/__pycache__/multiVarStore.cpython-312.pyc,, +fontTools/varLib/__pycache__/mutator.cpython-312.pyc,, +fontTools/varLib/__pycache__/mvar.cpython-312.pyc,, +fontTools/varLib/__pycache__/plot.cpython-312.pyc,, +fontTools/varLib/__pycache__/stat.cpython-312.pyc,, +fontTools/varLib/__pycache__/varStore.cpython-312.pyc,, +fontTools/varLib/avar/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +fontTools/varLib/avar/__main__.py,sha256=ew1fJpg81GpYbdkrp4I7ntWdbhkDeg5fNn1SrpF254k,1770 +fontTools/varLib/avar/__pycache__/__init__.cpython-312.pyc,, +fontTools/varLib/avar/__pycache__/__main__.cpython-312.pyc,, +fontTools/varLib/avar/__pycache__/build.cpython-312.pyc,, +fontTools/varLib/avar/__pycache__/map.cpython-312.pyc,, +fontTools/varLib/avar/__pycache__/plan.cpython-312.pyc,, +fontTools/varLib/avar/__pycache__/unbuild.cpython-312.pyc,, +fontTools/varLib/avar/build.py,sha256=YSMsRMjGkNW2zhKn29HgovNNG-jgiFSsK7zYH4un2Ws,2087 +fontTools/varLib/avar/map.py,sha256=UBeElT40SHXvJrQ7SflH5MNYKTTom7RLHwq2varb_oE,2867 +fontTools/varLib/avar/plan.py,sha256=qDKZjQq6OfeXBr3T3DZnrmauyMm5ek0rm5YVosAO_UA,27354 +fontTools/varLib/avar/unbuild.py,sha256=h3qb0JBN7SngijPrwjTjGEjueeENlhPYY9FyN02FLEk,10467 +fontTools/varLib/avarPlanner.py,sha256=CabG5xB57FMdwmN1acpVMvyOucVPxCdRdJSK2gu7gb4,109 +fontTools/varLib/builder.py,sha256=mSKOCcnnw-WzmZs15FayoqCDh77Ts7o9Tre9psh8CUc,6609 +fontTools/varLib/cff.py,sha256=EVgaQcoROIrYQsRuftnxFuGGldEPYbrIh5yBckylJC4,22901 +fontTools/varLib/errors.py,sha256=dMo8eGj76I7H4hrBEiNbYrGs2J1K1SwdsUyTHpkVOrQ,6934 +fontTools/varLib/featureVars.py,sha256=kp4gPjKyyGRu7yBWxgd1N4OkKnU9V45QwiWSHs6OWd0,26180 +fontTools/varLib/hvar.py,sha256=1IvL5BneTkg8jJYicH0TSQViB6D0vBEesLdlfqoLBX4,3695 +fontTools/varLib/instancer/__init__.py,sha256=KVXejhVDoIx8cqQuJKOB_N2q8uE2stsMQ68tqp268PI,75502 +fontTools/varLib/instancer/__main__.py,sha256=zfULwcP01FhplS1IlcMgNQnLxk5RVfmOuinWjqeid-g,104 +fontTools/varLib/instancer/__pycache__/__init__.cpython-312.pyc,, +fontTools/varLib/instancer/__pycache__/__main__.cpython-312.pyc,, +fontTools/varLib/instancer/__pycache__/featureVars.cpython-312.pyc,, +fontTools/varLib/instancer/__pycache__/names.cpython-312.pyc,, +fontTools/varLib/instancer/__pycache__/solver.cpython-312.pyc,, +fontTools/varLib/instancer/featureVars.py,sha256=oPqSlnHLMDTtOsmQMi6gkzLox7ymCrqlRAkvC_EJ4bc,7110 +fontTools/varLib/instancer/names.py,sha256=IPRqel_M8zVU0jl30WsfgufxUm9PBBQDQCY3VHapeHc,14950 +fontTools/varLib/instancer/solver.py,sha256=uMePwX0BVT5F94kUvDglsI4_F0nEH67F7RFuJ6tQwQ0,11002 +fontTools/varLib/interpolatable.py,sha256=Bhlq_LhEZ-sXfLNY8aFEChFrsKuT2kzmnuMfG5qi0v4,45221 +fontTools/varLib/interpolatableHelpers.py,sha256=cTFgTqDjggSCqNfTM77__5b9Sja_g7xWWMiB-pXDx84,11672 +fontTools/varLib/interpolatablePlot.py,sha256=w393P6mGLRhYkIjSxMww3qyoYxAUZzCXlmPBbI_84C0,44375 +fontTools/varLib/interpolatableTestContourOrder.py,sha256=mHJ9Ry7Rm7H3zHDwEUQEtEIDseiUzOxjg4MveW_FSiU,3021 +fontTools/varLib/interpolatableTestStartingPoint.py,sha256=K6OYKBspim6BXc91pfLTbGLyi5XZukfMuBc6hRpENG8,4296 +fontTools/varLib/interpolate_layout.py,sha256=22VjGZuV2YiAe2MpdTf0xPVz1x2G84bcOL0vOeBpGQM,3689 +fontTools/varLib/iup.c,sha256=64aSGfKmZAN_uGJrYxvtmzstksJYD3ak2iv8iyjW3bM,834812 +fontTools/varLib/iup.cpython-312-x86_64-linux-gnu.so,sha256=avtYY7S7fCAM2bP8jjTATzKZ2dGVBRc98J22k_fxPxU,1650280 +fontTools/varLib/iup.py,sha256=mKq_GRWuUg4yTmw2V32nu0v2r-SzzN7xS7rIbV0mYuc,14984 +fontTools/varLib/merger.py,sha256=E59oli4AwqWZ-FgnuStMSBvsB-FHe-55esXTYUqGeJ8,60802 +fontTools/varLib/models.py,sha256=mkwlucmKDkyhPZWEpab7ofcuA19hGlqjHtY9MwmYtOI,22697 +fontTools/varLib/multiVarStore.py,sha256=eQEuWNY01YF5zDpy1UwNtvOYyD6c0FLxpH-QFpX1i78,8305 +fontTools/varLib/mutator.py,sha256=kzXiLFxRLgU2pcHzOzh9u0n0KkO3DuBk06xZ_RPhWz8,19804 +fontTools/varLib/mvar.py,sha256=LTV77vH_3Ecg_qKBO5xQzjLOlJir_ppEr7mPVZRgad8,2449 +fontTools/varLib/plot.py,sha256=NoSZkJ5ndxNcDvJIvd5pQ9_jX6X1oM1K2G_tR4sdPVs,7494 +fontTools/varLib/stat.py,sha256=XuNKKZxGlBrl4OGFDAwVXhpBwJi23U3BdHmNTKoJnvE,4811 +fontTools/varLib/varStore.py,sha256=2QA9SDI6jQyQ_zq82OOwa3FBkfl-ksaSo1KGmVFpa9Q,24069 +fontTools/voltLib/__init__.py,sha256=ZZ1AsTx1VlDn40Kupce-fM3meOWugy3RZraBW9LG-9M,151 +fontTools/voltLib/__main__.py,sha256=uVtABLzMeHtvKL8zetf4rpC4aB8BkYr5QLSegNjZZZI,5928 +fontTools/voltLib/__pycache__/__init__.cpython-312.pyc,, +fontTools/voltLib/__pycache__/__main__.cpython-312.pyc,, +fontTools/voltLib/__pycache__/ast.cpython-312.pyc,, +fontTools/voltLib/__pycache__/error.cpython-312.pyc,, +fontTools/voltLib/__pycache__/lexer.cpython-312.pyc,, +fontTools/voltLib/__pycache__/parser.cpython-312.pyc,, +fontTools/voltLib/__pycache__/voltToFea.cpython-312.pyc,, +fontTools/voltLib/ast.py,sha256=arA9W3Gqo6OqljwNNKnMAojz-C5LStbC5SgjJh7buKk,13300 +fontTools/voltLib/error.py,sha256=phcQOQj-xOspCXu9hBJQRhSOBDzxHRgZd3fWQOFNJzw,395 +fontTools/voltLib/lexer.py,sha256=OvuETOSvlS6v7iCVeJ3IdH2Cg71n3OJoEyiB3-h6vhE,3368 +fontTools/voltLib/parser.py,sha256=rkw2IHBZPsrhGVC7Kw7V501m0u52kh1JSM5HXp-xchM,25396 +fontTools/voltLib/voltToFea.py,sha256=Z2yvnaZLQXzPLT86Uta0zRsXIYgj6NnvZtSWt5xmw2s,36549 +fonttools-4.61.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +fonttools-4.61.1.dist-info/METADATA,sha256=sF0h1Ul5KoQam4cNjcDOXJ1D2Z1nYh_xN5sYvop8OUA,114185 +fonttools-4.61.1.dist-info/RECORD,, +fonttools-4.61.1.dist-info/WHEEL,sha256=kaf3GlniTxElItJKanA5i_DdhhelZhGNqtfvwRkHPG0,224 +fonttools-4.61.1.dist-info/entry_points.txt,sha256=8kVHddxfFWA44FSD4mBpmC-4uCynQnkoz_9aNJb227Y,147 +fonttools-4.61.1.dist-info/licenses/LICENSE,sha256=Z4cgj4P2Wcy8IiOy_elS_6b36KymLxqKK_W8UbsbI4M,1072 +fonttools-4.61.1.dist-info/licenses/LICENSE.external,sha256=lKg6ruBymg8wLTSsxKzsvZ1YNm8mJCkHX-VX5KVLLmk,20022 +fonttools-4.61.1.dist-info/top_level.txt,sha256=rRgRylrXzekqWOsrhygzib12pQ7WILf7UGjqEwkIFDM,10 diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/WHEEL b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4afdbc4f127375c632d0a0df6c686b5925cd3b32 --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/WHEEL @@ -0,0 +1,8 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: false +Tag: cp312-cp312-manylinux_2_5_x86_64 +Tag: cp312-cp312-manylinux1_x86_64 +Tag: cp312-cp312-manylinux_2_17_x86_64 +Tag: cp312-cp312-manylinux2014_x86_64 + diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/entry_points.txt b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..87ae781f169a63f0cf672a9050474035bfa5add4 --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/entry_points.txt @@ -0,0 +1,5 @@ +[console_scripts] +fonttools = fontTools.__main__:main +pyftmerge = fontTools.merge:main +pyftsubset = fontTools.subset:main +ttx = fontTools.ttx:main diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..cc633905d333c4b42c1a0c8b34e9f734adeb6e1e --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2017 Just van Rossum + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE.external b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE.external new file mode 100644 index 0000000000000000000000000000000000000000..5c45052f870940119ee6bc2b51ab3a7bc43d01d7 --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/licenses/LICENSE.external @@ -0,0 +1,388 @@ +FontTools includes the following font projects for testing purposes, which are +under SIL Open Font License, Version 1.1: + +Lobster + Copyright (c) 2010, Pablo Impallari (www.impallari.com|impallari@gmail.com), + with Reserved Font Name Lobster. + This Font Software is licensed under the SIL Open Font License, Version 1.1. + +Noto Fonts + This Font Software is licensed under the SIL Open Font License, Version 1.1. + +XITS font project + Copyright (c) 2001-2010 by the STI Pub Companies, consisting of the American + Institute of Physics, the American Chemical Society, the American + Mathematical Society, the American Physical Society, Elsevier, Inc., and The + Institute of Electrical and Electronic Engineers, Inc. (www.stixfonts.org), + with Reserved Font Name STIX Fonts, STIX Fonts (TM) is a trademark of The + Institute of Electrical and Electronics Engineers, Inc. + + Portions copyright (c) 1998-2003 by MicroPress, Inc. + (www.micropress-inc.com), with Reserved Font Name TM Math. To obtain + additional mathematical fonts, please contact MicroPress, Inc., 68-30 Harrow + Street, Forest Hills, NY 11375, USA, Phone: (718) 575-1816. + + Portions copyright (c) 1990 by Elsevier, Inc. + + This Font Software is licensed under the SIL Open Font License, Version 1.1. + +Iosevka + Copyright (c) 2015-2020 Belleve Invis (belleve@typeof.net). + This Font Software is licensed under the SIL Open Font License, Version 1.1. + +This license is copied below, and is also available with a FAQ at: +http://scripts.sil.org/OFL + +----------------------------------------------------------- +SIL OPEN FONT LICENSE Version 1.1 - 26 February 2007 +----------------------------------------------------------- + +PREAMBLE +The goals of the Open Font License (OFL) are to stimulate worldwide +development of collaborative font projects, to support the font +creation efforts of academic and linguistic communities, and to +provide a free and open framework in which fonts may be shared and +improved in partnership with others. + +The OFL allows the licensed fonts to be used, studied, modified and +redistributed freely as long as they are not sold by themselves. The +fonts, including any derivative works, can be bundled, embedded, +redistributed and/or sold with any software provided that any reserved +names are not used by derivative works. The fonts and derivatives, +however, cannot be released under any other type of license. The +requirement for fonts to remain under this license does not apply to +any document created using the fonts or their derivatives. + +DEFINITIONS +"Font Software" refers to the set of files released by the Copyright +Holder(s) under this license and clearly marked as such. This may +include source files, build scripts and documentation. + +"Reserved Font Name" refers to any names specified as such after the +copyright statement(s). + +"Original Version" refers to the collection of Font Software +components as distributed by the Copyright Holder(s). + +"Modified Version" refers to any derivative made by adding to, +deleting, or substituting -- in part or in whole -- any of the +components of the Original Version, by changing formats or by porting +the Font Software to a new environment. + +"Author" refers to any designer, engineer, programmer, technical +writer or other person who contributed to the Font Software. + +PERMISSION & CONDITIONS +Permission is hereby granted, free of charge, to any person obtaining +a copy of the Font Software, to use, study, copy, merge, embed, +modify, redistribute, and sell modified and unmodified copies of the +Font Software, subject to the following conditions: + +1) Neither the Font Software nor any of its individual components, in +Original or Modified Versions, may be sold by itself. + +2) Original or Modified Versions of the Font Software may be bundled, +redistributed and/or sold with any software, provided that each copy +contains the above copyright notice and this license. These can be +included either as stand-alone text files, human-readable headers or +in the appropriate machine-readable metadata fields within text or +binary files as long as those fields can be easily viewed by the user. + +3) No Modified Version of the Font Software may use the Reserved Font +Name(s) unless explicit written permission is granted by the +corresponding Copyright Holder. This restriction only applies to the +primary font name as presented to the users. + +4) The name(s) of the Copyright Holder(s) or the Author(s) of the Font +Software shall not be used to promote, endorse or advertise any +Modified Version, except to acknowledge the contribution(s) of the +Copyright Holder(s) and the Author(s) or with their explicit written +permission. + +5) The Font Software, modified or unmodified, in part or in whole, +must be distributed entirely under this license, and must not be +distributed under any other license. The requirement for fonts to +remain under this license does not apply to any document created using +the Font Software. + +TERMINATION +This license becomes null and void if any of the above conditions are +not met. + +DISCLAIMER +THE FONT SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO ANY WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT +OF COPYRIGHT, PATENT, TRADEMARK, OR OTHER RIGHT. IN NO EVENT SHALL THE +COPYRIGHT HOLDER BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, +INCLUDING ANY GENERAL, SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL +DAMAGES, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF THE USE OR INABILITY TO USE THE FONT SOFTWARE OR FROM +OTHER DEALINGS IN THE FONT SOFTWARE. + +===== + +FontTools includes Adobe AGL & AGLFN, which is under 3-clauses BSD license: + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +Redistributions of source code must retain the above copyright notice, +this list of conditions and the following disclaimer. + +Redistributions in binary form must reproduce the above copyright +notice, this list of conditions and the following disclaimer in the +documentation and/or other materials provided with the distribution. + +Neither the name of Adobe Systems Incorporated nor the names of its +contributors may be used to endorse or promote products derived from +this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +===== + +FontTools includes cu2qu, which is Copyright 2016 Google Inc. All Rights Reserved. +Licensed under the Apache License, Version 2.0, a copy of which is reproduced below: + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. + +===== + +FontTools includes code in `fontTools.misc.filesystem` which is derived from: + +PyFilesystem2 (i.e. the `fs` package) by Will McGugan +Licensed under the MIT License +https://github.com/PyFilesystem/pyfilesystem2 + +Copyright (c) 2017-2021 The PyFilesystem2 contributors +Copyright (c) 2016-2019 Will McGugan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/top_level.txt b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..9af65ba39d292309497df4accdc44bd6f8143d10 --- /dev/null +++ b/lib/python3.12/site-packages/fonttools-4.61.1.dist-info/top_level.txt @@ -0,0 +1 @@ +fontTools diff --git a/lib/python3.12/site-packages/gitdb/__init__.py b/lib/python3.12/site-packages/gitdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1fb7df8939f16d8378bad6f26e124c071dc7d1a5 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/__init__.py @@ -0,0 +1,16 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Initialize the object database module""" + +__author__ = "Sebastian Thiel" +__contact__ = "byronimo@gmail.com" +__homepage__ = "https://github.com/gitpython-developers/gitdb" +version_info = (4, 0, 12) +__version__ = '.'.join(str(i) for i in version_info) + +# default imports +from gitdb.base import * +from gitdb.db import * +from gitdb.stream import * diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2906db2b7bbf2adaf80a74919be5fb6d9768c6b1 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6819610d02fb9e9af76d64ed93fc04ff220d2f76 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/const.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/const.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c74b7fb4ae5e49e6fc88657a8a91f68cd2f681c Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/const.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/exc.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/exc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35b01c5c0b85a9316f0da761cf296dcf0190d7a2 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/exc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/fun.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/fun.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3aea792b611d3f89f7aef0d8bc258907e30f88f3 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/fun.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/pack.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/pack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7467be11ee567baab53b90f01e06252d919e9731 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/pack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/stream.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/stream.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9ea85188c5b81ca3890aa3788a1f93e468be086 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/stream.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/typ.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/typ.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62b4ccc4732136d1b3ac4d076d45d212e9490fca Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/typ.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..925f120bce63776c34f19f4dc66c7ffacabd10e5 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/base.py b/lib/python3.12/site-packages/gitdb/base.py new file mode 100644 index 0000000000000000000000000000000000000000..9a23a4f71acb85106cff88d7cab619915d67d39e --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/base.py @@ -0,0 +1,315 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Module with basic data structures - they are designed to be lightweight and fast""" +from gitdb.util import bin_to_hex + +from gitdb.fun import ( + type_id_to_type_map, + type_to_type_id_map +) + +__all__ = ('OInfo', 'OPackInfo', 'ODeltaPackInfo', + 'OStream', 'OPackStream', 'ODeltaPackStream', + 'IStream', 'InvalidOInfo', 'InvalidOStream') + +#{ ODB Bases + + +class OInfo(tuple): + + """Carries information about an object in an ODB, providing information + about the binary sha of the object, the type_string as well as the uncompressed size + in bytes. + + It can be accessed using tuple notation and using attribute access notation:: + + assert dbi[0] == dbi.binsha + assert dbi[1] == dbi.type + assert dbi[2] == dbi.size + + The type is designed to be as lightweight as possible.""" + __slots__ = tuple() + + def __new__(cls, sha, type, size): + return tuple.__new__(cls, (sha, type, size)) + + def __init__(self, *args): + tuple.__init__(self) + + #{ Interface + @property + def binsha(self): + """:return: our sha as binary, 20 bytes""" + return self[0] + + @property + def hexsha(self): + """:return: our sha, hex encoded, 40 bytes""" + return bin_to_hex(self[0]) + + @property + def type(self): + return self[1] + + @property + def type_id(self): + return type_to_type_id_map[self[1]] + + @property + def size(self): + return self[2] + #} END interface + + +class OPackInfo(tuple): + + """As OInfo, but provides a type_id property to retrieve the numerical type id, and + does not include a sha. + + Additionally, the pack_offset is the absolute offset into the packfile at which + all object information is located. The data_offset property points to the absolute + location in the pack at which that actual data stream can be found.""" + __slots__ = tuple() + + def __new__(cls, packoffset, type, size): + return tuple.__new__(cls, (packoffset, type, size)) + + def __init__(self, *args): + tuple.__init__(self) + + #{ Interface + + @property + def pack_offset(self): + return self[0] + + @property + def type(self): + return type_id_to_type_map[self[1]] + + @property + def type_id(self): + return self[1] + + @property + def size(self): + return self[2] + + #} END interface + + +class ODeltaPackInfo(OPackInfo): + + """Adds delta specific information, + Either the 20 byte sha which points to some object in the database, + or the negative offset from the pack_offset, so that pack_offset - delta_info yields + the pack offset of the base object""" + __slots__ = tuple() + + def __new__(cls, packoffset, type, size, delta_info): + return tuple.__new__(cls, (packoffset, type, size, delta_info)) + + #{ Interface + @property + def delta_info(self): + return self[3] + #} END interface + + +class OStream(OInfo): + + """Base for object streams retrieved from the database, providing additional + information about the stream. + Generally, ODB streams are read-only as objects are immutable""" + __slots__ = tuple() + + def __new__(cls, sha, type, size, stream, *args, **kwargs): + """Helps with the initialization of subclasses""" + return tuple.__new__(cls, (sha, type, size, stream)) + + def __init__(self, *args, **kwargs): + tuple.__init__(self) + + #{ Stream Reader Interface + + def read(self, size=-1): + return self[3].read(size) + + @property + def stream(self): + return self[3] + + #} END stream reader interface + + +class ODeltaStream(OStream): + + """Uses size info of its stream, delaying reads""" + + def __new__(cls, sha, type, size, stream, *args, **kwargs): + """Helps with the initialization of subclasses""" + return tuple.__new__(cls, (sha, type, size, stream)) + + #{ Stream Reader Interface + + @property + def size(self): + return self[3].size + + #} END stream reader interface + + +class OPackStream(OPackInfo): + + """Next to pack object information, a stream outputting an undeltified base object + is provided""" + __slots__ = tuple() + + def __new__(cls, packoffset, type, size, stream, *args): + """Helps with the initialization of subclasses""" + return tuple.__new__(cls, (packoffset, type, size, stream)) + + #{ Stream Reader Interface + def read(self, size=-1): + return self[3].read(size) + + @property + def stream(self): + return self[3] + #} END stream reader interface + + +class ODeltaPackStream(ODeltaPackInfo): + + """Provides a stream outputting the uncompressed offset delta information""" + __slots__ = tuple() + + def __new__(cls, packoffset, type, size, delta_info, stream): + return tuple.__new__(cls, (packoffset, type, size, delta_info, stream)) + + #{ Stream Reader Interface + def read(self, size=-1): + return self[4].read(size) + + @property + def stream(self): + return self[4] + #} END stream reader interface + + +class IStream(list): + + """Represents an input content stream to be fed into the ODB. It is mutable to allow + the ODB to record information about the operations outcome right in this instance. + + It provides interfaces for the OStream and a StreamReader to allow the instance + to blend in without prior conversion. + + The only method your content stream must support is 'read'""" + __slots__ = tuple() + + def __new__(cls, type, size, stream, sha=None): + return list.__new__(cls, (sha, type, size, stream, None)) + + def __init__(self, type, size, stream, sha=None): + list.__init__(self, (sha, type, size, stream, None)) + + #{ Interface + @property + def hexsha(self): + """:return: our sha, hex encoded, 40 bytes""" + return bin_to_hex(self[0]) + + def _error(self): + """:return: the error that occurred when processing the stream, or None""" + return self[4] + + def _set_error(self, exc): + """Set this input stream to the given exc, may be None to reset the error""" + self[4] = exc + + error = property(_error, _set_error) + + #} END interface + + #{ Stream Reader Interface + + def read(self, size=-1): + """Implements a simple stream reader interface, passing the read call on + to our internal stream""" + return self[3].read(size) + + #} END stream reader interface + + #{ interface + + def _set_binsha(self, binsha): + self[0] = binsha + + def _binsha(self): + return self[0] + + binsha = property(_binsha, _set_binsha) + + def _type(self): + return self[1] + + def _set_type(self, type): + self[1] = type + + type = property(_type, _set_type) + + def _size(self): + return self[2] + + def _set_size(self, size): + self[2] = size + + size = property(_size, _set_size) + + def _stream(self): + return self[3] + + def _set_stream(self, stream): + self[3] = stream + + stream = property(_stream, _set_stream) + + #} END odb info interface + + +class InvalidOInfo(tuple): + + """Carries information about a sha identifying an object which is invalid in + the queried database. The exception attribute provides more information about + the cause of the issue""" + __slots__ = tuple() + + def __new__(cls, sha, exc): + return tuple.__new__(cls, (sha, exc)) + + def __init__(self, sha, exc): + tuple.__init__(self, (sha, exc)) + + @property + def binsha(self): + return self[0] + + @property + def hexsha(self): + return bin_to_hex(self[0]) + + @property + def error(self): + """:return: exception instance explaining the failure""" + return self[1] + + +class InvalidOStream(InvalidOInfo): + + """Carries information about an invalid ODB stream""" + __slots__ = tuple() + +#} END ODB Bases diff --git a/lib/python3.12/site-packages/gitdb/const.py b/lib/python3.12/site-packages/gitdb/const.py new file mode 100644 index 0000000000000000000000000000000000000000..6391d796f3fa3714802c8858960f5cfb3947c230 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/const.py @@ -0,0 +1,4 @@ +BYTE_SPACE = b' ' +NULL_BYTE = b'\0' +NULL_HEX_SHA = "0" * 40 +NULL_BIN_SHA = NULL_BYTE * 20 diff --git a/lib/python3.12/site-packages/gitdb/db/__init__.py b/lib/python3.12/site-packages/gitdb/db/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..20fd2280db7a978d8204e7b1566967ab7d3e772a --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/__init__.py @@ -0,0 +1,11 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ + +from gitdb.db.base import * +from gitdb.db.loose import * +from gitdb.db.mem import * +from gitdb.db.pack import * +from gitdb.db.git import * +from gitdb.db.ref import * diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca03d0056fa54f53b33173ac9b4a980662e31101 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56af6535ce283a15b010833643e9041d12cbbf01 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/git.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/git.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acd3bd0b023f0a9f7bda01e6d4ceda1d70d85963 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/git.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/loose.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/loose.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..264700d81e7c93c4cf4411a3783b9154dcaf4244 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/loose.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/mem.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/mem.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdc87a22836d736f8b61f35c89aa8278750bd6d5 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/mem.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/pack.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/pack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55170d67d85e96a05e5e038112e4a78c8e3caebb Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/pack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/__pycache__/ref.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/db/__pycache__/ref.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b50331b6889dd62cac7330687610ca3ce5250aa Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/db/__pycache__/ref.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/db/base.py b/lib/python3.12/site-packages/gitdb/db/base.py new file mode 100644 index 0000000000000000000000000000000000000000..7312fe00f3ec772e7abe508b2eeff63c5e79a7e3 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/base.py @@ -0,0 +1,278 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Contains implementations of database retrieveing objects""" +from gitdb.util import ( + join, + LazyMixin, + hex_to_bin +) + +from gitdb.utils.encoding import force_text +from gitdb.exc import ( + BadObject, + AmbiguousObjectName +) + +from itertools import chain +from functools import reduce + + +__all__ = ('ObjectDBR', 'ObjectDBW', 'FileDBBase', 'CompoundDB', 'CachingDB') + + +class ObjectDBR: + + """Defines an interface for object database lookup. + Objects are identified either by their 20 byte bin sha""" + + def __contains__(self, sha): + return self.has_obj + + #{ Query Interface + def has_object(self, sha): + """ + Whether the object identified by the given 20 bytes + binary sha is contained in the database + + :return: True if the object identified by the given 20 bytes + binary sha is contained in the database""" + raise NotImplementedError("To be implemented in subclass") + + def info(self, sha): + """ :return: OInfo instance + :param sha: bytes binary sha + :raise BadObject:""" + raise NotImplementedError("To be implemented in subclass") + + def stream(self, sha): + """:return: OStream instance + :param sha: 20 bytes binary sha + :raise BadObject:""" + raise NotImplementedError("To be implemented in subclass") + + def size(self): + """:return: amount of objects in this database""" + raise NotImplementedError() + + def sha_iter(self): + """Return iterator yielding 20 byte shas for all objects in this data base""" + raise NotImplementedError() + + #} END query interface + + +class ObjectDBW: + + """Defines an interface to create objects in the database""" + + def __init__(self, *args, **kwargs): + self._ostream = None + + #{ Edit Interface + def set_ostream(self, stream): + """ + Adjusts the stream to which all data should be sent when storing new objects + + :param stream: if not None, the stream to use, if None the default stream + will be used. + :return: previously installed stream, or None if there was no override + :raise TypeError: if the stream doesn't have the supported functionality""" + cstream = self._ostream + self._ostream = stream + return cstream + + def ostream(self): + """ + Return the output stream + + :return: overridden output stream this instance will write to, or None + if it will write to the default stream""" + return self._ostream + + def store(self, istream): + """ + Create a new object in the database + :return: the input istream object with its sha set to its corresponding value + + :param istream: IStream compatible instance. If its sha is already set + to a value, the object will just be stored in the our database format, + in which case the input stream is expected to be in object format ( header + contents ). + :raise IOError: if data could not be written""" + raise NotImplementedError("To be implemented in subclass") + + #} END edit interface + + +class FileDBBase: + + """Provides basic facilities to retrieve files of interest, including + caching facilities to help mapping hexsha's to objects""" + + def __init__(self, root_path): + """Initialize this instance to look for its files at the given root path + All subsequent operations will be relative to this path + :raise InvalidDBRoot: + **Note:** The base will not perform any accessablity checking as the base + might not yet be accessible, but become accessible before the first + access.""" + super().__init__() + self._root_path = root_path + + #{ Interface + def root_path(self): + """:return: path at which this db operates""" + return self._root_path + + def db_path(self, rela_path): + """ + :return: the given relative path relative to our database root, allowing + to pontentially access datafiles""" + return join(self._root_path, force_text(rela_path)) + #} END interface + + +class CachingDB: + + """A database which uses caches to speed-up access""" + + #{ Interface + def update_cache(self, force=False): + """ + Call this method if the underlying data changed to trigger an update + of the internal caching structures. + + :param force: if True, the update must be performed. Otherwise the implementation + may decide not to perform an update if it thinks nothing has changed. + :return: True if an update was performed as something change indeed""" + + # END interface + + +def _databases_recursive(database, output): + """Fill output list with database from db, in order. Deals with Loose, Packed + and compound databases.""" + if isinstance(database, CompoundDB): + dbs = database.databases() + output.extend(db for db in dbs if not isinstance(db, CompoundDB)) + for cdb in (db for db in dbs if isinstance(db, CompoundDB)): + _databases_recursive(cdb, output) + else: + output.append(database) + # END handle database type + + +class CompoundDB(ObjectDBR, LazyMixin, CachingDB): + + """A database which delegates calls to sub-databases. + + Databases are stored in the lazy-loaded _dbs attribute. + Define _set_cache_ to update it with your databases""" + + def _set_cache_(self, attr): + if attr == '_dbs': + self._dbs = list() + elif attr == '_db_cache': + self._db_cache = dict() + else: + super()._set_cache_(attr) + + def _db_query(self, sha): + """:return: database containing the given 20 byte sha + :raise BadObject:""" + # most databases use binary representations, prevent converting + # it every time a database is being queried + try: + return self._db_cache[sha] + except KeyError: + pass + # END first level cache + + for db in self._dbs: + if db.has_object(sha): + self._db_cache[sha] = db + return db + # END for each database + raise BadObject(sha) + + #{ ObjectDBR interface + + def has_object(self, sha): + try: + self._db_query(sha) + return True + except BadObject: + return False + # END handle exceptions + + def info(self, sha): + return self._db_query(sha).info(sha) + + def stream(self, sha): + return self._db_query(sha).stream(sha) + + def size(self): + """:return: total size of all contained databases""" + return reduce(lambda x, y: x + y, (db.size() for db in self._dbs), 0) + + def sha_iter(self): + return chain(*(db.sha_iter() for db in self._dbs)) + + #} END object DBR Interface + + #{ Interface + + def databases(self): + """:return: tuple of database instances we use for lookups""" + return tuple(self._dbs) + + def update_cache(self, force=False): + # something might have changed, clear everything + self._db_cache.clear() + stat = False + for db in self._dbs: + if isinstance(db, CachingDB): + stat |= db.update_cache(force) + # END if is caching db + # END for each database to update + return stat + + def partial_to_complete_sha_hex(self, partial_hexsha): + """ + :return: 20 byte binary sha1 from the given less-than-40 byte hexsha (bytes or str) + :param partial_hexsha: hexsha with less than 40 byte + :raise AmbiguousObjectName: """ + databases = list() + _databases_recursive(self, databases) + partial_hexsha = force_text(partial_hexsha) + len_partial_hexsha = len(partial_hexsha) + if len_partial_hexsha % 2 != 0: + partial_binsha = hex_to_bin(partial_hexsha + "0") + else: + partial_binsha = hex_to_bin(partial_hexsha) + # END assure successful binary conversion + + candidate = None + for db in databases: + full_bin_sha = None + try: + if hasattr(db, 'partial_to_complete_sha_hex'): + full_bin_sha = db.partial_to_complete_sha_hex(partial_hexsha) + else: + full_bin_sha = db.partial_to_complete_sha(partial_binsha, len_partial_hexsha) + # END handle database type + except BadObject: + continue + # END ignore bad objects + if full_bin_sha: + if candidate and candidate != full_bin_sha: + raise AmbiguousObjectName(partial_hexsha) + candidate = full_bin_sha + # END handle candidate + # END for each db + if not candidate: + raise BadObject(partial_binsha) + return candidate + + #} END interface diff --git a/lib/python3.12/site-packages/gitdb/db/git.py b/lib/python3.12/site-packages/gitdb/db/git.py new file mode 100644 index 0000000000000000000000000000000000000000..a1ed1428054f94e5ccda150ce7d8267b31ce55fd --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/git.py @@ -0,0 +1,85 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +from gitdb.db.base import ( + CompoundDB, + ObjectDBW, + FileDBBase +) + +from gitdb.db.loose import LooseObjectDB +from gitdb.db.pack import PackedDB +from gitdb.db.ref import ReferenceDB + +from gitdb.exc import InvalidDBRoot + +import os + +__all__ = ('GitDB', ) + + +class GitDB(FileDBBase, ObjectDBW, CompoundDB): + + """A git-style object database, which contains all objects in the 'objects' + subdirectory + + ``IMPORTANT``: The usage of this implementation is highly discouraged as it fails to release file-handles. + This can be a problem with long-running processes and/or big repositories. + """ + # Configuration + PackDBCls = PackedDB + LooseDBCls = LooseObjectDB + ReferenceDBCls = ReferenceDB + + # Directories + packs_dir = 'pack' + loose_dir = '' + alternates_dir = os.path.join('info', 'alternates') + + def __init__(self, root_path): + """Initialize ourselves on a git objects directory""" + super().__init__(root_path) + + def _set_cache_(self, attr): + if attr == '_dbs' or attr == '_loose_db': + self._dbs = list() + loose_db = None + for subpath, dbcls in ((self.packs_dir, self.PackDBCls), + (self.loose_dir, self.LooseDBCls), + (self.alternates_dir, self.ReferenceDBCls)): + path = self.db_path(subpath) + if os.path.exists(path): + self._dbs.append(dbcls(path)) + if dbcls is self.LooseDBCls: + loose_db = self._dbs[-1] + # END remember loose db + # END check path exists + # END for each db type + + # should have at least one subdb + if not self._dbs: + raise InvalidDBRoot(self.root_path()) + # END handle error + + # we the first one should have the store method + assert loose_db is not None and hasattr(loose_db, 'store'), "First database needs store functionality" + + # finally set the value + self._loose_db = loose_db + else: + super()._set_cache_(attr) + # END handle attrs + + #{ ObjectDBW interface + + def store(self, istream): + return self._loose_db.store(istream) + + def ostream(self): + return self._loose_db.ostream() + + def set_ostream(self, ostream): + return self._loose_db.set_ostream(ostream) + + #} END objectdbw interface diff --git a/lib/python3.12/site-packages/gitdb/db/loose.py b/lib/python3.12/site-packages/gitdb/db/loose.py new file mode 100644 index 0000000000000000000000000000000000000000..03d387e868b2ca49585803e8f9a4aed8854273ed --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/loose.py @@ -0,0 +1,268 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +from contextlib import suppress + +from gitdb.db.base import ( + FileDBBase, + ObjectDBR, + ObjectDBW +) + +from gitdb.exc import ( + BadObject, + AmbiguousObjectName +) + +from gitdb.stream import ( + DecompressMemMapReader, + FDCompressedSha1Writer, + FDStream, + Sha1Writer +) + +from gitdb.base import ( + OStream, + OInfo +) + +from gitdb.util import ( + file_contents_ro_filepath, + ENOENT, + hex_to_bin, + bin_to_hex, + exists, + chmod, + isfile, + remove, + rename, + dirname, + basename, + join +) + +from gitdb.fun import ( + chunk_size, + loose_object_header_info, + write_object, + stream_copy +) + +from gitdb.utils.encoding import force_bytes + +import tempfile +import os +import sys +import time + + +__all__ = ('LooseObjectDB', ) + + +class LooseObjectDB(FileDBBase, ObjectDBR, ObjectDBW): + + """A database which operates on loose object files""" + + # CONFIGURATION + # chunks in which data will be copied between streams + stream_chunk_size = chunk_size + + # On windows we need to keep it writable, otherwise it cannot be removed + # either + new_objects_mode = int("444", 8) + if os.name == 'nt': + new_objects_mode = int("644", 8) + + def __init__(self, root_path): + super().__init__(root_path) + self._hexsha_to_file = dict() + # Additional Flags - might be set to 0 after the first failure + # Depending on the root, this might work for some mounts, for others not, which + # is why it is per instance + self._fd_open_flags = getattr(os, 'O_NOATIME', 0) + + #{ Interface + def object_path(self, hexsha): + """ + :return: path at which the object with the given hexsha would be stored, + relative to the database root""" + return join(hexsha[:2], hexsha[2:]) + + def readable_db_object_path(self, hexsha): + """ + :return: readable object path to the object identified by hexsha + :raise BadObject: If the object file does not exist""" + with suppress(KeyError): + return self._hexsha_to_file[hexsha] + # END ignore cache misses + + # try filesystem + path = self.db_path(self.object_path(hexsha)) + if exists(path): + self._hexsha_to_file[hexsha] = path + return path + # END handle cache + raise BadObject(hexsha) + + def partial_to_complete_sha_hex(self, partial_hexsha): + """:return: 20 byte binary sha1 string which matches the given name uniquely + :param name: hexadecimal partial name (bytes or ascii string) + :raise AmbiguousObjectName: + :raise BadObject: """ + candidate = None + for binsha in self.sha_iter(): + if bin_to_hex(binsha).startswith(force_bytes(partial_hexsha)): + # it can't ever find the same object twice + if candidate is not None: + raise AmbiguousObjectName(partial_hexsha) + candidate = binsha + # END for each object + if candidate is None: + raise BadObject(partial_hexsha) + return candidate + + #} END interface + + def _map_loose_object(self, sha): + """ + :return: memory map of that file to allow random read access + :raise BadObject: if object could not be located""" + db_path = self.db_path(self.object_path(bin_to_hex(sha))) + try: + return file_contents_ro_filepath(db_path, flags=self._fd_open_flags) + except OSError as e: + if e.errno != ENOENT: + # try again without noatime + try: + return file_contents_ro_filepath(db_path) + except OSError as new_e: + raise BadObject(sha) from new_e + # didn't work because of our flag, don't try it again + self._fd_open_flags = 0 + else: + raise BadObject(sha) from e + # END handle error + # END exception handling + + def set_ostream(self, stream): + """:raise TypeError: if the stream does not support the Sha1Writer interface""" + if stream is not None and not isinstance(stream, Sha1Writer): + raise TypeError("Output stream musst support the %s interface" % Sha1Writer.__name__) + return super().set_ostream(stream) + + def info(self, sha): + m = self._map_loose_object(sha) + try: + typ, size = loose_object_header_info(m) + return OInfo(sha, typ, size) + finally: + if hasattr(m, 'close'): + m.close() + # END assure release of system resources + + def stream(self, sha): + m = self._map_loose_object(sha) + type, size, stream = DecompressMemMapReader.new(m, close_on_deletion=True) + return OStream(sha, type, size, stream) + + def has_object(self, sha): + try: + self.readable_db_object_path(bin_to_hex(sha)) + return True + except BadObject: + return False + # END check existence + + def store(self, istream): + """note: The sha we produce will be hex by nature""" + tmp_path = None + writer = self.ostream() + if writer is None: + # open a tmp file to write the data to + fd, tmp_path = tempfile.mkstemp(prefix='obj', dir=self._root_path) + + if istream.binsha is None: + writer = FDCompressedSha1Writer(fd) + else: + writer = FDStream(fd) + # END handle direct stream copies + # END handle custom writer + + try: + try: + if istream.binsha is not None: + # copy as much as possible, the actual uncompressed item size might + # be smaller than the compressed version + stream_copy(istream.read, writer.write, sys.maxsize, self.stream_chunk_size) + else: + # write object with header, we have to make a new one + write_object(istream.type, istream.size, istream.read, writer.write, + chunk_size=self.stream_chunk_size) + # END handle direct stream copies + finally: + if tmp_path: + writer.close() + # END assure target stream is closed + except: + if tmp_path: + remove(tmp_path) + raise + # END assure tmpfile removal on error + + hexsha = None + if istream.binsha: + hexsha = istream.hexsha + else: + hexsha = writer.sha(as_hex=True) + # END handle sha + + if tmp_path: + obj_path = self.db_path(self.object_path(hexsha)) + obj_dir = dirname(obj_path) + os.makedirs(obj_dir, exist_ok=True) + # END handle destination directory + # rename onto existing doesn't work on NTFS + if isfile(obj_path): + remove(tmp_path) + else: + rename(tmp_path, obj_path) + # end rename only if needed + + # Ensure rename is actually done and file is stable + # Retry up to 14 times - exponential wait & retry in ms. + # The total maximum wait time is 1000ms, which should be vastly enough for the + # OS to return and commit the file to disk. + for exp_backoff_ms in [1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 181]: + with suppress(PermissionError): + # make sure its readable for all ! It started out as rw-- tmp file + # but needs to be rwrr + chmod(obj_path, self.new_objects_mode) + break + time.sleep(exp_backoff_ms / 1000.0) + else: + raise PermissionError( + "Impossible to apply `chmod` to file {}".format(obj_path) + ) + + # END handle dry_run + + istream.binsha = hex_to_bin(hexsha) + return istream + + def sha_iter(self): + # find all files which look like an object, extract sha from there + for root, dirs, files in os.walk(self.root_path()): + root_base = basename(root) + if len(root_base) != 2: + continue + + for f in files: + if len(f) != 38: + continue + yield hex_to_bin(root_base + f) + # END for each file + # END for each walk iteration + + def size(self): + return len(tuple(self.sha_iter())) diff --git a/lib/python3.12/site-packages/gitdb/db/mem.py b/lib/python3.12/site-packages/gitdb/db/mem.py new file mode 100644 index 0000000000000000000000000000000000000000..d4772fdb5e5d98b14c5fb329f4198e2cc99f8906 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/mem.py @@ -0,0 +1,110 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Contains the MemoryDatabase implementation""" +from gitdb.db.loose import LooseObjectDB +from gitdb.db.base import ( + ObjectDBR, + ObjectDBW +) + +from gitdb.base import ( + OStream, + IStream, +) + +from gitdb.exc import ( + BadObject, + UnsupportedOperation +) + +from gitdb.stream import ( + ZippedStoreShaWriter, + DecompressMemMapReader, +) + +from io import BytesIO + +__all__ = ("MemoryDB", ) + + +class MemoryDB(ObjectDBR, ObjectDBW): + + """A memory database stores everything to memory, providing fast IO and object + retrieval. It should be used to buffer results and obtain SHAs before writing + it to the actual physical storage, as it allows to query whether object already + exists in the target storage before introducing actual IO""" + + def __init__(self): + super().__init__() + self._db = LooseObjectDB("path/doesnt/matter") + + # maps 20 byte shas to their OStream objects + self._cache = dict() + + def set_ostream(self, stream): + raise UnsupportedOperation("MemoryDB's always stream into memory") + + def store(self, istream): + zstream = ZippedStoreShaWriter() + self._db.set_ostream(zstream) + + istream = self._db.store(istream) + zstream.close() # close to flush + zstream.seek(0) + + # don't provide a size, the stream is written in object format, hence the + # header needs decompression + decomp_stream = DecompressMemMapReader(zstream.getvalue(), close_on_deletion=False) + self._cache[istream.binsha] = OStream(istream.binsha, istream.type, istream.size, decomp_stream) + + return istream + + def has_object(self, sha): + return sha in self._cache + + def info(self, sha): + # we always return streams, which are infos as well + return self.stream(sha) + + def stream(self, sha): + try: + ostream = self._cache[sha] + # rewind stream for the next one to read + ostream.stream.seek(0) + return ostream + except KeyError as e: + raise BadObject(sha) from e + # END exception handling + + def size(self): + return len(self._cache) + + def sha_iter(self): + return self._cache.keys() + + #{ Interface + def stream_copy(self, sha_iter, odb): + """Copy the streams as identified by sha's yielded by sha_iter into the given odb + The streams will be copied directly + **Note:** the object will only be written if it did not exist in the target db + + :return: amount of streams actually copied into odb. If smaller than the amount + of input shas, one or more objects did already exist in odb""" + count = 0 + for sha in sha_iter: + if odb.has_object(sha): + continue + # END check object existence + + ostream = self.stream(sha) + # compressed data including header + sio = BytesIO(ostream.stream.data()) + istream = IStream(ostream.type, ostream.size, sio, sha) + + odb.store(istream) + count += 1 + # END for each sha + return count + #} END interface diff --git a/lib/python3.12/site-packages/gitdb/db/pack.py b/lib/python3.12/site-packages/gitdb/db/pack.py new file mode 100644 index 0000000000000000000000000000000000000000..274ea59904ee6b83ffa98d13df130e5a1e0d75f2 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/pack.py @@ -0,0 +1,206 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Module containing a database to deal with packs""" +from gitdb.db.base import ( + FileDBBase, + ObjectDBR, + CachingDB +) + +from gitdb.util import LazyMixin + +from gitdb.exc import ( + BadObject, + UnsupportedOperation, + AmbiguousObjectName +) + +from gitdb.pack import PackEntity + +from functools import reduce + +import os +import glob + +__all__ = ('PackedDB', ) + +#{ Utilities + + +class PackedDB(FileDBBase, ObjectDBR, CachingDB, LazyMixin): + + """A database operating on a set of object packs""" + + # sort the priority list every N queries + # Higher values are better, performance tests don't show this has + # any effect, but it should have one + _sort_interval = 500 + + def __init__(self, root_path): + super().__init__(root_path) + # list of lists with three items: + # * hits - number of times the pack was hit with a request + # * entity - Pack entity instance + # * sha_to_index - PackIndexFile.sha_to_index method for direct cache query + # self._entities = list() # lazy loaded list + self._hit_count = 0 # amount of hits + self._st_mtime = 0 # last modification data of our root path + + def _set_cache_(self, attr): + if attr == '_entities': + self._entities = list() + self.update_cache(force=True) + # END handle entities initialization + + def _sort_entities(self): + self._entities.sort(key=lambda l: l[0], reverse=True) + + def _pack_info(self, sha): + """:return: tuple(entity, index) for an item at the given sha + :param sha: 20 or 40 byte sha + :raise BadObject: + **Note:** This method is not thread-safe, but may be hit in multi-threaded + operation. The worst thing that can happen though is a counter that + was not incremented, or the list being in wrong order. So we safe + the time for locking here, lets see how that goes""" + # presort ? + if self._hit_count % self._sort_interval == 0: + self._sort_entities() + # END update sorting + + for item in self._entities: + index = item[2](sha) + if index is not None: + item[0] += 1 # one hit for you + self._hit_count += 1 # general hit count + return (item[1], index) + # END index found in pack + # END for each item + + # no hit, see whether we have to update packs + # NOTE: considering packs don't change very often, we safe this call + # and leave it to the super-caller to trigger that + raise BadObject(sha) + + #{ Object DB Read + + def has_object(self, sha): + try: + self._pack_info(sha) + return True + except BadObject: + return False + # END exception handling + + def info(self, sha): + entity, index = self._pack_info(sha) + return entity.info_at_index(index) + + def stream(self, sha): + entity, index = self._pack_info(sha) + return entity.stream_at_index(index) + + def sha_iter(self): + for entity in self.entities(): + index = entity.index() + sha_by_index = index.sha + for index in range(index.size()): + yield sha_by_index(index) + # END for each index + # END for each entity + + def size(self): + sizes = [item[1].index().size() for item in self._entities] + return reduce(lambda x, y: x + y, sizes, 0) + + #} END object db read + + #{ object db write + + def store(self, istream): + """Storing individual objects is not feasible as a pack is designed to + hold multiple objects. Writing or rewriting packs for single objects is + inefficient""" + raise UnsupportedOperation() + + #} END object db write + + #{ Interface + + def update_cache(self, force=False): + """ + Update our cache with the actually existing packs on disk. Add new ones, + and remove deleted ones. We keep the unchanged ones + + :param force: If True, the cache will be updated even though the directory + does not appear to have changed according to its modification timestamp. + :return: True if the packs have been updated so there is new information, + False if there was no change to the pack database""" + stat = os.stat(self.root_path()) + if not force and stat.st_mtime <= self._st_mtime: + return False + # END abort early on no change + self._st_mtime = stat.st_mtime + + # packs are supposed to be prefixed with pack- by git-convention + # get all pack files, figure out what changed + pack_files = set(glob.glob(os.path.join(self.root_path(), "pack-*.pack"))) + our_pack_files = {item[1].pack().path() for item in self._entities} + + # new packs + for pack_file in (pack_files - our_pack_files): + # init the hit-counter/priority with the size, a good measure for hit- + # probability. Its implemented so that only 12 bytes will be read + entity = PackEntity(pack_file) + self._entities.append([entity.pack().size(), entity, entity.index().sha_to_index]) + # END for each new packfile + + # removed packs + for pack_file in (our_pack_files - pack_files): + del_index = -1 + for i, item in enumerate(self._entities): + if item[1].pack().path() == pack_file: + del_index = i + break + # END found index + # END for each entity + assert del_index != -1 + del(self._entities[del_index]) + # END for each removed pack + + # reinitialize prioritiess + self._sort_entities() + return True + + def entities(self): + """:return: list of pack entities operated upon by this database""" + return [item[1] for item in self._entities] + + def partial_to_complete_sha(self, partial_binsha, canonical_length): + """:return: 20 byte sha as inferred by the given partial binary sha + :param partial_binsha: binary sha with less than 20 bytes + :param canonical_length: length of the corresponding canonical representation. + It is required as binary sha's cannot display whether the original hex sha + had an odd or even number of characters + :raise AmbiguousObjectName: + :raise BadObject: """ + candidate = None + for item in self._entities: + item_index = item[1].index().partial_sha_to_index(partial_binsha, canonical_length) + if item_index is not None: + sha = item[1].index().sha(item_index) + if candidate and candidate != sha: + raise AmbiguousObjectName(partial_binsha) + candidate = sha + # END handle full sha could be found + # END for each entity + + if candidate: + return candidate + + # still not found ? + raise BadObject(partial_binsha) + + #} END interface diff --git a/lib/python3.12/site-packages/gitdb/db/ref.py b/lib/python3.12/site-packages/gitdb/db/ref.py new file mode 100644 index 0000000000000000000000000000000000000000..bd3015602f2967441ca8b27813424b7bfab9160b --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/db/ref.py @@ -0,0 +1,82 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +import codecs +from gitdb.db.base import ( + CompoundDB, +) + +__all__ = ('ReferenceDB', ) + + +class ReferenceDB(CompoundDB): + + """A database consisting of database referred to in a file""" + + # Configuration + # Specifies the object database to use for the paths found in the alternates + # file. If None, it defaults to the GitDB + ObjectDBCls = None + + def __init__(self, ref_file): + super().__init__() + self._ref_file = ref_file + + def _set_cache_(self, attr): + if attr == '_dbs': + self._dbs = list() + self._update_dbs_from_ref_file() + else: + super()._set_cache_(attr) + # END handle attrs + + def _update_dbs_from_ref_file(self): + dbcls = self.ObjectDBCls + if dbcls is None: + # late import + from gitdb.db.git import GitDB + dbcls = GitDB + # END get db type + + # try to get as many as possible, don't fail if some are unavailable + ref_paths = list() + try: + with codecs.open(self._ref_file, 'r', encoding="utf-8") as f: + ref_paths = [l.strip() for l in f] + except OSError: + pass + # END handle alternates + + ref_paths_set = set(ref_paths) + cur_ref_paths_set = {db.root_path() for db in self._dbs} + + # remove existing + for path in (cur_ref_paths_set - ref_paths_set): + for i, db in enumerate(self._dbs[:]): + if db.root_path() == path: + del(self._dbs[i]) + continue + # END del matching db + # END for each path to remove + + # add new + # sort them to maintain order + added_paths = sorted(ref_paths_set - cur_ref_paths_set, key=lambda p: ref_paths.index(p)) + for path in added_paths: + try: + db = dbcls(path) + # force an update to verify path + if isinstance(db, CompoundDB): + db.databases() + # END verification + self._dbs.append(db) + except Exception: + # ignore invalid paths or issues + pass + # END for each path to add + + def update_cache(self, force=False): + # re-read alternates and update databases + self._update_dbs_from_ref_file() + return super().update_cache(force) diff --git a/lib/python3.12/site-packages/gitdb/exc.py b/lib/python3.12/site-packages/gitdb/exc.py new file mode 100644 index 0000000000000000000000000000000000000000..752dafdbbc837e577617f9e5f1abaf101d93cdd3 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/exc.py @@ -0,0 +1,57 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Module with common exceptions""" +from gitdb.util import to_hex_sha + +__all__ = [ + 'AmbiguousObjectName', + 'BadName', + 'BadObject', + 'BadObjectType', + 'InvalidDBRoot', + 'ODBError', + 'ParseError', + 'UnsupportedOperation', + 'to_hex_sha', +] + +class ODBError(Exception): + """All errors thrown by the object database""" + + +class InvalidDBRoot(ODBError): + """Thrown if an object database cannot be initialized at the given path""" + + +class BadObject(ODBError): + """The object with the given SHA does not exist. Instantiate with the + failed sha""" + + def __str__(self): + return "BadObject: %s" % to_hex_sha(self.args[0]) + + +class BadName(ODBError): + """A name provided to rev_parse wasn't understood""" + + def __str__(self): + return "Ref '%s' did not resolve to an object" % self.args[0] + + +class ParseError(ODBError): + """Thrown if the parsing of a file failed due to an invalid format""" + + +class AmbiguousObjectName(ODBError): + """Thrown if a possibly shortened name does not uniquely represent a single object + in the database""" + + +class BadObjectType(ODBError): + """The object had an unsupported type""" + + +class UnsupportedOperation(ODBError): + """Thrown if the given operation cannot be supported by the object database""" diff --git a/lib/python3.12/site-packages/gitdb/fun.py b/lib/python3.12/site-packages/gitdb/fun.py new file mode 100644 index 0000000000000000000000000000000000000000..a272e5caa85e27b5f7f2c4adf3b396c3ce5c0b21 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/fun.py @@ -0,0 +1,704 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Contains basic c-functions which usually contain performance critical code +Keeping this code separate from the beginning makes it easier to out-source +it into c later, if required""" + +import zlib +from gitdb.util import byte_ord +decompressobj = zlib.decompressobj + +import mmap +from itertools import islice +from functools import reduce + +from gitdb.const import NULL_BYTE, BYTE_SPACE +from gitdb.utils.encoding import force_text +from gitdb.typ import ( + str_blob_type, + str_commit_type, + str_tree_type, + str_tag_type, +) + +from io import StringIO + +# INVARIANTS +OFS_DELTA = 6 +REF_DELTA = 7 +delta_types = (OFS_DELTA, REF_DELTA) + +type_id_to_type_map = { + 0: b'', # EXT 1 + 1: str_commit_type, + 2: str_tree_type, + 3: str_blob_type, + 4: str_tag_type, + 5: b'', # EXT 2 + OFS_DELTA: "OFS_DELTA", # OFFSET DELTA + REF_DELTA: "REF_DELTA" # REFERENCE DELTA +} + +type_to_type_id_map = { + str_commit_type: 1, + str_tree_type: 2, + str_blob_type: 3, + str_tag_type: 4, + "OFS_DELTA": OFS_DELTA, + "REF_DELTA": REF_DELTA, +} + +# used when dealing with larger streams +chunk_size = 1000 * mmap.PAGESIZE + +__all__ = ('is_loose_object', 'loose_object_header_info', 'msb_size', 'pack_object_header_info', + 'write_object', 'loose_object_header', 'stream_copy', 'apply_delta_data', + 'is_equal_canonical_sha', 'connect_deltas', 'DeltaChunkList', 'create_pack_object_header') + + +#{ Structures + +def _set_delta_rbound(d, size): + """Truncate the given delta to the given size + :param size: size relative to our target offset, may not be 0, must be smaller or equal + to our size + :return: d""" + d.ts = size + + # NOTE: data is truncated automatically when applying the delta + # MUST NOT DO THIS HERE + return d + + +def _move_delta_lbound(d, bytes): + """Move the delta by the given amount of bytes, reducing its size so that its + right bound stays static + :param bytes: amount of bytes to move, must be smaller than delta size + :return: d""" + if bytes == 0: + return + + d.to += bytes + d.so += bytes + d.ts -= bytes + if d.data is not None: + d.data = d.data[bytes:] + # END handle data + + return d + + +def delta_duplicate(src): + return DeltaChunk(src.to, src.ts, src.so, src.data) + + +def delta_chunk_apply(dc, bbuf, write): + """Apply own data to the target buffer + :param bbuf: buffer providing source bytes for copy operations + :param write: write method to call with data to write""" + if dc.data is None: + # COPY DATA FROM SOURCE + write(bbuf[dc.so:dc.so + dc.ts]) + else: + # APPEND DATA + # what's faster: if + 4 function calls or just a write with a slice ? + # Considering data can be larger than 127 bytes now, it should be worth it + if dc.ts < len(dc.data): + write(dc.data[:dc.ts]) + else: + write(dc.data) + # END handle truncation + # END handle chunk mode + + +class DeltaChunk: + + """Represents a piece of a delta, it can either add new data, or copy existing + one from a source buffer""" + __slots__ = ( + 'to', # start offset in the target buffer in bytes + 'ts', # size of this chunk in the target buffer in bytes + 'so', # start offset in the source buffer in bytes or None + 'data', # chunk of bytes to be added to the target buffer, + # DeltaChunkList to use as base, or None + ) + + def __init__(self, to, ts, so, data): + self.to = to + self.ts = ts + self.so = so + self.data = data + + def __repr__(self): + return "DeltaChunk(%i, %i, %s, %s)" % (self.to, self.ts, self.so, self.data or "") + + #{ Interface + + def rbound(self): + return self.to + self.ts + + def has_data(self): + """:return: True if the instance has data to add to the target stream""" + return self.data is not None + + #} END interface + + +def _closest_index(dcl, absofs): + """:return: index at which the given absofs should be inserted. The index points + to the DeltaChunk with a target buffer absofs that equals or is greater than + absofs. + **Note:** global method for performance only, it belongs to DeltaChunkList""" + lo = 0 + hi = len(dcl) + while lo < hi: + mid = (lo + hi) / 2 + dc = dcl[mid] + if dc.to > absofs: + hi = mid + elif dc.rbound() > absofs or dc.to == absofs: + return mid + else: + lo = mid + 1 + # END handle bound + # END for each delta absofs + return len(dcl) - 1 + + +def delta_list_apply(dcl, bbuf, write): + """Apply the chain's changes and write the final result using the passed + write function. + :param bbuf: base buffer containing the base of all deltas contained in this + list. It will only be used if the chunk in question does not have a base + chain. + :param write: function taking a string of bytes to write to the output""" + for dc in dcl: + delta_chunk_apply(dc, bbuf, write) + # END for each dc + + +def delta_list_slice(dcl, absofs, size, ndcl): + """:return: Subsection of this list at the given absolute offset, with the given + size in bytes. + :return: None""" + cdi = _closest_index(dcl, absofs) # delta start index + cd = dcl[cdi] + slen = len(dcl) + lappend = ndcl.append + + if cd.to != absofs: + tcd = DeltaChunk(cd.to, cd.ts, cd.so, cd.data) + _move_delta_lbound(tcd, absofs - cd.to) + tcd.ts = min(tcd.ts, size) + lappend(tcd) + size -= tcd.ts + cdi += 1 + # END lbound overlap handling + + while cdi < slen and size: + # are we larger than the current block + cd = dcl[cdi] + if cd.ts <= size: + lappend(DeltaChunk(cd.to, cd.ts, cd.so, cd.data)) + size -= cd.ts + else: + tcd = DeltaChunk(cd.to, cd.ts, cd.so, cd.data) + tcd.ts = size + lappend(tcd) + size -= tcd.ts + break + # END hadle size + cdi += 1 + # END for each chunk + + +class DeltaChunkList(list): + + """List with special functionality to deal with DeltaChunks. + There are two types of lists we represent. The one was created bottom-up, working + towards the latest delta, the other kind was created top-down, working from the + latest delta down to the earliest ancestor. This attribute is queryable + after all processing with is_reversed.""" + + __slots__ = tuple() + + def rbound(self): + """:return: rightmost extend in bytes, absolute""" + if len(self) == 0: + return 0 + return self[-1].rbound() + + def lbound(self): + """:return: leftmost byte at which this chunklist starts""" + if len(self) == 0: + return 0 + return self[0].to + + def size(self): + """:return: size of bytes as measured by our delta chunks""" + return self.rbound() - self.lbound() + + def apply(self, bbuf, write): + """Only used by public clients, internally we only use the global routines + for performance""" + return delta_list_apply(self, bbuf, write) + + def compress(self): + """Alter the list to reduce the amount of nodes. Currently we concatenate + add-chunks + :return: self""" + slen = len(self) + if slen < 2: + return self + i = 0 + + first_data_index = None + while i < slen: + dc = self[i] + i += 1 + if dc.data is None: + if first_data_index is not None and i - 2 - first_data_index > 1: + # if first_data_index is not None: + nd = StringIO() # new data + so = self[first_data_index].to # start offset in target buffer + for x in range(first_data_index, i - 1): + xdc = self[x] + nd.write(xdc.data[:xdc.ts]) + # END collect data + + del(self[first_data_index:i - 1]) + buf = nd.getvalue() + self.insert(first_data_index, DeltaChunk(so, len(buf), 0, buf)) + + slen = len(self) + i = first_data_index + 1 + + # END concatenate data + first_data_index = None + continue + # END skip non-data chunks + + if first_data_index is None: + first_data_index = i - 1 + # END iterate list + + # if slen_orig != len(self): + # print "INFO: Reduced delta list len to %f %% of former size" % ((float(len(self)) / slen_orig) * 100) + return self + + def check_integrity(self, target_size=-1): + """Verify the list has non-overlapping chunks only, and the total size matches + target_size + :param target_size: if not -1, the total size of the chain must be target_size + :raise AssertionError: if the size doesn't match""" + if target_size > -1: + assert self[-1].rbound() == target_size + assert reduce(lambda x, y: x + y, (d.ts for d in self), 0) == target_size + # END target size verification + + if len(self) < 2: + return + + # check data + for dc in self: + assert dc.ts > 0 + if dc.has_data(): + assert len(dc.data) >= dc.ts + # END for each dc + + left = islice(self, 0, len(self) - 1) + right = iter(self) + right.next() + # this is very pythonic - we might have just use index based access here, + # but this could actually be faster + for lft, rgt in zip(left, right): + assert lft.rbound() == rgt.to + assert lft.to + lft.ts == rgt.to + # END for each pair + + +class TopdownDeltaChunkList(DeltaChunkList): + + """Represents a list which is generated by feeding its ancestor streams one by + one""" + __slots__ = tuple() + + def connect_with_next_base(self, bdcl): + """Connect this chain with the next level of our base delta chunklist. + The goal in this game is to mark as many of our chunks rigid, hence they + cannot be changed by any of the upcoming bases anymore. Once all our + chunks are marked like that, we can stop all processing + :param bdcl: data chunk list being one of our bases. They must be fed in + consecutively and in order, towards the earliest ancestor delta + :return: True if processing was done. Use it to abort processing of + remaining streams if False is returned""" + nfc = 0 # number of frozen chunks + dci = 0 # delta chunk index + slen = len(self) # len of self + ccl = list() # temporary list + while dci < slen: + dc = self[dci] + dci += 1 + + # all add-chunks which are already topmost don't need additional processing + if dc.data is not None: + nfc += 1 + continue + # END skip add chunks + + # copy chunks + # integrate the portion of the base list into ourselves. Lists + # dont support efficient insertion ( just one at a time ), but for now + # we live with it. Internally, its all just a 32/64bit pointer, and + # the portions of moved memory should be smallish. Maybe we just rebuild + # ourselves in order to reduce the amount of insertions ... + del(ccl[:]) + delta_list_slice(bdcl, dc.so, dc.ts, ccl) + + # move the target bounds into place to match with our chunk + ofs = dc.to - dc.so + for cdc in ccl: + cdc.to += ofs + # END update target bounds + + if len(ccl) == 1: + self[dci - 1] = ccl[0] + else: + # maybe try to compute the expenses here, and pick the right algorithm + # It would normally be faster than copying everything physically though + # TODO: Use a deque here, and decide by the index whether to extend + # or extend left ! + post_dci = self[dci:] + del(self[dci - 1:]) # include deletion of dc + self.extend(ccl) + self.extend(post_dci) + + slen = len(self) + dci += len(ccl) - 1 # deleted dc, added rest + + # END handle chunk replacement + # END for each chunk + + if nfc == slen: + return False + # END handle completeness + return True + + +#} END structures + +#{ Routines + +def is_loose_object(m): + """ + :return: True the file contained in memory map m appears to be a loose object. + Only the first two bytes are needed""" + b0, b1 = map(ord, m[:2]) + word = (b0 << 8) + b1 + return b0 == 0x78 and (word % 31) == 0 + + +def loose_object_header_info(m): + """ + :return: tuple(type_string, uncompressed_size_in_bytes) the type string of the + object as well as its uncompressed size in bytes. + :param m: memory map from which to read the compressed object data""" + decompress_size = 8192 # is used in cgit as well + hdr = decompressobj().decompress(m, decompress_size) + type_name, size = hdr[:hdr.find(NULL_BYTE)].split(BYTE_SPACE) + + return type_name, int(size) + + +def pack_object_header_info(data): + """ + :return: tuple(type_id, uncompressed_size_in_bytes, byte_offset) + The type_id should be interpreted according to the ``type_id_to_type_map`` map + The byte-offset specifies the start of the actual zlib compressed datastream + :param m: random-access memory, like a string or memory map""" + c = byte_ord(data[0]) # first byte + i = 1 # next char to read + type_id = (c >> 4) & 7 # numeric type + size = c & 15 # starting size + s = 4 # starting bit-shift size + while c & 0x80: + c = byte_ord(data[i]) + i += 1 + size += (c & 0x7f) << s + s += 7 + # END character loop + # end performance at expense of maintenance ... + return (type_id, size, i) + + +def create_pack_object_header(obj_type, obj_size): + """ + :return: string defining the pack header comprised of the object type + and its incompressed size in bytes + + :param obj_type: pack type_id of the object + :param obj_size: uncompressed size in bytes of the following object stream""" + c = 0 # 1 byte + hdr = bytearray() # output string + + c = (obj_type << 4) | (obj_size & 0xf) + obj_size >>= 4 + while obj_size: + hdr.append(c | 0x80) + c = obj_size & 0x7f + obj_size >>= 7 + # END until size is consumed + hdr.append(c) + # end handle interpreter + return hdr + + +def msb_size(data, offset=0): + """ + :return: tuple(read_bytes, size) read the msb size from the given random + access data starting at the given byte offset""" + size = 0 + i = 0 + l = len(data) + hit_msb = False + while i < l: + c = data[i + offset] + size |= (c & 0x7f) << i * 7 + i += 1 + if not c & 0x80: + hit_msb = True + break + # END check msb bit + # END while in range + # end performance ... + if not hit_msb: + raise AssertionError("Could not find terminating MSB byte in data stream") + return i + offset, size + + +def loose_object_header(type, size): + """ + :return: bytes representing the loose object header, which is immediately + followed by the content stream of size 'size'""" + return ('%s %i\0' % (force_text(type), size)).encode('ascii') + + +def write_object(type, size, read, write, chunk_size=chunk_size): + """ + Write the object as identified by type, size and source_stream into the + target_stream + + :param type: type string of the object + :param size: amount of bytes to write from source_stream + :param read: read method of a stream providing the content data + :param write: write method of the output stream + :param close_target_stream: if True, the target stream will be closed when + the routine exits, even if an error is thrown + :return: The actual amount of bytes written to stream, which includes the header and a trailing newline""" + tbw = 0 # total num bytes written + + # WRITE HEADER: type SP size NULL + tbw += write(loose_object_header(type, size)) + tbw += stream_copy(read, write, size, chunk_size) + + return tbw + + +def stream_copy(read, write, size, chunk_size): + """ + Copy a stream up to size bytes using the provided read and write methods, + in chunks of chunk_size + + **Note:** its much like stream_copy utility, but operates just using methods""" + dbw = 0 # num data bytes written + + # WRITE ALL DATA UP TO SIZE + while True: + cs = min(chunk_size, size - dbw) + # NOTE: not all write methods return the amount of written bytes, like + # mmap.write. Its bad, but we just deal with it ... perhaps its not + # even less efficient + # data_len = write(read(cs)) + # dbw += data_len + data = read(cs) + data_len = len(data) + dbw += data_len + write(data) + if data_len < cs or dbw == size: + break + # END check for stream end + # END duplicate data + return dbw + + +def connect_deltas(dstreams): + """ + Read the condensed delta chunk information from dstream and merge its information + into a list of existing delta chunks + + :param dstreams: iterable of delta stream objects, the delta to be applied last + comes first, then all its ancestors in order + :return: DeltaChunkList, containing all operations to apply""" + tdcl = None # topmost dcl + + dcl = tdcl = TopdownDeltaChunkList() + for dsi, ds in enumerate(dstreams): + # print "Stream", dsi + db = ds.read() + delta_buf_size = ds.size + + # read header + i, base_size = msb_size(db) + i, target_size = msb_size(db, i) + + # interpret opcodes + tbw = 0 # amount of target bytes written + while i < delta_buf_size: + c = ord(db[i]) + i += 1 + if c & 0x80: + cp_off, cp_size = 0, 0 + if (c & 0x01): + cp_off = ord(db[i]) + i += 1 + if (c & 0x02): + cp_off |= (ord(db[i]) << 8) + i += 1 + if (c & 0x04): + cp_off |= (ord(db[i]) << 16) + i += 1 + if (c & 0x08): + cp_off |= (ord(db[i]) << 24) + i += 1 + if (c & 0x10): + cp_size = ord(db[i]) + i += 1 + if (c & 0x20): + cp_size |= (ord(db[i]) << 8) + i += 1 + if (c & 0x40): + cp_size |= (ord(db[i]) << 16) + i += 1 + + if not cp_size: + cp_size = 0x10000 + + rbound = cp_off + cp_size + if (rbound < cp_size or + rbound > base_size): + break + + dcl.append(DeltaChunk(tbw, cp_size, cp_off, None)) + tbw += cp_size + elif c: + # NOTE: in C, the data chunks should probably be concatenated here. + # In python, we do it as a post-process + dcl.append(DeltaChunk(tbw, c, 0, db[i:i + c])) + i += c + tbw += c + else: + raise ValueError("unexpected delta opcode 0") + # END handle command byte + # END while processing delta data + + dcl.compress() + + # merge the lists ! + if dsi > 0: + if not tdcl.connect_with_next_base(dcl): + break + # END handle merge + + # prepare next base + dcl = DeltaChunkList() + # END for each delta stream + + return tdcl + + +def apply_delta_data(src_buf, src_buf_size, delta_buf, delta_buf_size, write): + """ + Apply data from a delta buffer using a source buffer to the target file + + :param src_buf: random access data from which the delta was created + :param src_buf_size: size of the source buffer in bytes + :param delta_buf_size: size for the delta buffer in bytes + :param delta_buf: random access delta data + :param write: write method taking a chunk of bytes + + **Note:** transcribed to python from the similar routine in patch-delta.c""" + i = 0 + db = delta_buf + while i < delta_buf_size: + c = db[i] + i += 1 + if c & 0x80: + cp_off, cp_size = 0, 0 + if (c & 0x01): + cp_off = db[i] + i += 1 + if (c & 0x02): + cp_off |= (db[i] << 8) + i += 1 + if (c & 0x04): + cp_off |= (db[i] << 16) + i += 1 + if (c & 0x08): + cp_off |= (db[i] << 24) + i += 1 + if (c & 0x10): + cp_size = db[i] + i += 1 + if (c & 0x20): + cp_size |= (db[i] << 8) + i += 1 + if (c & 0x40): + cp_size |= (db[i] << 16) + i += 1 + + if not cp_size: + cp_size = 0x10000 + + rbound = cp_off + cp_size + if (rbound < cp_size or + rbound > src_buf_size): + break + write(src_buf[cp_off:cp_off + cp_size]) + elif c: + write(db[i:i + c]) + i += c + else: + raise ValueError("unexpected delta opcode 0") + # END handle command byte + # END while processing delta data + + # yes, lets use the exact same error message that git uses :) + assert i == delta_buf_size, "delta replay has gone wild" + + +def is_equal_canonical_sha(canonical_length, match, sha1): + """ + :return: True if the given lhs and rhs 20 byte binary shas + The comparison will take the canonical_length of the match sha into account, + hence the comparison will only use the last 4 bytes for uneven canonical representations + :param match: less than 20 byte sha + :param sha1: 20 byte sha""" + binary_length = canonical_length // 2 + if match[:binary_length] != sha1[:binary_length]: + return False + + if canonical_length - binary_length and \ + (byte_ord(match[-1]) ^ byte_ord(sha1[len(match) - 1])) & 0xf0: + return False + # END handle uneven canonnical length + return True + +#} END routines + + +try: + from gitdb_speedups._perf import connect_deltas +except ImportError: + pass diff --git a/lib/python3.12/site-packages/gitdb/pack.py b/lib/python3.12/site-packages/gitdb/pack.py new file mode 100644 index 0000000000000000000000000000000000000000..e559e113d4d31135fa72c39d4b5f189997ab211c --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/pack.py @@ -0,0 +1,1031 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Contains PackIndexFile and PackFile implementations""" +import zlib + +from gitdb.exc import ( + BadObject, + AmbiguousObjectName, + UnsupportedOperation, + ParseError +) + +from gitdb.util import ( + mman, + LazyMixin, + unpack_from, + bin_to_hex, + byte_ord, +) + +from gitdb.fun import ( + create_pack_object_header, + pack_object_header_info, + is_equal_canonical_sha, + type_id_to_type_map, + write_object, + stream_copy, + chunk_size, + delta_types, + OFS_DELTA, + REF_DELTA, + msb_size +) + +try: + from gitdb_speedups._perf import PackIndexFile_sha_to_index +except ImportError: + pass +# END try c module + +from gitdb.base import ( # Amazing ! + OInfo, + OStream, + OPackInfo, + OPackStream, + ODeltaStream, + ODeltaPackInfo, + ODeltaPackStream, +) + +from gitdb.stream import ( + DecompressMemMapReader, + DeltaApplyReader, + Sha1Writer, + NullStream, + FlexibleSha1Writer +) + +from struct import pack +from binascii import crc32 + +from gitdb.const import NULL_BYTE + +import tempfile +import array +import os +import sys + +__all__ = ('PackIndexFile', 'PackFile', 'PackEntity') + + +#{ Utilities + +def pack_object_at(cursor, offset, as_stream): + """ + :return: Tuple(abs_data_offset, PackInfo|PackStream) + an object of the correct type according to the type_id of the object. + If as_stream is True, the object will contain a stream, allowing the + data to be read decompressed. + :param data: random accessible data containing all required information + :parma offset: offset in to the data at which the object information is located + :param as_stream: if True, a stream object will be returned that can read + the data, otherwise you receive an info object only""" + data = cursor.use_region(offset).buffer() + type_id, uncomp_size, data_rela_offset = pack_object_header_info(data) + total_rela_offset = None # set later, actual offset until data stream begins + delta_info = None + + # OFFSET DELTA + if type_id == OFS_DELTA: + i = data_rela_offset + c = byte_ord(data[i]) + i += 1 + delta_offset = c & 0x7f + while c & 0x80: + c = byte_ord(data[i]) + i += 1 + delta_offset += 1 + delta_offset = (delta_offset << 7) + (c & 0x7f) + # END character loop + delta_info = delta_offset + total_rela_offset = i + # REF DELTA + elif type_id == REF_DELTA: + total_rela_offset = data_rela_offset + 20 + delta_info = data[data_rela_offset:total_rela_offset] + # BASE OBJECT + else: + # assume its a base object + total_rela_offset = data_rela_offset + # END handle type id + abs_data_offset = offset + total_rela_offset + if as_stream: + stream = DecompressMemMapReader(data[total_rela_offset:], False, uncomp_size) + if delta_info is None: + return abs_data_offset, OPackStream(offset, type_id, uncomp_size, stream) + else: + return abs_data_offset, ODeltaPackStream(offset, type_id, uncomp_size, delta_info, stream) + else: + if delta_info is None: + return abs_data_offset, OPackInfo(offset, type_id, uncomp_size) + else: + return abs_data_offset, ODeltaPackInfo(offset, type_id, uncomp_size, delta_info) + # END handle info + # END handle stream + + +def write_stream_to_pack(read, write, zstream, base_crc=None): + """Copy a stream as read from read function, zip it, and write the result. + Count the number of written bytes and return it + :param base_crc: if not None, the crc will be the base for all compressed data + we consecutively write and generate a crc32 from. If None, no crc will be generated + :return: tuple(no bytes read, no bytes written, crc32) crc might be 0 if base_crc + was false""" + br = 0 # bytes read + bw = 0 # bytes written + want_crc = base_crc is not None + crc = 0 + if want_crc: + crc = base_crc + # END initialize crc + + while True: + chunk = read(chunk_size) + br += len(chunk) + compressed = zstream.compress(chunk) + bw += len(compressed) + write(compressed) # cannot assume return value + + if want_crc: + crc = crc32(compressed, crc) + # END handle crc + + if len(chunk) != chunk_size: + break + # END copy loop + + compressed = zstream.flush() + bw += len(compressed) + write(compressed) + if want_crc: + crc = crc32(compressed, crc) + # END handle crc + + return (br, bw, crc) + + +#} END utilities + + +class IndexWriter: + + """Utility to cache index information, allowing to write all information later + in one go to the given stream + **Note:** currently only writes v2 indices""" + __slots__ = '_objs' + + def __init__(self): + self._objs = list() + + def append(self, binsha, crc, offset): + """Append one piece of object information""" + self._objs.append((binsha, crc, offset)) + + def write(self, pack_sha, write): + """Write the index file using the given write method + :param pack_sha: binary sha over the whole pack that we index + :return: sha1 binary sha over all index file contents""" + # sort for sha1 hash + self._objs.sort(key=lambda o: o[0]) + + sha_writer = FlexibleSha1Writer(write) + sha_write = sha_writer.write + sha_write(PackIndexFile.index_v2_signature) + sha_write(pack(">L", PackIndexFile.index_version_default)) + + # fanout + tmplist = list((0,) * 256) # fanout or list with 64 bit offsets + for t in self._objs: + tmplist[byte_ord(t[0][0])] += 1 + # END prepare fanout + for i in range(255): + v = tmplist[i] + sha_write(pack('>L', v)) + tmplist[i + 1] += v + # END write each fanout entry + sha_write(pack('>L', tmplist[255])) + + # sha1 ordered + # save calls, that is push them into c + sha_write(b''.join(t[0] for t in self._objs)) + + # crc32 + for t in self._objs: + sha_write(pack('>L', t[1] & 0xffffffff)) + # END for each crc + + tmplist = list() + # offset 32 + for t in self._objs: + ofs = t[2] + if ofs > 0x7fffffff: + tmplist.append(ofs) + ofs = 0x80000000 + len(tmplist) - 1 + # END handle 64 bit offsets + sha_write(pack('>L', ofs & 0xffffffff)) + # END for each offset + + # offset 64 + for ofs in tmplist: + sha_write(pack(">Q", ofs)) + # END for each offset + + # trailer + assert(len(pack_sha) == 20) + sha_write(pack_sha) + sha = sha_writer.sha(as_hex=False) + write(sha) + return sha + + +class PackIndexFile(LazyMixin): + + """A pack index provides offsets into the corresponding pack, allowing to find + locations for offsets faster.""" + + # Dont use slots as we dynamically bind functions for each version, need a dict for this + # The slots you see here are just to keep track of our instance variables + # __slots__ = ('_indexpath', '_fanout_table', '_cursor', '_version', + # '_sha_list_offset', '_crc_list_offset', '_pack_offset', '_pack_64_offset') + + # used in v2 indices + _sha_list_offset = 8 + 1024 + index_v2_signature = b'\xfftOc' + index_version_default = 2 + + def __init__(self, indexpath): + super().__init__() + self._indexpath = indexpath + + def close(self): + mman.force_map_handle_removal_win(self._indexpath) + self._cursor = None + + def _set_cache_(self, attr): + if attr == "_packfile_checksum": + self._packfile_checksum = self._cursor.map()[-40:-20] + elif attr == "_packfile_checksum": + self._packfile_checksum = self._cursor.map()[-20:] + elif attr == "_cursor": + # Note: We don't lock the file when reading as we cannot be sure + # that we can actually write to the location - it could be a read-only + # alternate for instance + self._cursor = mman.make_cursor(self._indexpath).use_region() + # We will assume that the index will always fully fit into memory ! + if mman.window_size() > 0 and self._cursor.file_size() > mman.window_size(): + raise AssertionError("The index file at %s is too large to fit into a mapped window (%i > %i). This is a limitation of the implementation" % ( + self._indexpath, self._cursor.file_size(), mman.window_size())) + # END assert window size + else: + # now its time to initialize everything - if we are here, someone wants + # to access the fanout table or related properties + + # CHECK VERSION + mmap = self._cursor.map() + self._version = (mmap[:4] == self.index_v2_signature and 2) or 1 + if self._version == 2: + version_id = unpack_from(">L", mmap, 4)[0] + assert version_id == self._version, "Unsupported index version: %i" % version_id + # END assert version + + # SETUP FUNCTIONS + # setup our functions according to the actual version + for fname in ('entry', 'offset', 'sha', 'crc'): + setattr(self, fname, getattr(self, "_%s_v%i" % (fname, self._version))) + # END for each function to initialize + + # INITIALIZE DATA + # byte offset is 8 if version is 2, 0 otherwise + self._initialize() + # END handle attributes + + #{ Access V1 + + def _entry_v1(self, i): + """:return: tuple(offset, binsha, 0)""" + return unpack_from(">L20s", self._cursor.map(), 1024 + i * 24) + (0, ) + + def _offset_v1(self, i): + """see ``_offset_v2``""" + return unpack_from(">L", self._cursor.map(), 1024 + i * 24)[0] + + def _sha_v1(self, i): + """see ``_sha_v2``""" + base = 1024 + (i * 24) + 4 + return self._cursor.map()[base:base + 20] + + def _crc_v1(self, i): + """unsupported""" + return 0 + + #} END access V1 + + #{ Access V2 + def _entry_v2(self, i): + """:return: tuple(offset, binsha, crc)""" + return (self._offset_v2(i), self._sha_v2(i), self._crc_v2(i)) + + def _offset_v2(self, i): + """:return: 32 or 64 byte offset into pack files. 64 byte offsets will only + be returned if the pack is larger than 4 GiB, or 2^32""" + offset = unpack_from(">L", self._cursor.map(), self._pack_offset + i * 4)[0] + + # if the high-bit is set, this indicates that we have to lookup the offset + # in the 64 bit region of the file. The current offset ( lower 31 bits ) + # are the index into it + if offset & 0x80000000: + offset = unpack_from(">Q", self._cursor.map(), self._pack_64_offset + (offset & ~0x80000000) * 8)[0] + # END handle 64 bit offset + + return offset + + def _sha_v2(self, i): + """:return: sha at the given index of this file index instance""" + base = self._sha_list_offset + i * 20 + return self._cursor.map()[base:base + 20] + + def _crc_v2(self, i): + """:return: 4 bytes crc for the object at index i""" + return unpack_from(">L", self._cursor.map(), self._crc_list_offset + i * 4)[0] + + #} END access V2 + + #{ Initialization + + def _initialize(self): + """initialize base data""" + self._fanout_table = self._read_fanout((self._version == 2) * 8) + + if self._version == 2: + self._crc_list_offset = self._sha_list_offset + self.size() * 20 + self._pack_offset = self._crc_list_offset + self.size() * 4 + self._pack_64_offset = self._pack_offset + self.size() * 4 + # END setup base + + def _read_fanout(self, byte_offset): + """Generate a fanout table from our data""" + d = self._cursor.map() + out = list() + append = out.append + for i in range(256): + append(unpack_from('>L', d, byte_offset + i * 4)[0]) + # END for each entry + return out + + #} END initialization + + #{ Properties + def version(self): + return self._version + + def size(self): + """:return: amount of objects referred to by this index""" + return self._fanout_table[255] + + def path(self): + """:return: path to the packindexfile""" + return self._indexpath + + def packfile_checksum(self): + """:return: 20 byte sha representing the sha1 hash of the pack file""" + return self._cursor.map()[-40:-20] + + def indexfile_checksum(self): + """:return: 20 byte sha representing the sha1 hash of this index file""" + return self._cursor.map()[-20:] + + def offsets(self): + """:return: sequence of all offsets in the order in which they were written + + **Note:** return value can be random accessed, but may be immmutable""" + if self._version == 2: + # read stream to array, convert to tuple + a = array.array('I') # 4 byte unsigned int, long are 8 byte on 64 bit it appears + a.frombytes(self._cursor.map()[self._pack_offset:self._pack_64_offset]) + + # networkbyteorder to something array likes more + if sys.byteorder == 'little': + a.byteswap() + return a + else: + return tuple(self.offset(index) for index in range(self.size())) + # END handle version + + def sha_to_index(self, sha): + """ + :return: index usable with the ``offset`` or ``entry`` method, or None + if the sha was not found in this pack index + :param sha: 20 byte sha to lookup""" + first_byte = byte_ord(sha[0]) + get_sha = self.sha + lo = 0 # lower index, the left bound of the bisection + if first_byte != 0: + lo = self._fanout_table[first_byte - 1] + hi = self._fanout_table[first_byte] # the upper, right bound of the bisection + + # bisect until we have the sha + while lo < hi: + mid = (lo + hi) // 2 + mid_sha = get_sha(mid) + if sha < mid_sha: + hi = mid + elif sha == mid_sha: + return mid + else: + lo = mid + 1 + # END handle midpoint + # END bisect + return None + + def partial_sha_to_index(self, partial_bin_sha, canonical_length): + """ + :return: index as in `sha_to_index` or None if the sha was not found in this + index file + :param partial_bin_sha: an at least two bytes of a partial binary sha as bytes + :param canonical_length: length of the original hexadecimal representation of the + given partial binary sha + :raise AmbiguousObjectName:""" + if len(partial_bin_sha) < 2: + raise ValueError("Require at least 2 bytes of partial sha") + + assert isinstance(partial_bin_sha, bytes), "partial_bin_sha must be bytes" + first_byte = byte_ord(partial_bin_sha[0]) + + get_sha = self.sha + lo = 0 # lower index, the left bound of the bisection + if first_byte != 0: + lo = self._fanout_table[first_byte - 1] + hi = self._fanout_table[first_byte] # the upper, right bound of the bisection + + # fill the partial to full 20 bytes + filled_sha = partial_bin_sha + NULL_BYTE * (20 - len(partial_bin_sha)) + + # find lowest + while lo < hi: + mid = (lo + hi) // 2 + mid_sha = get_sha(mid) + if filled_sha < mid_sha: + hi = mid + elif filled_sha == mid_sha: + # perfect match + lo = mid + break + else: + lo = mid + 1 + # END handle midpoint + # END bisect + + if lo < self.size(): + cur_sha = get_sha(lo) + if is_equal_canonical_sha(canonical_length, partial_bin_sha, cur_sha): + next_sha = None + if lo + 1 < self.size(): + next_sha = get_sha(lo + 1) + if next_sha and next_sha == cur_sha: + raise AmbiguousObjectName(partial_bin_sha) + return lo + # END if we have a match + # END if we found something + return None + + if 'PackIndexFile_sha_to_index' in globals(): + # NOTE: Its just about 25% faster, the major bottleneck might be the attr + # accesses + def sha_to_index(self, sha): + return PackIndexFile_sha_to_index(self, sha) + # END redefine heavy-hitter with c version + + #} END properties + + +class PackFile(LazyMixin): + + """A pack is a file written according to the Version 2 for git packs + + As we currently use memory maps, it could be assumed that the maximum size of + packs therefore is 32 bit on 32 bit systems. On 64 bit systems, this should be + fine though. + + **Note:** at some point, this might be implemented using streams as well, or + streams are an alternate path in the case memory maps cannot be created + for some reason - one clearly doesn't want to read 10GB at once in that + case""" + + __slots__ = ('_packpath', '_cursor', '_size', '_version') + pack_signature = 0x5041434b # 'PACK' + pack_version_default = 2 + + # offset into our data at which the first object starts + first_object_offset = 3 * 4 # header bytes + footer_size = 20 # final sha + + def __init__(self, packpath): + self._packpath = packpath + + def close(self): + mman.force_map_handle_removal_win(self._packpath) + self._cursor = None + + def _set_cache_(self, attr): + # we fill the whole cache, whichever attribute gets queried first + self._cursor = mman.make_cursor(self._packpath).use_region() + + # read the header information + type_id, self._version, self._size = unpack_from(">LLL", self._cursor.map(), 0) + + # TODO: figure out whether we should better keep the lock, or maybe + # add a .keep file instead ? + if type_id != self.pack_signature: + raise ParseError("Invalid pack signature: %i" % type_id) + + def _iter_objects(self, start_offset, as_stream=True): + """Handle the actual iteration of objects within this pack""" + c = self._cursor + content_size = c.file_size() - self.footer_size + cur_offset = start_offset or self.first_object_offset + + null = NullStream() + while cur_offset < content_size: + data_offset, ostream = pack_object_at(c, cur_offset, True) + # scrub the stream to the end - this decompresses the object, but yields + # the amount of compressed bytes we need to get to the next offset + + stream_copy(ostream.read, null.write, ostream.size, chunk_size) + assert ostream.stream._br == ostream.size + cur_offset += (data_offset - ostream.pack_offset) + ostream.stream.compressed_bytes_read() + + # if a stream is requested, reset it beforehand + # Otherwise return the Stream object directly, its derived from the + # info object + if as_stream: + ostream.stream.seek(0) + yield ostream + # END until we have read everything + + #{ Pack Information + + def size(self): + """:return: The amount of objects stored in this pack""" + return self._size + + def version(self): + """:return: the version of this pack""" + return self._version + + def data(self): + """ + :return: read-only data of this pack. It provides random access and usually + is a memory map. + :note: This method is unsafe as it returns a window into a file which might be larger than than the actual window size""" + # can use map as we are starting at offset 0. Otherwise we would have to use buffer() + return self._cursor.use_region().map() + + def checksum(self): + """:return: 20 byte sha1 hash on all object sha's contained in this file""" + return self._cursor.use_region(self._cursor.file_size() - 20).buffer()[:] + + def path(self): + """:return: path to the packfile""" + return self._packpath + #} END pack information + + #{ Pack Specific + + def collect_streams(self, offset): + """ + :return: list of pack streams which are required to build the object + at the given offset. The first entry of the list is the object at offset, + the last one is either a full object, or a REF_Delta stream. The latter + type needs its reference object to be locked up in an ODB to form a valid + delta chain. + If the object at offset is no delta, the size of the list is 1. + :param offset: specifies the first byte of the object within this pack""" + out = list() + c = self._cursor + while True: + ostream = pack_object_at(c, offset, True)[1] + out.append(ostream) + if ostream.type_id == OFS_DELTA: + offset = ostream.pack_offset - ostream.delta_info + else: + # the only thing we can lookup are OFFSET deltas. Everything + # else is either an object, or a ref delta, in the latter + # case someone else has to find it + break + # END handle type + # END while chaining streams + return out + + #} END pack specific + + #{ Read-Database like Interface + + def info(self, offset): + """Retrieve information about the object at the given file-absolute offset + + :param offset: byte offset + :return: OPackInfo instance, the actual type differs depending on the type_id attribute""" + return pack_object_at(self._cursor, offset or self.first_object_offset, False)[1] + + def stream(self, offset): + """Retrieve an object at the given file-relative offset as stream along with its information + + :param offset: byte offset + :return: OPackStream instance, the actual type differs depending on the type_id attribute""" + return pack_object_at(self._cursor, offset or self.first_object_offset, True)[1] + + def stream_iter(self, start_offset=0): + """ + :return: iterator yielding OPackStream compatible instances, allowing + to access the data in the pack directly. + :param start_offset: offset to the first object to iterate. If 0, iteration + starts at the very first object in the pack. + + **Note:** Iterating a pack directly is costly as the datastream has to be decompressed + to determine the bounds between the objects""" + return self._iter_objects(start_offset, as_stream=True) + + #} END Read-Database like Interface + + +class PackEntity(LazyMixin): + + """Combines the PackIndexFile and the PackFile into one, allowing the + actual objects to be resolved and iterated""" + + __slots__ = ('_index', # our index file + '_pack', # our pack file + '_offset_map' # on demand dict mapping one offset to the next consecutive one + ) + + IndexFileCls = PackIndexFile + PackFileCls = PackFile + + def __init__(self, pack_or_index_path): + """Initialize ourselves with the path to the respective pack or index file""" + basename, ext = os.path.splitext(pack_or_index_path) + self._index = self.IndexFileCls("%s.idx" % basename) # PackIndexFile instance + self._pack = self.PackFileCls("%s.pack" % basename) # corresponding PackFile instance + + def close(self): + self._index.close() + self._pack.close() + + def _set_cache_(self, attr): + # currently this can only be _offset_map + # TODO: make this a simple sorted offset array which can be bisected + # to find the respective entry, from which we can take a +1 easily + # This might be slower, but should also be much lighter in memory ! + offsets_sorted = sorted(self._index.offsets()) + last_offset = len(self._pack.data()) - self._pack.footer_size + assert offsets_sorted, "Cannot handle empty indices" + + offset_map = None + if len(offsets_sorted) == 1: + offset_map = {offsets_sorted[0]: last_offset} + else: + iter_offsets = iter(offsets_sorted) + iter_offsets_plus_one = iter(offsets_sorted) + next(iter_offsets_plus_one) + consecutive = zip(iter_offsets, iter_offsets_plus_one) + + offset_map = dict(consecutive) + + # the last offset is not yet set + offset_map[offsets_sorted[-1]] = last_offset + # END handle offset amount + self._offset_map = offset_map + + def _sha_to_index(self, sha): + """:return: index for the given sha, or raise""" + index = self._index.sha_to_index(sha) + if index is None: + raise BadObject(sha) + return index + + def _iter_objects(self, as_stream): + """Iterate over all objects in our index and yield their OInfo or OStream instences""" + _sha = self._index.sha + _object = self._object + for index in range(self._index.size()): + yield _object(_sha(index), as_stream, index) + # END for each index + + def _object(self, sha, as_stream, index=-1): + """:return: OInfo or OStream object providing information about the given sha + :param index: if not -1, its assumed to be the sha's index in the IndexFile""" + # its a little bit redundant here, but it needs to be efficient + if index < 0: + index = self._sha_to_index(sha) + if sha is None: + sha = self._index.sha(index) + # END assure sha is present ( in output ) + offset = self._index.offset(index) + type_id, uncomp_size, data_rela_offset = pack_object_header_info(self._pack._cursor.use_region(offset).buffer()) + if as_stream: + if type_id not in delta_types: + packstream = self._pack.stream(offset) + return OStream(sha, packstream.type, packstream.size, packstream.stream) + # END handle non-deltas + + # produce a delta stream containing all info + # To prevent it from applying the deltas when querying the size, + # we extract it from the delta stream ourselves + streams = self.collect_streams_at_offset(offset) + dstream = DeltaApplyReader.new(streams) + + return ODeltaStream(sha, dstream.type, None, dstream) + else: + if type_id not in delta_types: + return OInfo(sha, type_id_to_type_map[type_id], uncomp_size) + # END handle non-deltas + + # deltas are a little tougher - unpack the first bytes to obtain + # the actual target size, as opposed to the size of the delta data + streams = self.collect_streams_at_offset(offset) + buf = streams[0].read(512) + offset, src_size = msb_size(buf) + offset, target_size = msb_size(buf, offset) + + # collect the streams to obtain the actual object type + if streams[-1].type_id in delta_types: + raise BadObject(sha, "Could not resolve delta object") + return OInfo(sha, streams[-1].type, target_size) + # END handle stream + + #{ Read-Database like Interface + + def info(self, sha): + """Retrieve information about the object identified by the given sha + + :param sha: 20 byte sha1 + :raise BadObject: + :return: OInfo instance, with 20 byte sha""" + return self._object(sha, False) + + def stream(self, sha): + """Retrieve an object stream along with its information as identified by the given sha + + :param sha: 20 byte sha1 + :raise BadObject: + :return: OStream instance, with 20 byte sha""" + return self._object(sha, True) + + def info_at_index(self, index): + """As ``info``, but uses a PackIndexFile compatible index to refer to the object""" + return self._object(None, False, index) + + def stream_at_index(self, index): + """As ``stream``, but uses a PackIndexFile compatible index to refer to the + object""" + return self._object(None, True, index) + + #} END Read-Database like Interface + + #{ Interface + + def pack(self): + """:return: the underlying pack file instance""" + return self._pack + + def index(self): + """:return: the underlying pack index file instance""" + return self._index + + def is_valid_stream(self, sha, use_crc=False): + """ + Verify that the stream at the given sha is valid. + + :param use_crc: if True, the index' crc is run over the compressed stream of + the object, which is much faster than checking the sha1. It is also + more prone to unnoticed corruption or manipulation. + :param sha: 20 byte sha1 of the object whose stream to verify + whether the compressed stream of the object is valid. If it is + a delta, this only verifies that the delta's data is valid, not the + data of the actual undeltified object, as it depends on more than + just this stream. + If False, the object will be decompressed and the sha generated. It must + match the given sha + + :return: True if the stream is valid + :raise UnsupportedOperation: If the index is version 1 only + :raise BadObject: sha was not found""" + if use_crc: + if self._index.version() < 2: + raise UnsupportedOperation("Version 1 indices do not contain crc's, verify by sha instead") + # END handle index version + + index = self._sha_to_index(sha) + offset = self._index.offset(index) + next_offset = self._offset_map[offset] + crc_value = self._index.crc(index) + + # create the current crc value, on the compressed object data + # Read it in chunks, without copying the data + crc_update = zlib.crc32 + pack_data = self._pack.data() + cur_pos = offset + this_crc_value = 0 + while cur_pos < next_offset: + rbound = min(cur_pos + chunk_size, next_offset) + size = rbound - cur_pos + this_crc_value = crc_update(pack_data[cur_pos:cur_pos + size], this_crc_value) + cur_pos += size + # END window size loop + + # crc returns signed 32 bit numbers, the AND op forces it into unsigned + # mode ... wow, sneaky, from dulwich. + return (this_crc_value & 0xffffffff) == crc_value + else: + shawriter = Sha1Writer() + stream = self._object(sha, as_stream=True) + # write a loose object, which is the basis for the sha + write_object(stream.type, stream.size, stream.read, shawriter.write) + + assert shawriter.sha(as_hex=False) == sha + return shawriter.sha(as_hex=False) == sha + # END handle crc/sha verification + return True + + def info_iter(self): + """ + :return: Iterator over all objects in this pack. The iterator yields + OInfo instances""" + return self._iter_objects(as_stream=False) + + def stream_iter(self): + """ + :return: iterator over all objects in this pack. The iterator yields + OStream instances""" + return self._iter_objects(as_stream=True) + + def collect_streams_at_offset(self, offset): + """ + As the version in the PackFile, but can resolve REF deltas within this pack + For more info, see ``collect_streams`` + + :param offset: offset into the pack file at which the object can be found""" + streams = self._pack.collect_streams(offset) + + # try to resolve the last one if needed. It is assumed to be either + # a REF delta, or a base object, as OFFSET deltas are resolved by the pack + if streams[-1].type_id == REF_DELTA: + stream = streams[-1] + while stream.type_id in delta_types: + if stream.type_id == REF_DELTA: + # smmap can return memory view objects, which can't be compared as buffers/bytes can ... + if isinstance(stream.delta_info, memoryview): + sindex = self._index.sha_to_index(stream.delta_info.tobytes()) + else: + sindex = self._index.sha_to_index(stream.delta_info) + if sindex is None: + break + stream = self._pack.stream(self._index.offset(sindex)) + streams.append(stream) + else: + # must be another OFS DELTA - this could happen if a REF + # delta we resolve previously points to an OFS delta. Who + # would do that ;) ? We can handle it though + stream = self._pack.stream(stream.delta_info) + streams.append(stream) + # END handle ref delta + # END resolve ref streams + # END resolve streams + + return streams + + def collect_streams(self, sha): + """ + As ``PackFile.collect_streams``, but takes a sha instead of an offset. + Additionally, ref_delta streams will be resolved within this pack. + If this is not possible, the stream will be left alone, hence it is adivsed + to check for unresolved ref-deltas and resolve them before attempting to + construct a delta stream. + + :param sha: 20 byte sha1 specifying the object whose related streams you want to collect + :return: list of streams, first being the actual object delta, the last being + a possibly unresolved base object. + :raise BadObject:""" + return self.collect_streams_at_offset(self._index.offset(self._sha_to_index(sha))) + + @classmethod + def write_pack(cls, object_iter, pack_write, index_write=None, + object_count=None, zlib_compression=zlib.Z_BEST_SPEED): + """ + Create a new pack by putting all objects obtained by the object_iterator + into a pack which is written using the pack_write method. + The respective index is produced as well if index_write is not Non. + + :param object_iter: iterator yielding odb output objects + :param pack_write: function to receive strings to write into the pack stream + :param indx_write: if not None, the function writes the index file corresponding + to the pack. + :param object_count: if you can provide the amount of objects in your iteration, + this would be the place to put it. Otherwise we have to pre-iterate and store + all items into a list to get the number, which uses more memory than necessary. + :param zlib_compression: the zlib compression level to use + :return: tuple(pack_sha, index_binsha) binary sha over all the contents of the pack + and over all contents of the index. If index_write was None, index_binsha will be None + + **Note:** The destination of the write functions is up to the user. It could + be a socket, or a file for instance + + **Note:** writes only undeltified objects""" + objs = object_iter + if not object_count: + if not isinstance(object_iter, (tuple, list)): + objs = list(object_iter) + # END handle list type + object_count = len(objs) + # END handle object + + pack_writer = FlexibleSha1Writer(pack_write) + pwrite = pack_writer.write + ofs = 0 # current offset into the pack file + index = None + wants_index = index_write is not None + + # write header + pwrite(pack('>LLL', PackFile.pack_signature, PackFile.pack_version_default, object_count)) + ofs += 12 + + if wants_index: + index = IndexWriter() + # END handle index header + + actual_count = 0 + for obj in objs: + actual_count += 1 + crc = 0 + + # object header + hdr = create_pack_object_header(obj.type_id, obj.size) + if index_write: + crc = crc32(hdr) + else: + crc = None + # END handle crc + pwrite(hdr) + + # data stream + zstream = zlib.compressobj(zlib_compression) + ostream = obj.stream + br, bw, crc = write_stream_to_pack(ostream.read, pwrite, zstream, base_crc=crc) + assert(br == obj.size) + if wants_index: + index.append(obj.binsha, crc, ofs) + # END handle index + + ofs += len(hdr) + bw + if actual_count == object_count: + break + # END abort once we are done + # END for each object + + if actual_count != object_count: + raise ValueError( + "Expected to write %i objects into pack, but received only %i from iterators" % (object_count, actual_count)) + # END count assertion + + # write footer + pack_sha = pack_writer.sha(as_hex=False) + assert len(pack_sha) == 20 + pack_write(pack_sha) + ofs += len(pack_sha) # just for completeness ;) + + index_sha = None + if wants_index: + index_sha = index.write(pack_sha, index_write) + # END handle index + + return pack_sha, index_sha + + @classmethod + def create(cls, object_iter, base_dir, object_count=None, zlib_compression=zlib.Z_BEST_SPEED): + """Create a new on-disk entity comprised of a properly named pack file and a properly named + and corresponding index file. The pack contains all OStream objects contained in object iter. + :param base_dir: directory which is to contain the files + :return: PackEntity instance initialized with the new pack + + **Note:** for more information on the other parameters see the write_pack method""" + pack_fd, pack_path = tempfile.mkstemp('', 'pack', base_dir) + index_fd, index_path = tempfile.mkstemp('', 'index', base_dir) + pack_write = lambda d: os.write(pack_fd, d) + index_write = lambda d: os.write(index_fd, d) + + pack_binsha, index_binsha = cls.write_pack(object_iter, pack_write, index_write, object_count, zlib_compression) + os.close(pack_fd) + os.close(index_fd) + + fmt = "pack-%s.%s" + new_pack_path = os.path.join(base_dir, fmt % (bin_to_hex(pack_binsha), 'pack')) + new_index_path = os.path.join(base_dir, fmt % (bin_to_hex(pack_binsha), 'idx')) + os.rename(pack_path, new_pack_path) + os.rename(index_path, new_index_path) + + return cls(new_pack_path) + + #} END interface diff --git a/lib/python3.12/site-packages/gitdb/stream.py b/lib/python3.12/site-packages/gitdb/stream.py new file mode 100644 index 0000000000000000000000000000000000000000..1e0be84fd10f7173bb3dccea7a59cc6355909ffd --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/stream.py @@ -0,0 +1,730 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ + +from io import BytesIO + +import mmap +import os +import sys +import zlib + +from gitdb.fun import ( + msb_size, + stream_copy, + apply_delta_data, + connect_deltas, + delta_types +) + +from gitdb.util import ( + allocate_memory, + LazyMixin, + make_sha, + write, + close, +) + +from gitdb.const import NULL_BYTE, BYTE_SPACE +from gitdb.utils.encoding import force_bytes + +has_perf_mod = False +try: + from gitdb_speedups._perf import apply_delta as c_apply_delta + has_perf_mod = True +except ImportError: + pass + +__all__ = ('DecompressMemMapReader', 'FDCompressedSha1Writer', 'DeltaApplyReader', + 'Sha1Writer', 'FlexibleSha1Writer', 'ZippedStoreShaWriter', 'FDCompressedSha1Writer', + 'FDStream', 'NullStream') + + +#{ RO Streams + +class DecompressMemMapReader(LazyMixin): + + """Reads data in chunks from a memory map and decompresses it. The client sees + only the uncompressed data, respective file-like read calls are handling on-demand + buffered decompression accordingly + + A constraint on the total size of bytes is activated, simulating + a logical file within a possibly larger physical memory area + + To read efficiently, you clearly don't want to read individual bytes, instead, + read a few kilobytes at least. + + **Note:** The chunk-size should be carefully selected as it will involve quite a bit + of string copying due to the way the zlib is implemented. Its very wasteful, + hence we try to find a good tradeoff between allocation time and number of + times we actually allocate. An own zlib implementation would be good here + to better support streamed reading - it would only need to keep the mmap + and decompress it into chunks, that's all ... """ + __slots__ = ('_m', '_zip', '_buf', '_buflen', '_br', '_cws', '_cwe', '_s', '_close', + '_cbr', '_phi') + + max_read_size = 512 * 1024 # currently unused + + def __init__(self, m, close_on_deletion, size=None): + """Initialize with mmap for stream reading + :param m: must be content data - use new if you have object data and no size""" + self._m = m + self._zip = zlib.decompressobj() + self._buf = None # buffer of decompressed bytes + self._buflen = 0 # length of bytes in buffer + if size is not None: + self._s = size # size of uncompressed data to read in total + self._br = 0 # num uncompressed bytes read + self._cws = 0 # start byte of compression window + self._cwe = 0 # end byte of compression window + self._cbr = 0 # number of compressed bytes read + self._phi = False # is True if we parsed the header info + self._close = close_on_deletion # close the memmap on deletion ? + + def _set_cache_(self, attr): + assert attr == '_s' + # only happens for size, which is a marker to indicate we still + # have to parse the header from the stream + self._parse_header_info() + + def __del__(self): + self.close() + + def _parse_header_info(self): + """If this stream contains object data, parse the header info and skip the + stream to a point where each read will yield object content + + :return: parsed type_string, size""" + # read header + # should really be enough, cgit uses 8192 I believe + # And for good reason !! This needs to be that high for the header to be read correctly in all cases + maxb = 8192 + self._s = maxb + hdr = self.read(maxb) + hdrend = hdr.find(NULL_BYTE) + typ, size = hdr[:hdrend].split(BYTE_SPACE) + size = int(size) + self._s = size + + # adjust internal state to match actual header length that we ignore + # The buffer will be depleted first on future reads + self._br = 0 + hdrend += 1 + self._buf = BytesIO(hdr[hdrend:]) + self._buflen = len(hdr) - hdrend + + self._phi = True + + return typ, size + + #{ Interface + + @classmethod + def new(self, m, close_on_deletion=False): + """Create a new DecompressMemMapReader instance for acting as a read-only stream + This method parses the object header from m and returns the parsed + type and size, as well as the created stream instance. + + :param m: memory map on which to operate. It must be object data ( header + contents ) + :param close_on_deletion: if True, the memory map will be closed once we are + being deleted""" + inst = DecompressMemMapReader(m, close_on_deletion, 0) + typ, size = inst._parse_header_info() + return typ, size, inst + + def data(self): + """:return: random access compatible data we are working on""" + return self._m + + def close(self): + """Close our underlying stream of compressed bytes if this was allowed during initialization + :return: True if we closed the underlying stream + :note: can be called safely + """ + if self._close: + if hasattr(self._m, 'close'): + self._m.close() + self._close = False + # END handle resource freeing + + def compressed_bytes_read(self): + """ + :return: number of compressed bytes read. This includes the bytes it + took to decompress the header ( if there was one )""" + # ABSTRACT: When decompressing a byte stream, it can be that the first + # x bytes which were requested match the first x bytes in the loosely + # compressed datastream. This is the worst-case assumption that the reader + # does, it assumes that it will get at least X bytes from X compressed bytes + # in call cases. + # The caveat is that the object, according to our known uncompressed size, + # is already complete, but there are still some bytes left in the compressed + # stream that contribute to the amount of compressed bytes. + # How can we know that we are truly done, and have read all bytes we need + # to read ? + # Without help, we cannot know, as we need to obtain the status of the + # decompression. If it is not finished, we need to decompress more data + # until it is finished, to yield the actual number of compressed bytes + # belonging to the decompressed object + # We are using a custom zlib module for this, if its not present, + # we try to put in additional bytes up for decompression if feasible + # and check for the unused_data. + + # Only scrub the stream forward if we are officially done with the + # bytes we were to have. + if self._br == self._s and not self._zip.unused_data: + # manipulate the bytes-read to allow our own read method to continue + # but keep the window at its current position + self._br = 0 + if hasattr(self._zip, 'status'): + while self._zip.status == zlib.Z_OK: + self.read(mmap.PAGESIZE) + # END scrub-loop custom zlib + else: + # pass in additional pages, until we have unused data + while not self._zip.unused_data and self._cbr != len(self._m): + self.read(mmap.PAGESIZE) + # END scrub-loop default zlib + # END handle stream scrubbing + + # reset bytes read, just to be sure + self._br = self._s + # END handle stream scrubbing + + # unused data ends up in the unconsumed tail, which was removed + # from the count already + return self._cbr + + #} END interface + + def seek(self, offset, whence=getattr(os, 'SEEK_SET', 0)): + """Allows to reset the stream to restart reading + :raise ValueError: If offset and whence are not 0""" + if offset != 0 or whence != getattr(os, 'SEEK_SET', 0): + raise ValueError("Can only seek to position 0") + # END handle offset + + self._zip = zlib.decompressobj() + self._br = self._cws = self._cwe = self._cbr = 0 + if self._phi: + self._phi = False + del(self._s) # trigger header parsing on first access + # END skip header + + def read(self, size=-1): + if size < 1: + size = self._s - self._br + else: + size = min(size, self._s - self._br) + # END clamp size + + if size == 0: + return b'' + # END handle depletion + + # deplete the buffer, then just continue using the decompress object + # which has an own buffer. We just need this to transparently parse the + # header from the zlib stream + dat = b'' + if self._buf: + if self._buflen >= size: + # have enough data + dat = self._buf.read(size) + self._buflen -= size + self._br += size + return dat + else: + dat = self._buf.read() # ouch, duplicates data + size -= self._buflen + self._br += self._buflen + + self._buflen = 0 + self._buf = None + # END handle buffer len + # END handle buffer + + # decompress some data + # Abstract: zlib needs to operate on chunks of our memory map ( which may + # be large ), as it will otherwise and always fill in the 'unconsumed_tail' + # attribute which possible reads our whole map to the end, forcing + # everything to be read from disk even though just a portion was requested. + # As this would be a nogo, we workaround it by passing only chunks of data, + # moving the window into the memory map along as we decompress, which keeps + # the tail smaller than our chunk-size. This causes 'only' the chunk to be + # copied once, and another copy of a part of it when it creates the unconsumed + # tail. We have to use it to hand in the appropriate amount of bytes during + # the next read. + tail = self._zip.unconsumed_tail + if tail: + # move the window, make it as large as size demands. For code-clarity, + # we just take the chunk from our map again instead of reusing the unconsumed + # tail. The latter one would safe some memory copying, but we could end up + # with not getting enough data uncompressed, so we had to sort that out as well. + # Now we just assume the worst case, hence the data is uncompressed and the window + # needs to be as large as the uncompressed bytes we want to read. + self._cws = self._cwe - len(tail) + self._cwe = self._cws + size + else: + cws = self._cws + self._cws = self._cwe + self._cwe = cws + size + # END handle tail + + # if window is too small, make it larger so zip can decompress something + if self._cwe - self._cws < 8: + self._cwe = self._cws + 8 + # END adjust winsize + + # takes a slice, but doesn't copy the data, it says ... + indata = self._m[self._cws:self._cwe] + + # get the actual window end to be sure we don't use it for computations + self._cwe = self._cws + len(indata) + dcompdat = self._zip.decompress(indata, size) + # update the amount of compressed bytes read + # We feed possibly overlapping chunks, which is why the unconsumed tail + # has to be taken into consideration, as well as the unused data + # if we hit the end of the stream + # NOTE: Behavior changed in PY2.7 onward, which requires special handling to make the tests work properly. + # They are thorough, and I assume it is truly working. + # Why is this logic as convoluted as it is ? Please look at the table in + # https://github.com/gitpython-developers/gitdb/issues/19 to learn about the test-results. + # Basically, on py2.6, you want to use branch 1, whereas on all other python version, the second branch + # will be the one that works. + # However, the zlib VERSIONs as well as the platform check is used to further match the entries in the + # table in the github issue. This is it ... it was the only way I could make this work everywhere. + # IT's CERTAINLY GOING TO BITE US IN THE FUTURE ... . + if getattr(zlib, 'ZLIB_RUNTIME_VERSION', zlib.ZLIB_VERSION) in ('1.2.7', '1.2.5') and not sys.platform == 'darwin': + unused_datalen = len(self._zip.unconsumed_tail) + else: + unused_datalen = len(self._zip.unconsumed_tail) + len(self._zip.unused_data) + # # end handle very special case ... + + self._cbr += len(indata) - unused_datalen + self._br += len(dcompdat) + + if dat: + dcompdat = dat + dcompdat + # END prepend our cached data + + # it can happen, depending on the compression, that we get less bytes + # than ordered as it needs the final portion of the data as well. + # Recursively resolve that. + # Note: dcompdat can be empty even though we still appear to have bytes + # to read, if we are called by compressed_bytes_read - it manipulates + # us to empty the stream + if dcompdat and (len(dcompdat) - len(dat)) < size and self._br < self._s: + dcompdat += self.read(size - len(dcompdat)) + # END handle special case + return dcompdat + + +class DeltaApplyReader(LazyMixin): + + """A reader which dynamically applies pack deltas to a base object, keeping the + memory demands to a minimum. + + The size of the final object is only obtainable once all deltas have been + applied, unless it is retrieved from a pack index. + + The uncompressed Delta has the following layout (MSB being a most significant + bit encoded dynamic size): + + * MSB Source Size - the size of the base against which the delta was created + * MSB Target Size - the size of the resulting data after the delta was applied + * A list of one byte commands (cmd) which are followed by a specific protocol: + + * cmd & 0x80 - copy delta_data[offset:offset+size] + + * Followed by an encoded offset into the delta data + * Followed by an encoded size of the chunk to copy + + * cmd & 0x7f - insert + + * insert cmd bytes from the delta buffer into the output stream + + * cmd == 0 - invalid operation ( or error in delta stream ) + """ + __slots__ = ( + "_bstream", # base stream to which to apply the deltas + "_dstreams", # tuple of delta stream readers + "_mm_target", # memory map of the delta-applied data + "_size", # actual number of bytes in _mm_target + "_br" # number of bytes read + ) + + #{ Configuration + k_max_memory_move = 250 * 1000 * 1000 + #} END configuration + + def __init__(self, stream_list): + """Initialize this instance with a list of streams, the first stream being + the delta to apply on top of all following deltas, the last stream being the + base object onto which to apply the deltas""" + assert len(stream_list) > 1, "Need at least one delta and one base stream" + + self._bstream = stream_list[-1] + self._dstreams = tuple(stream_list[:-1]) + self._br = 0 + + def _set_cache_too_slow_without_c(self, attr): + # the direct algorithm is fastest and most direct if there is only one + # delta. Also, the extra overhead might not be worth it for items smaller + # than X - definitely the case in python, every function call costs + # huge amounts of time + # if len(self._dstreams) * self._bstream.size < self.k_max_memory_move: + if len(self._dstreams) == 1: + return self._set_cache_brute_(attr) + + # Aggregate all deltas into one delta in reverse order. Hence we take + # the last delta, and reverse-merge its ancestor delta, until we receive + # the final delta data stream. + dcl = connect_deltas(self._dstreams) + + # call len directly, as the (optional) c version doesn't implement the sequence + # protocol + if dcl.rbound() == 0: + self._size = 0 + self._mm_target = allocate_memory(0) + return + # END handle empty list + + self._size = dcl.rbound() + self._mm_target = allocate_memory(self._size) + + bbuf = allocate_memory(self._bstream.size) + stream_copy(self._bstream.read, bbuf.write, self._bstream.size, 256 * mmap.PAGESIZE) + + # APPLY CHUNKS + write = self._mm_target.write + dcl.apply(bbuf, write) + + self._mm_target.seek(0) + + def _set_cache_brute_(self, attr): + """If we are here, we apply the actual deltas""" + # TODO: There should be a special case if there is only one stream + # Then the default-git algorithm should perform a tad faster, as the + # delta is not peaked into, causing less overhead. + buffer_info_list = list() + max_target_size = 0 + for dstream in self._dstreams: + buf = dstream.read(512) # read the header information + X + offset, src_size = msb_size(buf) + offset, target_size = msb_size(buf, offset) + buffer_info_list.append((buf[offset:], offset, src_size, target_size)) + max_target_size = max(max_target_size, target_size) + # END for each delta stream + + # sanity check - the first delta to apply should have the same source + # size as our actual base stream + base_size = self._bstream.size + target_size = max_target_size + + # if we have more than 1 delta to apply, we will swap buffers, hence we must + # assure that all buffers we use are large enough to hold all the results + if len(self._dstreams) > 1: + base_size = target_size = max(base_size, max_target_size) + # END adjust buffer sizes + + # Allocate private memory map big enough to hold the first base buffer + # We need random access to it + bbuf = allocate_memory(base_size) + stream_copy(self._bstream.read, bbuf.write, base_size, 256 * mmap.PAGESIZE) + + # allocate memory map large enough for the largest (intermediate) target + # We will use it as scratch space for all delta ops. If the final + # target buffer is smaller than our allocated space, we just use parts + # of it upon return. + tbuf = allocate_memory(target_size) + + # for each delta to apply, memory map the decompressed delta and + # work on the op-codes to reconstruct everything. + # For the actual copying, we use a seek and write pattern of buffer + # slices. + final_target_size = None + for (dbuf, offset, src_size, target_size), dstream in zip(reversed(buffer_info_list), reversed(self._dstreams)): + # allocate a buffer to hold all delta data - fill in the data for + # fast access. We do this as we know that reading individual bytes + # from our stream would be slower than necessary ( although possible ) + # The dbuf buffer contains commands after the first two MSB sizes, the + # offset specifies the amount of bytes read to get the sizes. + ddata = allocate_memory(dstream.size - offset) + ddata.write(dbuf) + # read the rest from the stream. The size we give is larger than necessary + stream_copy(dstream.read, ddata.write, dstream.size, 256 * mmap.PAGESIZE) + + ####################################################################### + if 'c_apply_delta' in globals(): + c_apply_delta(bbuf, ddata, tbuf) + else: + apply_delta_data(bbuf, src_size, ddata, len(ddata), tbuf.write) + ####################################################################### + + # finally, swap out source and target buffers. The target is now the + # base for the next delta to apply + bbuf, tbuf = tbuf, bbuf + bbuf.seek(0) + tbuf.seek(0) + final_target_size = target_size + # END for each delta to apply + + # its already seeked to 0, constrain it to the actual size + # NOTE: in the end of the loop, it swaps buffers, hence our target buffer + # is not tbuf, but bbuf ! + self._mm_target = bbuf + self._size = final_target_size + + #{ Configuration + if not has_perf_mod: + _set_cache_ = _set_cache_brute_ + else: + _set_cache_ = _set_cache_too_slow_without_c + + #} END configuration + + def read(self, count=0): + bl = self._size - self._br # bytes left + if count < 1 or count > bl: + count = bl + # NOTE: we could check for certain size limits, and possibly + # return buffers instead of strings to prevent byte copying + data = self._mm_target.read(count) + self._br += len(data) + return data + + def seek(self, offset, whence=getattr(os, 'SEEK_SET', 0)): + """Allows to reset the stream to restart reading + + :raise ValueError: If offset and whence are not 0""" + if offset != 0 or whence != getattr(os, 'SEEK_SET', 0): + raise ValueError("Can only seek to position 0") + # END handle offset + self._br = 0 + self._mm_target.seek(0) + + #{ Interface + + @classmethod + def new(cls, stream_list): + """ + Convert the given list of streams into a stream which resolves deltas + when reading from it. + + :param stream_list: two or more stream objects, first stream is a Delta + to the object that you want to resolve, followed by N additional delta + streams. The list's last stream must be a non-delta stream. + + :return: Non-Delta OPackStream object whose stream can be used to obtain + the decompressed resolved data + :raise ValueError: if the stream list cannot be handled""" + if len(stream_list) < 2: + raise ValueError("Need at least two streams") + # END single object special handling + + if stream_list[-1].type_id in delta_types: + raise ValueError( + "Cannot resolve deltas if there is no base object stream, last one was type: %s" % stream_list[-1].type) + # END check stream + return cls(stream_list) + + #} END interface + + #{ OInfo like Interface + + @property + def type(self): + return self._bstream.type + + @property + def type_id(self): + return self._bstream.type_id + + @property + def size(self): + """:return: number of uncompressed bytes in the stream""" + return self._size + + #} END oinfo like interface + + +#} END RO streams + + +#{ W Streams + +class Sha1Writer: + + """Simple stream writer which produces a sha whenever you like as it degests + everything it is supposed to write""" + __slots__ = "sha1" + + def __init__(self): + self.sha1 = make_sha() + + #{ Stream Interface + + def write(self, data): + """:raise IOError: If not all bytes could be written + :param data: byte object + :return: length of incoming data""" + + self.sha1.update(data) + + return len(data) + + # END stream interface + + #{ Interface + + def sha(self, as_hex=False): + """:return: sha so far + :param as_hex: if True, sha will be hex-encoded, binary otherwise""" + if as_hex: + return self.sha1.hexdigest() + return self.sha1.digest() + + #} END interface + + +class FlexibleSha1Writer(Sha1Writer): + + """Writer producing a sha1 while passing on the written bytes to the given + write function""" + __slots__ = 'writer' + + def __init__(self, writer): + Sha1Writer.__init__(self) + self.writer = writer + + def write(self, data): + Sha1Writer.write(self, data) + self.writer(data) + + +class ZippedStoreShaWriter(Sha1Writer): + + """Remembers everything someone writes to it and generates a sha""" + __slots__ = ('buf', 'zip') + + def __init__(self): + Sha1Writer.__init__(self) + self.buf = BytesIO() + self.zip = zlib.compressobj(zlib.Z_BEST_SPEED) + + def __getattr__(self, attr): + return getattr(self.buf, attr) + + def write(self, data): + alen = Sha1Writer.write(self, data) + self.buf.write(self.zip.compress(data)) + + return alen + + def close(self): + self.buf.write(self.zip.flush()) + + def seek(self, offset, whence=getattr(os, 'SEEK_SET', 0)): + """Seeking currently only supports to rewind written data + Multiple writes are not supported""" + if offset != 0 or whence != getattr(os, 'SEEK_SET', 0): + raise ValueError("Can only seek to position 0") + # END handle offset + self.buf.seek(0) + + def getvalue(self): + """:return: string value from the current stream position to the end""" + return self.buf.getvalue() + + +class FDCompressedSha1Writer(Sha1Writer): + + """Digests data written to it, making the sha available, then compress the + data and write it to the file descriptor + + **Note:** operates on raw file descriptors + **Note:** for this to work, you have to use the close-method of this instance""" + __slots__ = ("fd", "sha1", "zip") + + # default exception + exc = IOError("Failed to write all bytes to filedescriptor") + + def __init__(self, fd): + super().__init__() + self.fd = fd + self.zip = zlib.compressobj(zlib.Z_BEST_SPEED) + + #{ Stream Interface + + def write(self, data): + """:raise IOError: If not all bytes could be written + :return: length of incoming data""" + self.sha1.update(data) + cdata = self.zip.compress(data) + bytes_written = write(self.fd, cdata) + + if bytes_written != len(cdata): + raise self.exc + + return len(data) + + def close(self): + remainder = self.zip.flush() + if write(self.fd, remainder) != len(remainder): + raise self.exc + return close(self.fd) + + #} END stream interface + + +class FDStream: + + """A simple wrapper providing the most basic functions on a file descriptor + with the fileobject interface. Cannot use os.fdopen as the resulting stream + takes ownership""" + __slots__ = ("_fd", '_pos') + + def __init__(self, fd): + self._fd = fd + self._pos = 0 + + def write(self, data): + self._pos += len(data) + os.write(self._fd, data) + + def read(self, count=0): + if count == 0: + count = os.path.getsize(self._filepath) + # END handle read everything + + bytes = os.read(self._fd, count) + self._pos += len(bytes) + return bytes + + def fileno(self): + return self._fd + + def tell(self): + return self._pos + + def close(self): + close(self._fd) + + +class NullStream: + + """A stream that does nothing but providing a stream interface. + Use it like /dev/null""" + __slots__ = tuple() + + def read(self, size=0): + return '' + + def close(self): + pass + + def write(self, data): + return len(data) + + +#} END W streams diff --git a/lib/python3.12/site-packages/gitdb/test/__init__.py b/lib/python3.12/site-packages/gitdb/test/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..03bd406be478bd82f20b45c589d8d029367c5f1a --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a748af7e87107a7ffa02fc5d3ef16e053b2c5ce2 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/lib.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/lib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf113be028e44c8ce15b753da9ebaddd0d875319 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/lib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/test_base.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7a042679a64dd369bcfdc9aab1d02fdbe8113c8 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/test_example.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7c017586918a5ae4f7684978fd1c96d4d003701 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/test_pack.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_pack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09a34186a595f08bc0b89f92c3e3fab0d9abf573 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_pack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/test_stream.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_stream.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b8825f11efe788964c1ac4a5605690009a337d2 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_stream.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/__pycache__/test_util.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c40e5168830304c3293e9ab8d3f3014af83608c5 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/test/__pycache__/test_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/test/lib.py b/lib/python3.12/site-packages/gitdb/test/lib.py new file mode 100644 index 0000000000000000000000000000000000000000..8e602342ca59ef9d556254c2bb089c1cd4c9bcf2 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/lib.py @@ -0,0 +1,192 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Utilities used in ODB testing""" +from gitdb import OStream + +import sys +import random +from array import array + +from io import BytesIO + +import glob +import unittest +import tempfile +import shutil +import os +import gc +import logging +from functools import wraps + + +#{ Bases + +class TestBase(unittest.TestCase): + """Base class for all tests + + TestCase providing access to readonly repositories using the following member variables. + + * gitrepopath + + * read-only base path of the git source repository, i.e. .../git/.git + """ + + #{ Invvariants + k_env_git_repo = "GITDB_TEST_GIT_REPO_BASE" + #} END invariants + + @classmethod + def setUpClass(cls): + try: + super().setUpClass() + except AttributeError: + pass + + cls.gitrepopath = os.environ.get(cls.k_env_git_repo) + if not cls.gitrepopath: + logging.info( + "You can set the %s environment variable to a .git repository of your choice - defaulting to the gitdb repository", cls.k_env_git_repo) + ospd = os.path.dirname + cls.gitrepopath = os.path.join(ospd(ospd(ospd(__file__))), '.git') + # end assure gitrepo is set + assert cls.gitrepopath.endswith('.git') + + +#} END bases + +#{ Decorators + +def with_rw_directory(func): + """Create a temporary directory which can be written to, remove it if the + test succeeds, but leave it otherwise to aid additional debugging""" + + def wrapper(self): + path = tempfile.mktemp(prefix=func.__name__) + os.mkdir(path) + keep = False + try: + try: + return func(self, path) + except Exception: + sys.stderr.write(f"Test {type(self).__name__}.{func.__name__} failed, output is at {path!r}\n") + keep = True + raise + finally: + # Need to collect here to be sure all handles have been closed. It appears + # a windows-only issue. In fact things should be deleted, as well as + # memory maps closed, once objects go out of scope. For some reason + # though this is not the case here unless we collect explicitly. + if not keep: + gc.collect() + shutil.rmtree(path) + # END handle exception + # END wrapper + + wrapper.__name__ = func.__name__ + return wrapper + + +def with_packs_rw(func): + """Function that provides a path into which the packs for testing should be + copied. Will pass on the path to the actual function afterwards""" + + def wrapper(self, path): + src_pack_glob = fixture_path('packs/*') + copy_files_globbed(src_pack_glob, path, hard_link_ok=True) + return func(self, path) + # END wrapper + + wrapper.__name__ = func.__name__ + return wrapper + +#} END decorators + +#{ Routines + + +def fixture_path(relapath=''): + """:return: absolute path into the fixture directory + :param relapath: relative path into the fixtures directory, or '' + to obtain the fixture directory itself""" + return os.path.join(os.path.dirname(__file__), 'fixtures', relapath) + + +def copy_files_globbed(source_glob, target_dir, hard_link_ok=False): + """Copy all files found according to the given source glob into the target directory + :param hard_link_ok: if True, hard links will be created if possible. Otherwise + the files will be copied""" + for src_file in glob.glob(source_glob): + if hard_link_ok and hasattr(os, 'link'): + target = os.path.join(target_dir, os.path.basename(src_file)) + try: + os.link(src_file, target) + except OSError: + shutil.copy(src_file, target_dir) + # END handle cross device links ( and resulting failure ) + else: + shutil.copy(src_file, target_dir) + # END try hard link + # END for each file to copy + + +def make_bytes(size_in_bytes, randomize=False): + """:return: string with given size in bytes + :param randomize: try to produce a very random stream""" + actual_size = size_in_bytes // 4 + producer = range(actual_size) + if randomize: + producer = list(producer) + random.shuffle(producer) + # END randomize + a = array('i', producer) + return a.tobytes() + + +def make_object(type, data): + """:return: bytes resembling an uncompressed object""" + odata = "blob %i\0" % len(data) + return odata.encode("ascii") + data + + +def make_memory_file(size_in_bytes, randomize=False): + """:return: tuple(size_of_stream, stream) + :param randomize: try to produce a very random stream""" + d = make_bytes(size_in_bytes, randomize) + return len(d), BytesIO(d) + +#} END routines + +#{ Stream Utilities + + +class DummyStream: + + def __init__(self): + self.was_read = False + self.bytes = 0 + self.closed = False + + def read(self, size): + self.was_read = True + self.bytes = size + + def close(self): + self.closed = True + + def _assert(self): + assert self.was_read + + +class DeriveTest(OStream): + + def __init__(self, sha, type, size, stream, *args, **kwargs): + self.myarg = kwargs.pop('myarg') + self.args = args + + def _assert(self): + assert self.args + assert self.myarg + +#} END stream utilitiess diff --git a/lib/python3.12/site-packages/gitdb/test/test_base.py b/lib/python3.12/site-packages/gitdb/test/test_base.py new file mode 100644 index 0000000000000000000000000000000000000000..17906c9c205b622a0a34b5ac13eb16dacfbef4fb --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/test_base.py @@ -0,0 +1,105 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Test for object db""" +from gitdb.test.lib import ( + TestBase, + DummyStream, + DeriveTest, +) + +from gitdb import ( + OInfo, + OPackInfo, + ODeltaPackInfo, + OStream, + OPackStream, + ODeltaPackStream, + IStream +) +from gitdb.util import ( + NULL_BIN_SHA +) + +from gitdb.typ import ( + str_blob_type +) + + +class TestBaseTypes(TestBase): + + def test_streams(self): + # test info + sha = NULL_BIN_SHA + s = 20 + blob_id = 3 + + info = OInfo(sha, str_blob_type, s) + assert info.binsha == sha + assert info.type == str_blob_type + assert info.type_id == blob_id + assert info.size == s + + # test pack info + # provides type_id + pinfo = OPackInfo(0, blob_id, s) + assert pinfo.type == str_blob_type + assert pinfo.type_id == blob_id + assert pinfo.pack_offset == 0 + + dpinfo = ODeltaPackInfo(0, blob_id, s, sha) + assert dpinfo.type == str_blob_type + assert dpinfo.type_id == blob_id + assert dpinfo.delta_info == sha + assert dpinfo.pack_offset == 0 + + # test ostream + stream = DummyStream() + ostream = OStream(*(info + (stream, ))) + assert ostream.stream is stream + ostream.read(15) + stream._assert() + assert stream.bytes == 15 + ostream.read(20) + assert stream.bytes == 20 + + # test packstream + postream = OPackStream(*(pinfo + (stream, ))) + assert postream.stream is stream + postream.read(10) + stream._assert() + assert stream.bytes == 10 + + # test deltapackstream + dpostream = ODeltaPackStream(*(dpinfo + (stream, ))) + assert dpostream.stream is stream + dpostream.read(5) + stream._assert() + assert stream.bytes == 5 + + # derive with own args + DeriveTest(sha, str_blob_type, s, stream, 'mine', myarg=3)._assert() + + # test istream + istream = IStream(str_blob_type, s, stream) + assert istream.binsha == None + istream.binsha = sha + assert istream.binsha == sha + + assert len(istream.binsha) == 20 + assert len(istream.hexsha) == 40 + + assert istream.size == s + istream.size = s * 2 + assert istream.size == s * 2 + assert istream.type == str_blob_type + istream.type = "something" + assert istream.type == "something" + assert istream.stream is stream + istream.stream = None + assert istream.stream is None + + assert istream.error is None + istream.error = Exception() + assert isinstance(istream.error, Exception) diff --git a/lib/python3.12/site-packages/gitdb/test/test_example.py b/lib/python3.12/site-packages/gitdb/test/test_example.py new file mode 100644 index 0000000000000000000000000000000000000000..3b4c9084bb2cb10e58ffe656bc19f9b3142e3e65 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/test_example.py @@ -0,0 +1,43 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Module with examples from the tutorial section of the docs""" +import os +from gitdb.test.lib import TestBase +from gitdb import IStream +from gitdb.db import LooseObjectDB + +from io import BytesIO + + +class TestExamples(TestBase): + + def test_base(self): + ldb = LooseObjectDB(os.path.join(self.gitrepopath, 'objects')) + + for sha1 in ldb.sha_iter(): + oinfo = ldb.info(sha1) + ostream = ldb.stream(sha1) + assert oinfo[:3] == ostream[:3] + + assert len(ostream.read()) == ostream.size + assert ldb.has_object(oinfo.binsha) + # END for each sha in database + # assure we close all files + try: + del(ostream) + del(oinfo) + except UnboundLocalError: + pass + # END ignore exception if there are no loose objects + + data = b"my data" + istream = IStream("blob", len(data), BytesIO(data)) + + # the object does not yet have a sha + assert istream.binsha is None + ldb.store(istream) + # now the sha is set + assert len(istream.binsha) == 20 + assert ldb.has_object(istream.binsha) diff --git a/lib/python3.12/site-packages/gitdb/test/test_pack.py b/lib/python3.12/site-packages/gitdb/test/test_pack.py new file mode 100644 index 0000000000000000000000000000000000000000..e7234822805f05373c3122d010424c0f339d5027 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/test_pack.py @@ -0,0 +1,249 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Test everything about packs reading and writing""" +from gitdb.test.lib import ( + TestBase, + with_rw_directory, + fixture_path +) + +from gitdb.stream import DeltaApplyReader + +from gitdb.pack import ( + PackEntity, + PackIndexFile, + PackFile +) + +from gitdb.base import ( + OInfo, + OStream, +) + +from gitdb.fun import delta_types +from gitdb.exc import UnsupportedOperation +from gitdb.util import to_bin_sha + +import pytest + +import os +import tempfile + + +#{ Utilities +def bin_sha_from_filename(filename): + return to_bin_sha(os.path.splitext(os.path.basename(filename))[0][5:]) +#} END utilities + + +class TestPack(TestBase): + + packindexfile_v1 = (fixture_path('packs/pack-c0438c19fb16422b6bbcce24387b3264416d485b.idx'), 1, 67) + packindexfile_v2 = (fixture_path('packs/pack-11fdfa9e156ab73caae3b6da867192221f2089c2.idx'), 2, 30) + packindexfile_v2_3_ascii = (fixture_path('packs/pack-a2bf8e71d8c18879e499335762dd95119d93d9f1.idx'), 2, 42) + packfile_v2_1 = (fixture_path('packs/pack-c0438c19fb16422b6bbcce24387b3264416d485b.pack'), 2, packindexfile_v1[2]) + packfile_v2_2 = (fixture_path('packs/pack-11fdfa9e156ab73caae3b6da867192221f2089c2.pack'), 2, packindexfile_v2[2]) + packfile_v2_3_ascii = ( + fixture_path('packs/pack-a2bf8e71d8c18879e499335762dd95119d93d9f1.pack'), 2, packindexfile_v2_3_ascii[2]) + + def _assert_index_file(self, index, version, size): + assert index.packfile_checksum() != index.indexfile_checksum() + assert len(index.packfile_checksum()) == 20 + assert len(index.indexfile_checksum()) == 20 + assert index.version() == version + assert index.size() == size + assert len(index.offsets()) == size + + # get all data of all objects + for oidx in range(index.size()): + sha = index.sha(oidx) + assert oidx == index.sha_to_index(sha) + + entry = index.entry(oidx) + assert len(entry) == 3 + + assert entry[0] == index.offset(oidx) + assert entry[1] == sha + assert entry[2] == index.crc(oidx) + + # verify partial sha + for l in (4, 8, 11, 17, 20): + assert index.partial_sha_to_index(sha[:l], l * 2) == oidx + + # END for each object index in indexfile + self.assertRaises(ValueError, index.partial_sha_to_index, "\0", 2) + + def _assert_pack_file(self, pack, version, size): + assert pack.version() == 2 + assert pack.size() == size + assert len(pack.checksum()) == 20 + + num_obj = 0 + for obj in pack.stream_iter(): + num_obj += 1 + info = pack.info(obj.pack_offset) + stream = pack.stream(obj.pack_offset) + + assert info.pack_offset == stream.pack_offset + assert info.type_id == stream.type_id + assert hasattr(stream, 'read') + + # it should be possible to read from both streams + assert obj.read() == stream.read() + + streams = pack.collect_streams(obj.pack_offset) + assert streams + + # read the stream + try: + dstream = DeltaApplyReader.new(streams) + except ValueError: + # ignore these, old git versions use only ref deltas, + # which we haven't resolved ( as we are without an index ) + # Also ignore non-delta streams + continue + # END get deltastream + + # read all + data = dstream.read() + assert len(data) == dstream.size + + # test seek + dstream.seek(0) + assert dstream.read() == data + + # read chunks + # NOTE: the current implementation is safe, it basically transfers + # all calls to the underlying memory map + + # END for each object + assert num_obj == size + + def test_pack_index(self): + # check version 1 and 2 + for indexfile, version, size in (self.packindexfile_v1, self.packindexfile_v2): + index = PackIndexFile(indexfile) + self._assert_index_file(index, version, size) + # END run tests + + def test_pack(self): + # there is this special version 3, but apparently its like 2 ... + for packfile, version, size in (self.packfile_v2_3_ascii, self.packfile_v2_1, self.packfile_v2_2): + pack = PackFile(packfile) + self._assert_pack_file(pack, version, size) + # END for each pack to test + + @with_rw_directory + def test_pack_entity(self, rw_dir): + pack_objs = list() + for packinfo, indexinfo in ((self.packfile_v2_1, self.packindexfile_v1), + (self.packfile_v2_2, self.packindexfile_v2), + (self.packfile_v2_3_ascii, self.packindexfile_v2_3_ascii)): + packfile, version, size = packinfo + indexfile, version, size = indexinfo + entity = PackEntity(packfile) + assert entity.pack().path() == packfile + assert entity.index().path() == indexfile + pack_objs.extend(entity.stream_iter()) + + count = 0 + for info, stream in zip(entity.info_iter(), entity.stream_iter()): + count += 1 + assert info.binsha == stream.binsha + assert len(info.binsha) == 20 + assert info.type_id == stream.type_id + assert info.size == stream.size + + # we return fully resolved items, which is implied by the sha centric access + assert not info.type_id in delta_types + + # try all calls + assert len(entity.collect_streams(info.binsha)) + oinfo = entity.info(info.binsha) + assert isinstance(oinfo, OInfo) + assert oinfo.binsha is not None + ostream = entity.stream(info.binsha) + assert isinstance(ostream, OStream) + assert ostream.binsha is not None + + # verify the stream + try: + assert entity.is_valid_stream(info.binsha, use_crc=True) + except UnsupportedOperation: + pass + # END ignore version issues + assert entity.is_valid_stream(info.binsha, use_crc=False) + # END for each info, stream tuple + assert count == size + + # END for each entity + + # pack writing - write all packs into one + # index path can be None + pack_path1 = tempfile.mktemp('', "pack1", rw_dir) + pack_path2 = tempfile.mktemp('', "pack2", rw_dir) + index_path = tempfile.mktemp('', 'index', rw_dir) + iteration = 0 + + def rewind_streams(): + for obj in pack_objs: + obj.stream.seek(0) + # END utility + for ppath, ipath, num_obj in zip((pack_path1, pack_path2), + (index_path, None), + (len(pack_objs), None)): + iwrite = None + if ipath: + ifile = open(ipath, 'wb') + iwrite = ifile.write + # END handle ip + + # make sure we rewind the streams ... we work on the same objects over and over again + if iteration > 0: + rewind_streams() + # END rewind streams + iteration += 1 + + with open(ppath, 'wb') as pfile: + pack_sha, index_sha = PackEntity.write_pack(pack_objs, pfile.write, iwrite, object_count=num_obj) + assert os.path.getsize(ppath) > 100 + + # verify pack + pf = PackFile(ppath) + assert pf.size() == len(pack_objs) + assert pf.version() == PackFile.pack_version_default + assert pf.checksum() == pack_sha + pf.close() + + # verify index + if ipath is not None: + ifile.close() + assert os.path.getsize(ipath) > 100 + idx = PackIndexFile(ipath) + assert idx.version() == PackIndexFile.index_version_default + assert idx.packfile_checksum() == pack_sha + assert idx.indexfile_checksum() == index_sha + assert idx.size() == len(pack_objs) + idx.close() + # END verify files exist + # END for each packpath, indexpath pair + + # verify the packs thoroughly + rewind_streams() + entity = PackEntity.create(pack_objs, rw_dir) + count = 0 + for info in entity.info_iter(): + count += 1 + for use_crc in range(2): + assert entity.is_valid_stream(info.binsha, use_crc) + # END for each crc mode + # END for each info + assert count == len(pack_objs) + entity.close() + + def test_pack_64(self): + # TODO: hex-edit a pack helping us to verify that we can handle 64 byte offsets + # of course without really needing such a huge pack + pytest.skip('not implemented') diff --git a/lib/python3.12/site-packages/gitdb/test/test_stream.py b/lib/python3.12/site-packages/gitdb/test/test_stream.py new file mode 100644 index 0000000000000000000000000000000000000000..1e7e941d1cad57e7789de51941ec2c6b3587c1a5 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/test_stream.py @@ -0,0 +1,164 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Test for object db""" + +from gitdb.test.lib import ( + TestBase, + DummyStream, + make_bytes, + make_object, + fixture_path +) + +from gitdb import ( + DecompressMemMapReader, + FDCompressedSha1Writer, + LooseObjectDB, + Sha1Writer, + MemoryDB, + IStream, +) +from gitdb.util import hex_to_bin + +import zlib +from gitdb.typ import ( + str_blob_type +) + +import tempfile +import os +from io import BytesIO + + +class TestStream(TestBase): + + """Test stream classes""" + + data_sizes = (15, 10000, 1000 * 1024 + 512) + + def _assert_stream_reader(self, stream, cdata, rewind_stream=lambda s: None): + """Make stream tests - the orig_stream is seekable, allowing it to be + rewound and reused + :param cdata: the data we expect to read from stream, the contents + :param rewind_stream: function called to rewind the stream to make it ready + for reuse""" + ns = 10 + assert len(cdata) > ns - 1, "Data must be larger than %i, was %i" % (ns, len(cdata)) + + # read in small steps + ss = len(cdata) // ns + for i in range(ns): + data = stream.read(ss) + chunk = cdata[i * ss:(i + 1) * ss] + assert data == chunk + # END for each step + rest = stream.read() + if rest: + assert rest == cdata[-len(rest):] + # END handle rest + + if isinstance(stream, DecompressMemMapReader): + assert len(stream.data()) == stream.compressed_bytes_read() + # END handle special type + + rewind_stream(stream) + + # read everything + rdata = stream.read() + assert rdata == cdata + + if isinstance(stream, DecompressMemMapReader): + assert len(stream.data()) == stream.compressed_bytes_read() + # END handle special type + + def test_decompress_reader(self): + for close_on_deletion in range(2): + for with_size in range(2): + for ds in self.data_sizes: + cdata = make_bytes(ds, randomize=False) + + # zdata = zipped actual data + # cdata = original content data + + # create reader + if with_size: + # need object data + zdata = zlib.compress(make_object(str_blob_type, cdata)) + typ, size, reader = DecompressMemMapReader.new(zdata, close_on_deletion) + assert size == len(cdata) + assert typ == str_blob_type + + # even if we don't set the size, it will be set automatically on first read + test_reader = DecompressMemMapReader(zdata, close_on_deletion=False) + assert test_reader._s == len(cdata) + else: + # here we need content data + zdata = zlib.compress(cdata) + reader = DecompressMemMapReader(zdata, close_on_deletion, len(cdata)) + assert reader._s == len(cdata) + # END get reader + + self._assert_stream_reader(reader, cdata, lambda r: r.seek(0)) + + # put in a dummy stream for closing + dummy = DummyStream() + reader._m = dummy + + assert not dummy.closed + del(reader) + assert dummy.closed == close_on_deletion + # END for each datasize + # END whether size should be used + # END whether stream should be closed when deleted + + def test_sha_writer(self): + writer = Sha1Writer() + assert 2 == writer.write(b"hi") + assert len(writer.sha(as_hex=1)) == 40 + assert len(writer.sha(as_hex=0)) == 20 + + # make sure it does something ;) + prev_sha = writer.sha() + writer.write(b"hi again") + assert writer.sha() != prev_sha + + def test_compressed_writer(self): + for ds in self.data_sizes: + fd, path = tempfile.mkstemp() + ostream = FDCompressedSha1Writer(fd) + data = make_bytes(ds, randomize=False) + + # for now, just a single write, code doesn't care about chunking + assert len(data) == ostream.write(data) + ostream.close() + + # its closed already + self.assertRaises(OSError, os.close, fd) + + # read everything back, compare to data we zip + fd = os.open(path, os.O_RDONLY | getattr(os, 'O_BINARY', 0)) + written_data = os.read(fd, os.path.getsize(path)) + assert len(written_data) == os.path.getsize(path) + os.close(fd) + assert written_data == zlib.compress(data, 1) # best speed + + os.remove(path) + # END for each os + + def test_decompress_reader_special_case(self): + odb = LooseObjectDB(fixture_path('objects')) + mdb = MemoryDB() + for sha in (b'888401851f15db0eed60eb1bc29dec5ddcace911', + b'7bb839852ed5e3a069966281bb08d50012fb309b',): + ostream = odb.stream(hex_to_bin(sha)) + + # if there is a bug, we will be missing one byte exactly ! + data = ostream.read() + assert len(data) == ostream.size + + # Putting it back in should yield nothing new - after all, we have + dump = mdb.store(IStream(ostream.type, ostream.size, BytesIO(data))) + assert dump.hexsha == sha + # end for each loose object sha to test diff --git a/lib/python3.12/site-packages/gitdb/test/test_util.py b/lib/python3.12/site-packages/gitdb/test/test_util.py new file mode 100644 index 0000000000000000000000000000000000000000..166b33c3a6c8917fc737aefddbc85041aadc1b48 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/test/test_util.py @@ -0,0 +1,100 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Test for object db""" +import tempfile +import os + +from gitdb.test.lib import TestBase +from gitdb.util import ( + to_hex_sha, + to_bin_sha, + NULL_HEX_SHA, + LockedFD +) + + +class TestUtils(TestBase): + + def test_basics(self): + assert to_hex_sha(NULL_HEX_SHA) == NULL_HEX_SHA + assert len(to_bin_sha(NULL_HEX_SHA)) == 20 + assert to_hex_sha(to_bin_sha(NULL_HEX_SHA)) == NULL_HEX_SHA.encode("ascii") + + def _cmp_contents(self, file_path, data): + # raise if data from file at file_path + # does not match data string + with open(file_path, "rb") as fp: + assert fp.read() == data.encode("ascii") + + def test_lockedfd(self): + my_file = tempfile.mktemp() + orig_data = "hello" + new_data = "world" + with open(my_file, "wb") as my_file_fp: + my_file_fp.write(orig_data.encode("ascii")) + + try: + lfd = LockedFD(my_file) + lockfilepath = lfd._lockfilepath() + + # cannot end before it was started + self.assertRaises(AssertionError, lfd.rollback) + self.assertRaises(AssertionError, lfd.commit) + + # open for writing + assert not os.path.isfile(lockfilepath) + wfd = lfd.open(write=True) + assert lfd._fd is wfd + assert os.path.isfile(lockfilepath) + + # write data and fail + os.write(wfd, new_data.encode("ascii")) + lfd.rollback() + assert lfd._fd is None + self._cmp_contents(my_file, orig_data) + assert not os.path.isfile(lockfilepath) + + # additional call doesn't fail + lfd.commit() + lfd.rollback() + + # test reading + lfd = LockedFD(my_file) + rfd = lfd.open(write=False) + assert os.read(rfd, len(orig_data)) == orig_data.encode("ascii") + + assert os.path.isfile(lockfilepath) + # deletion rolls back + del(lfd) + assert not os.path.isfile(lockfilepath) + + # write data - concurrently + lfd = LockedFD(my_file) + olfd = LockedFD(my_file) + assert not os.path.isfile(lockfilepath) + wfdstream = lfd.open(write=True, stream=True) # this time as stream + assert os.path.isfile(lockfilepath) + # another one fails + self.assertRaises(IOError, olfd.open) + + wfdstream.write(new_data.encode("ascii")) + lfd.commit() + assert not os.path.isfile(lockfilepath) + self._cmp_contents(my_file, new_data) + + # could test automatic _end_writing on destruction + finally: + os.remove(my_file) + # END final cleanup + + # try non-existing file for reading + lfd = LockedFD(tempfile.mktemp()) + try: + lfd.open(write=False) + except OSError: + assert not os.path.exists(lfd._lockfilepath()) + else: + self.fail("expected OSError") + # END handle exceptions diff --git a/lib/python3.12/site-packages/gitdb/typ.py b/lib/python3.12/site-packages/gitdb/typ.py new file mode 100644 index 0000000000000000000000000000000000000000..314db50a73ecf8d3266ea513e82f43b71d14a25f --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/typ.py @@ -0,0 +1,10 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +"""Module containing information about types known to the database""" + +str_blob_type = b'blob' +str_commit_type = b'commit' +str_tree_type = b'tree' +str_tag_type = b'tag' diff --git a/lib/python3.12/site-packages/gitdb/util.py b/lib/python3.12/site-packages/gitdb/util.py new file mode 100644 index 0000000000000000000000000000000000000000..bb6d8797a051f810c9538d5dbfd776b0ef22324d --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/util.py @@ -0,0 +1,398 @@ +# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors +# +# This module is part of GitDB and is released under +# the New BSD License: https://opensource.org/license/bsd-3-clause/ +import binascii +import os +import mmap +import sys +import time +import errno + +from io import BytesIO + +from smmap import ( + StaticWindowMapManager, + SlidingWindowMapManager, + SlidingWindowMapBuffer +) + +# initialize our global memory manager instance +# Use it to free cached (and unused) resources. +mman = SlidingWindowMapManager() +# END handle mman + +import hashlib + +try: + from struct import unpack_from +except ImportError: + from struct import unpack, calcsize + __calcsize_cache = dict() + + def unpack_from(fmt, data, offset=0): + try: + size = __calcsize_cache[fmt] + except KeyError: + size = calcsize(fmt) + __calcsize_cache[fmt] = size + # END exception handling + return unpack(fmt, data[offset: offset + size]) + # END own unpack_from implementation + + +#{ Aliases + +hex_to_bin = binascii.a2b_hex +bin_to_hex = binascii.b2a_hex + +# errors +ENOENT = errno.ENOENT + +# os shortcuts +exists = os.path.exists +mkdir = os.mkdir +chmod = os.chmod +isdir = os.path.isdir +isfile = os.path.isfile +rename = os.rename +dirname = os.path.dirname +basename = os.path.basename +join = os.path.join +read = os.read +write = os.write +close = os.close +fsync = os.fsync + + +def _retry(func, *args, **kwargs): + # Wrapper around functions, that are problematic on "Windows". Sometimes + # the OS or someone else has still a handle to the file + if sys.platform == "win32": + for _ in range(10): + try: + return func(*args, **kwargs) + except Exception: + time.sleep(0.1) + return func(*args, **kwargs) + else: + return func(*args, **kwargs) + + +def remove(*args, **kwargs): + return _retry(os.remove, *args, **kwargs) + + +# Backwards compatibility imports +from gitdb.const import ( + NULL_BIN_SHA, + NULL_HEX_SHA +) + +#} END Aliases + +#{ compatibility stuff ... + + +class _RandomAccessBytesIO: + + """Wrapper to provide required functionality in case memory maps cannot or may + not be used. This is only really required in python 2.4""" + __slots__ = '_sio' + + def __init__(self, buf=''): + self._sio = BytesIO(buf) + + def __getattr__(self, attr): + return getattr(self._sio, attr) + + def __len__(self): + return len(self.getvalue()) + + def __getitem__(self, i): + return self.getvalue()[i] + + def __getslice__(self, start, end): + return self.getvalue()[start:end] + + +def byte_ord(b): + """ + Return the integer representation of the byte string. This supports Python + 3 byte arrays as well as standard strings. + """ + try: + return ord(b) + except TypeError: + return b + +#} END compatibility stuff ... + +#{ Routines + + +def make_sha(source=b''): + """A python2.4 workaround for the sha/hashlib module fiasco + + **Note** From the dulwich project """ + try: + return hashlib.sha1(source) + except NameError: + import sha + sha1 = sha.sha(source) + return sha1 + + +def allocate_memory(size): + """:return: a file-protocol accessible memory block of the given size""" + if size == 0: + return _RandomAccessBytesIO(b'') + # END handle empty chunks gracefully + + try: + return mmap.mmap(-1, size) # read-write by default + except OSError: + # setup real memory instead + # this of course may fail if the amount of memory is not available in + # one chunk - would only be the case in python 2.4, being more likely on + # 32 bit systems. + return _RandomAccessBytesIO(b"\0" * size) + # END handle memory allocation + + +def file_contents_ro(fd, stream=False, allow_mmap=True): + """:return: read-only contents of the file represented by the file descriptor fd + + :param fd: file descriptor opened for reading + :param stream: if False, random access is provided, otherwise the stream interface + is provided. + :param allow_mmap: if True, its allowed to map the contents into memory, which + allows large files to be handled and accessed efficiently. The file-descriptor + will change its position if this is False""" + try: + if allow_mmap: + # supports stream and random access + try: + return mmap.mmap(fd, 0, access=mmap.ACCESS_READ) + except OSError: + # python 2.4 issue, 0 wants to be the actual size + return mmap.mmap(fd, os.fstat(fd).st_size, access=mmap.ACCESS_READ) + # END handle python 2.4 + except OSError: + pass + # END exception handling + + # read manually + contents = os.read(fd, os.fstat(fd).st_size) + if stream: + return _RandomAccessBytesIO(contents) + return contents + + +def file_contents_ro_filepath(filepath, stream=False, allow_mmap=True, flags=0): + """Get the file contents at filepath as fast as possible + + :return: random access compatible memory of the given filepath + :param stream: see ``file_contents_ro`` + :param allow_mmap: see ``file_contents_ro`` + :param flags: additional flags to pass to os.open + :raise OSError: If the file could not be opened + + **Note** for now we don't try to use O_NOATIME directly as the right value needs to be + shared per database in fact. It only makes a real difference for loose object + databases anyway, and they use it with the help of the ``flags`` parameter""" + fd = os.open(filepath, os.O_RDONLY | getattr(os, 'O_BINARY', 0) | flags) + try: + return file_contents_ro(fd, stream, allow_mmap) + finally: + close(fd) + # END assure file is closed + + +def sliding_ro_buffer(filepath, flags=0): + """ + :return: a buffer compatible object which uses our mapped memory manager internally + ready to read the whole given filepath""" + return SlidingWindowMapBuffer(mman.make_cursor(filepath), flags=flags) + + +def to_hex_sha(sha): + """:return: hexified version of sha""" + if len(sha) == 40: + return sha + return bin_to_hex(sha) + + +def to_bin_sha(sha): + if len(sha) == 20: + return sha + return hex_to_bin(sha) + + +#} END routines + + +#{ Utilities + +class LazyMixin: + + """ + Base class providing an interface to lazily retrieve attribute values upon + first access. If slots are used, memory will only be reserved once the attribute + is actually accessed and retrieved the first time. All future accesses will + return the cached value as stored in the Instance's dict or slot. + """ + + __slots__ = tuple() + + def __getattr__(self, attr): + """ + Whenever an attribute is requested that we do not know, we allow it + to be created and set. Next time the same attribute is requested, it is simply + returned from our dict/slots. """ + self._set_cache_(attr) + # will raise in case the cache was not created + return object.__getattribute__(self, attr) + + def _set_cache_(self, attr): + """ + This method should be overridden in the derived class. + It should check whether the attribute named by attr can be created + and cached. Do nothing if you do not know the attribute or call your subclass + + The derived class may create as many additional attributes as it deems + necessary in case a git command returns more information than represented + in the single attribute.""" + pass + + +class LockedFD: + + """ + This class facilitates a safe read and write operation to a file on disk. + If we write to 'file', we obtain a lock file at 'file.lock' and write to + that instead. If we succeed, the lock file will be renamed to overwrite + the original file. + + When reading, we obtain a lock file, but to prevent other writers from + succeeding while we are reading the file. + + This type handles error correctly in that it will assure a consistent state + on destruction. + + **note** with this setup, parallel reading is not possible""" + __slots__ = ("_filepath", '_fd', '_write') + + def __init__(self, filepath): + """Initialize an instance with the givne filepath""" + self._filepath = filepath + self._fd = None + self._write = None # if True, we write a file + + def __del__(self): + # will do nothing if the file descriptor is already closed + if self._fd is not None: + self.rollback() + + def _lockfilepath(self): + return "%s.lock" % self._filepath + + def open(self, write=False, stream=False): + """ + Open the file descriptor for reading or writing, both in binary mode. + + :param write: if True, the file descriptor will be opened for writing. Other + wise it will be opened read-only. + :param stream: if True, the file descriptor will be wrapped into a simple stream + object which supports only reading or writing + :return: fd to read from or write to. It is still maintained by this instance + and must not be closed directly + :raise IOError: if the lock could not be retrieved + :raise OSError: If the actual file could not be opened for reading + + **note** must only be called once""" + if self._write is not None: + raise AssertionError("Called %s multiple times" % self.open) + + self._write = write + + # try to open the lock file + binary = getattr(os, 'O_BINARY', 0) + lockmode = os.O_WRONLY | os.O_CREAT | os.O_EXCL | binary + try: + fd = os.open(self._lockfilepath(), lockmode, int("600", 8)) + if not write: + os.close(fd) + else: + self._fd = fd + # END handle file descriptor + except OSError as e: + raise OSError("Lock at %r could not be obtained" % self._lockfilepath()) from e + # END handle lock retrieval + + # open actual file if required + if self._fd is None: + # we could specify exclusive here, as we obtained the lock anyway + try: + self._fd = os.open(self._filepath, os.O_RDONLY | binary) + except: + # assure we release our lockfile + remove(self._lockfilepath()) + raise + # END handle lockfile + # END open descriptor for reading + + if stream: + # need delayed import + from gitdb.stream import FDStream + return FDStream(self._fd) + else: + return self._fd + # END handle stream + + def commit(self): + """When done writing, call this function to commit your changes into the + actual file. + The file descriptor will be closed, and the lockfile handled. + + **Note** can be called multiple times""" + self._end_writing(successful=True) + + def rollback(self): + """Abort your operation without any changes. The file descriptor will be + closed, and the lock released. + + **Note** can be called multiple times""" + self._end_writing(successful=False) + + def _end_writing(self, successful=True): + """Handle the lock according to the write mode """ + if self._write is None: + raise AssertionError("Cannot end operation if it wasn't started yet") + + if self._fd is None: + return + + os.close(self._fd) + self._fd = None + + lockfile = self._lockfilepath() + if self._write and successful: + # on windows, rename does not silently overwrite the existing one + if sys.platform == "win32": + if isfile(self._filepath): + remove(self._filepath) + # END remove if exists + # END win32 special handling + os.rename(lockfile, self._filepath) + + # assure others can at least read the file - the tmpfile left it at rw-- + # We may also write that file, on windows that boils down to a remove- + # protection as well + chmod(self._filepath, int("644", 8)) + else: + # just delete the file so far, we failed + remove(lockfile) + # END successful handling + +#} END utilities diff --git a/lib/python3.12/site-packages/gitdb/utils/__init__.py b/lib/python3.12/site-packages/gitdb/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/gitdb/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f69a2d64b0dbdf43a554544e32ce8b88c0762aac Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/utils/__pycache__/encoding.cpython-312.pyc b/lib/python3.12/site-packages/gitdb/utils/__pycache__/encoding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a394f1a88bb4a213e24cef3707a622234696c83 Binary files /dev/null and b/lib/python3.12/site-packages/gitdb/utils/__pycache__/encoding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/gitdb/utils/encoding.py b/lib/python3.12/site-packages/gitdb/utils/encoding.py new file mode 100644 index 0000000000000000000000000000000000000000..b534ef7633e0c21b74a41b9b922083990804a2d8 --- /dev/null +++ b/lib/python3.12/site-packages/gitdb/utils/encoding.py @@ -0,0 +1,18 @@ +def force_bytes(data, encoding="utf-8"): + if isinstance(data, bytes): + return data + + if isinstance(data, str): + return data.encode(encoding) + + return data + + +def force_text(data, encoding="utf-8"): + if isinstance(data, str): + return data + + if isinstance(data, bytes): + return data.decode(encoding) + + return str(data, encoding) diff --git a/lib/python3.12/site-packages/huggingface_hub/__init__.py b/lib/python3.12/site-packages/huggingface_hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..185dc47e489a7176c27e78974028b3f3f0d6b73d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/__init__.py @@ -0,0 +1,1446 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# *********** +# `huggingface_hub` init has 2 modes: +# - Normal usage: +# If imported to use it, all modules and functions are lazy-loaded. This means +# they exist at top level in module but are imported only the first time they are +# used. This way, `from huggingface_hub import something` will import `something` +# quickly without the hassle of importing all the features from `huggingface_hub`. +# - Static check: +# If statically analyzed, all modules and functions are loaded normally. This way +# static typing check works properly as well as autocomplete in text editors and +# IDEs. +# +# The static model imports are done inside the `if TYPE_CHECKING:` statement at +# the bottom of this file. Since module/functions imports are duplicated, it is +# mandatory to make sure to add them twice when adding one. This is checked in the +# `make quality` command. +# +# To update the static imports, please run the following command and commit the changes. +# ``` +# # Use script +# python utils/check_static_imports.py --update-file +# +# # Or run style on codebase +# make style +# ``` +# +# *********** +# Lazy loader vendored from https://github.com/scientific-python/lazy_loader +import importlib +import os +import sys +from typing import TYPE_CHECKING + + +__version__ = "0.31.4" + +# Alphabetical order of definitions is ensured in tests +# WARNING: any comment added in this dictionary definition will be lost when +# re-generating the file ! +_SUBMOD_ATTRS = { + "_commit_scheduler": [ + "CommitScheduler", + ], + "_inference_endpoints": [ + "InferenceEndpoint", + "InferenceEndpointError", + "InferenceEndpointStatus", + "InferenceEndpointTimeoutError", + "InferenceEndpointType", + ], + "_login": [ + "auth_list", + "auth_switch", + "interpreter_login", + "login", + "logout", + "notebook_login", + ], + "_snapshot_download": [ + "snapshot_download", + ], + "_space_api": [ + "SpaceHardware", + "SpaceRuntime", + "SpaceStage", + "SpaceStorage", + "SpaceVariable", + ], + "_tensorboard_logger": [ + "HFSummaryWriter", + ], + "_webhooks_payload": [ + "WebhookPayload", + "WebhookPayloadComment", + "WebhookPayloadDiscussion", + "WebhookPayloadDiscussionChanges", + "WebhookPayloadEvent", + "WebhookPayloadMovedTo", + "WebhookPayloadRepo", + "WebhookPayloadUrl", + "WebhookPayloadWebhook", + ], + "_webhooks_server": [ + "WebhooksServer", + "webhook_endpoint", + ], + "community": [ + "Discussion", + "DiscussionComment", + "DiscussionCommit", + "DiscussionEvent", + "DiscussionStatusChange", + "DiscussionTitleChange", + "DiscussionWithDetails", + ], + "constants": [ + "CONFIG_NAME", + "FLAX_WEIGHTS_NAME", + "HUGGINGFACE_CO_URL_HOME", + "HUGGINGFACE_CO_URL_TEMPLATE", + "PYTORCH_WEIGHTS_NAME", + "REPO_TYPE_DATASET", + "REPO_TYPE_MODEL", + "REPO_TYPE_SPACE", + "TF2_WEIGHTS_NAME", + "TF_WEIGHTS_NAME", + ], + "fastai_utils": [ + "_save_pretrained_fastai", + "from_pretrained_fastai", + "push_to_hub_fastai", + ], + "file_download": [ + "HfFileMetadata", + "_CACHED_NO_EXIST", + "get_hf_file_metadata", + "hf_hub_download", + "hf_hub_url", + "try_to_load_from_cache", + ], + "hf_api": [ + "Collection", + "CollectionItem", + "CommitInfo", + "CommitOperation", + "CommitOperationAdd", + "CommitOperationCopy", + "CommitOperationDelete", + "DatasetInfo", + "GitCommitInfo", + "GitRefInfo", + "GitRefs", + "HfApi", + "ModelInfo", + "RepoUrl", + "SpaceInfo", + "User", + "UserLikes", + "WebhookInfo", + "WebhookWatchedItem", + "accept_access_request", + "add_collection_item", + "add_space_secret", + "add_space_variable", + "auth_check", + "cancel_access_request", + "change_discussion_status", + "comment_discussion", + "create_branch", + "create_collection", + "create_commit", + "create_discussion", + "create_inference_endpoint", + "create_inference_endpoint_from_catalog", + "create_pull_request", + "create_repo", + "create_tag", + "create_webhook", + "dataset_info", + "delete_branch", + "delete_collection", + "delete_collection_item", + "delete_file", + "delete_folder", + "delete_inference_endpoint", + "delete_repo", + "delete_space_secret", + "delete_space_storage", + "delete_space_variable", + "delete_tag", + "delete_webhook", + "disable_webhook", + "duplicate_space", + "edit_discussion_comment", + "enable_webhook", + "file_exists", + "get_collection", + "get_dataset_tags", + "get_discussion_details", + "get_full_repo_name", + "get_inference_endpoint", + "get_model_tags", + "get_paths_info", + "get_repo_discussions", + "get_safetensors_metadata", + "get_space_runtime", + "get_space_variables", + "get_token_permission", + "get_user_overview", + "get_webhook", + "grant_access", + "list_accepted_access_requests", + "list_collections", + "list_datasets", + "list_inference_catalog", + "list_inference_endpoints", + "list_lfs_files", + "list_liked_repos", + "list_models", + "list_organization_members", + "list_papers", + "list_pending_access_requests", + "list_rejected_access_requests", + "list_repo_commits", + "list_repo_files", + "list_repo_likers", + "list_repo_refs", + "list_repo_tree", + "list_spaces", + "list_user_followers", + "list_user_following", + "list_webhooks", + "merge_pull_request", + "model_info", + "move_repo", + "paper_info", + "parse_safetensors_file_metadata", + "pause_inference_endpoint", + "pause_space", + "permanently_delete_lfs_files", + "preupload_lfs_files", + "reject_access_request", + "rename_discussion", + "repo_exists", + "repo_info", + "repo_type_and_id_from_hf_id", + "request_space_hardware", + "request_space_storage", + "restart_space", + "resume_inference_endpoint", + "revision_exists", + "run_as_future", + "scale_to_zero_inference_endpoint", + "set_space_sleep_time", + "space_info", + "super_squash_history", + "unlike", + "update_collection_item", + "update_collection_metadata", + "update_inference_endpoint", + "update_repo_settings", + "update_repo_visibility", + "update_webhook", + "upload_file", + "upload_folder", + "upload_large_folder", + "whoami", + ], + "hf_file_system": [ + "HfFileSystem", + "HfFileSystemFile", + "HfFileSystemResolvedPath", + "HfFileSystemStreamFile", + ], + "hub_mixin": [ + "ModelHubMixin", + "PyTorchModelHubMixin", + ], + "inference._client": [ + "InferenceClient", + "InferenceTimeoutError", + ], + "inference._generated._async_client": [ + "AsyncInferenceClient", + ], + "inference._generated.types": [ + "AudioClassificationInput", + "AudioClassificationOutputElement", + "AudioClassificationOutputTransform", + "AudioClassificationParameters", + "AudioToAudioInput", + "AudioToAudioOutputElement", + "AutomaticSpeechRecognitionEarlyStoppingEnum", + "AutomaticSpeechRecognitionGenerationParameters", + "AutomaticSpeechRecognitionInput", + "AutomaticSpeechRecognitionOutput", + "AutomaticSpeechRecognitionOutputChunk", + "AutomaticSpeechRecognitionParameters", + "ChatCompletionInput", + "ChatCompletionInputFunctionDefinition", + "ChatCompletionInputFunctionName", + "ChatCompletionInputGrammarType", + "ChatCompletionInputGrammarTypeType", + "ChatCompletionInputMessage", + "ChatCompletionInputMessageChunk", + "ChatCompletionInputMessageChunkType", + "ChatCompletionInputStreamOptions", + "ChatCompletionInputTool", + "ChatCompletionInputToolCall", + "ChatCompletionInputToolChoiceClass", + "ChatCompletionInputToolChoiceEnum", + "ChatCompletionInputURL", + "ChatCompletionOutput", + "ChatCompletionOutputComplete", + "ChatCompletionOutputFunctionDefinition", + "ChatCompletionOutputLogprob", + "ChatCompletionOutputLogprobs", + "ChatCompletionOutputMessage", + "ChatCompletionOutputToolCall", + "ChatCompletionOutputTopLogprob", + "ChatCompletionOutputUsage", + "ChatCompletionStreamOutput", + "ChatCompletionStreamOutputChoice", + "ChatCompletionStreamOutputDelta", + "ChatCompletionStreamOutputDeltaToolCall", + "ChatCompletionStreamOutputFunction", + "ChatCompletionStreamOutputLogprob", + "ChatCompletionStreamOutputLogprobs", + "ChatCompletionStreamOutputTopLogprob", + "ChatCompletionStreamOutputUsage", + "DepthEstimationInput", + "DepthEstimationOutput", + "DocumentQuestionAnsweringInput", + "DocumentQuestionAnsweringInputData", + "DocumentQuestionAnsweringOutputElement", + "DocumentQuestionAnsweringParameters", + "FeatureExtractionInput", + "FeatureExtractionInputTruncationDirection", + "FillMaskInput", + "FillMaskOutputElement", + "FillMaskParameters", + "ImageClassificationInput", + "ImageClassificationOutputElement", + "ImageClassificationOutputTransform", + "ImageClassificationParameters", + "ImageSegmentationInput", + "ImageSegmentationOutputElement", + "ImageSegmentationParameters", + "ImageSegmentationSubtask", + "ImageToImageInput", + "ImageToImageOutput", + "ImageToImageParameters", + "ImageToImageTargetSize", + "ImageToTextEarlyStoppingEnum", + "ImageToTextGenerationParameters", + "ImageToTextInput", + "ImageToTextOutput", + "ImageToTextParameters", + "ObjectDetectionBoundingBox", + "ObjectDetectionInput", + "ObjectDetectionOutputElement", + "ObjectDetectionParameters", + "Padding", + "QuestionAnsweringInput", + "QuestionAnsweringInputData", + "QuestionAnsweringOutputElement", + "QuestionAnsweringParameters", + "SentenceSimilarityInput", + "SentenceSimilarityInputData", + "SummarizationInput", + "SummarizationOutput", + "SummarizationParameters", + "SummarizationTruncationStrategy", + "TableQuestionAnsweringInput", + "TableQuestionAnsweringInputData", + "TableQuestionAnsweringOutputElement", + "TableQuestionAnsweringParameters", + "Text2TextGenerationInput", + "Text2TextGenerationOutput", + "Text2TextGenerationParameters", + "Text2TextGenerationTruncationStrategy", + "TextClassificationInput", + "TextClassificationOutputElement", + "TextClassificationOutputTransform", + "TextClassificationParameters", + "TextGenerationInput", + "TextGenerationInputGenerateParameters", + "TextGenerationInputGrammarType", + "TextGenerationOutput", + "TextGenerationOutputBestOfSequence", + "TextGenerationOutputDetails", + "TextGenerationOutputFinishReason", + "TextGenerationOutputPrefillToken", + "TextGenerationOutputToken", + "TextGenerationStreamOutput", + "TextGenerationStreamOutputStreamDetails", + "TextGenerationStreamOutputToken", + "TextToAudioEarlyStoppingEnum", + "TextToAudioGenerationParameters", + "TextToAudioInput", + "TextToAudioOutput", + "TextToAudioParameters", + "TextToImageInput", + "TextToImageOutput", + "TextToImageParameters", + "TextToSpeechEarlyStoppingEnum", + "TextToSpeechGenerationParameters", + "TextToSpeechInput", + "TextToSpeechOutput", + "TextToSpeechParameters", + "TextToVideoInput", + "TextToVideoOutput", + "TextToVideoParameters", + "TokenClassificationAggregationStrategy", + "TokenClassificationInput", + "TokenClassificationOutputElement", + "TokenClassificationParameters", + "TranslationInput", + "TranslationOutput", + "TranslationParameters", + "TranslationTruncationStrategy", + "TypeEnum", + "VideoClassificationInput", + "VideoClassificationOutputElement", + "VideoClassificationOutputTransform", + "VideoClassificationParameters", + "VisualQuestionAnsweringInput", + "VisualQuestionAnsweringInputData", + "VisualQuestionAnsweringOutputElement", + "VisualQuestionAnsweringParameters", + "ZeroShotClassificationInput", + "ZeroShotClassificationOutputElement", + "ZeroShotClassificationParameters", + "ZeroShotImageClassificationInput", + "ZeroShotImageClassificationOutputElement", + "ZeroShotImageClassificationParameters", + "ZeroShotObjectDetectionBoundingBox", + "ZeroShotObjectDetectionInput", + "ZeroShotObjectDetectionOutputElement", + "ZeroShotObjectDetectionParameters", + ], + "inference_api": [ + "InferenceApi", + ], + "keras_mixin": [ + "KerasModelHubMixin", + "from_pretrained_keras", + "push_to_hub_keras", + "save_pretrained_keras", + ], + "repocard": [ + "DatasetCard", + "ModelCard", + "RepoCard", + "SpaceCard", + "metadata_eval_result", + "metadata_load", + "metadata_save", + "metadata_update", + ], + "repocard_data": [ + "CardData", + "DatasetCardData", + "EvalResult", + "ModelCardData", + "SpaceCardData", + ], + "repository": [ + "Repository", + ], + "serialization": [ + "StateDictSplit", + "get_tf_storage_size", + "get_torch_storage_id", + "get_torch_storage_size", + "load_state_dict_from_file", + "load_torch_model", + "save_torch_model", + "save_torch_state_dict", + "split_state_dict_into_shards_factory", + "split_tf_state_dict_into_shards", + "split_torch_state_dict_into_shards", + ], + "serialization._dduf": [ + "DDUFEntry", + "export_entries_as_dduf", + "export_folder_as_dduf", + "read_dduf_file", + ], + "utils": [ + "CacheNotFound", + "CachedFileInfo", + "CachedRepoInfo", + "CachedRevisionInfo", + "CorruptedCacheException", + "DeleteCacheStrategy", + "HFCacheInfo", + "HfFolder", + "cached_assets_path", + "configure_http_backend", + "dump_environment_info", + "get_session", + "get_token", + "logging", + "scan_cache_dir", + ], +} + +# WARNING: __all__ is generated automatically, Any manual edit will be lost when re-generating this file ! +# +# To update the static imports, please run the following command and commit the changes. +# ``` +# # Use script +# python utils/check_all_variable.py --update +# +# # Or run style on codebase +# make style +# ``` + +__all__ = [ + "AsyncInferenceClient", + "AudioClassificationInput", + "AudioClassificationOutputElement", + "AudioClassificationOutputTransform", + "AudioClassificationParameters", + "AudioToAudioInput", + "AudioToAudioOutputElement", + "AutomaticSpeechRecognitionEarlyStoppingEnum", + "AutomaticSpeechRecognitionGenerationParameters", + "AutomaticSpeechRecognitionInput", + "AutomaticSpeechRecognitionOutput", + "AutomaticSpeechRecognitionOutputChunk", + "AutomaticSpeechRecognitionParameters", + "CONFIG_NAME", + "CacheNotFound", + "CachedFileInfo", + "CachedRepoInfo", + "CachedRevisionInfo", + "CardData", + "ChatCompletionInput", + "ChatCompletionInputFunctionDefinition", + "ChatCompletionInputFunctionName", + "ChatCompletionInputGrammarType", + "ChatCompletionInputGrammarTypeType", + "ChatCompletionInputMessage", + "ChatCompletionInputMessageChunk", + "ChatCompletionInputMessageChunkType", + "ChatCompletionInputStreamOptions", + "ChatCompletionInputTool", + "ChatCompletionInputToolCall", + "ChatCompletionInputToolChoiceClass", + "ChatCompletionInputToolChoiceEnum", + "ChatCompletionInputURL", + "ChatCompletionOutput", + "ChatCompletionOutputComplete", + "ChatCompletionOutputFunctionDefinition", + "ChatCompletionOutputLogprob", + "ChatCompletionOutputLogprobs", + "ChatCompletionOutputMessage", + "ChatCompletionOutputToolCall", + "ChatCompletionOutputTopLogprob", + "ChatCompletionOutputUsage", + "ChatCompletionStreamOutput", + "ChatCompletionStreamOutputChoice", + "ChatCompletionStreamOutputDelta", + "ChatCompletionStreamOutputDeltaToolCall", + "ChatCompletionStreamOutputFunction", + "ChatCompletionStreamOutputLogprob", + "ChatCompletionStreamOutputLogprobs", + "ChatCompletionStreamOutputTopLogprob", + "ChatCompletionStreamOutputUsage", + "Collection", + "CollectionItem", + "CommitInfo", + "CommitOperation", + "CommitOperationAdd", + "CommitOperationCopy", + "CommitOperationDelete", + "CommitScheduler", + "CorruptedCacheException", + "DDUFEntry", + "DatasetCard", + "DatasetCardData", + "DatasetInfo", + "DeleteCacheStrategy", + "DepthEstimationInput", + "DepthEstimationOutput", + "Discussion", + "DiscussionComment", + "DiscussionCommit", + "DiscussionEvent", + "DiscussionStatusChange", + "DiscussionTitleChange", + "DiscussionWithDetails", + "DocumentQuestionAnsweringInput", + "DocumentQuestionAnsweringInputData", + "DocumentQuestionAnsweringOutputElement", + "DocumentQuestionAnsweringParameters", + "EvalResult", + "FLAX_WEIGHTS_NAME", + "FeatureExtractionInput", + "FeatureExtractionInputTruncationDirection", + "FillMaskInput", + "FillMaskOutputElement", + "FillMaskParameters", + "GitCommitInfo", + "GitRefInfo", + "GitRefs", + "HFCacheInfo", + "HFSummaryWriter", + "HUGGINGFACE_CO_URL_HOME", + "HUGGINGFACE_CO_URL_TEMPLATE", + "HfApi", + "HfFileMetadata", + "HfFileSystem", + "HfFileSystemFile", + "HfFileSystemResolvedPath", + "HfFileSystemStreamFile", + "HfFolder", + "ImageClassificationInput", + "ImageClassificationOutputElement", + "ImageClassificationOutputTransform", + "ImageClassificationParameters", + "ImageSegmentationInput", + "ImageSegmentationOutputElement", + "ImageSegmentationParameters", + "ImageSegmentationSubtask", + "ImageToImageInput", + "ImageToImageOutput", + "ImageToImageParameters", + "ImageToImageTargetSize", + "ImageToTextEarlyStoppingEnum", + "ImageToTextGenerationParameters", + "ImageToTextInput", + "ImageToTextOutput", + "ImageToTextParameters", + "InferenceApi", + "InferenceClient", + "InferenceEndpoint", + "InferenceEndpointError", + "InferenceEndpointStatus", + "InferenceEndpointTimeoutError", + "InferenceEndpointType", + "InferenceTimeoutError", + "KerasModelHubMixin", + "ModelCard", + "ModelCardData", + "ModelHubMixin", + "ModelInfo", + "ObjectDetectionBoundingBox", + "ObjectDetectionInput", + "ObjectDetectionOutputElement", + "ObjectDetectionParameters", + "PYTORCH_WEIGHTS_NAME", + "Padding", + "PyTorchModelHubMixin", + "QuestionAnsweringInput", + "QuestionAnsweringInputData", + "QuestionAnsweringOutputElement", + "QuestionAnsweringParameters", + "REPO_TYPE_DATASET", + "REPO_TYPE_MODEL", + "REPO_TYPE_SPACE", + "RepoCard", + "RepoUrl", + "Repository", + "SentenceSimilarityInput", + "SentenceSimilarityInputData", + "SpaceCard", + "SpaceCardData", + "SpaceHardware", + "SpaceInfo", + "SpaceRuntime", + "SpaceStage", + "SpaceStorage", + "SpaceVariable", + "StateDictSplit", + "SummarizationInput", + "SummarizationOutput", + "SummarizationParameters", + "SummarizationTruncationStrategy", + "TF2_WEIGHTS_NAME", + "TF_WEIGHTS_NAME", + "TableQuestionAnsweringInput", + "TableQuestionAnsweringInputData", + "TableQuestionAnsweringOutputElement", + "TableQuestionAnsweringParameters", + "Text2TextGenerationInput", + "Text2TextGenerationOutput", + "Text2TextGenerationParameters", + "Text2TextGenerationTruncationStrategy", + "TextClassificationInput", + "TextClassificationOutputElement", + "TextClassificationOutputTransform", + "TextClassificationParameters", + "TextGenerationInput", + "TextGenerationInputGenerateParameters", + "TextGenerationInputGrammarType", + "TextGenerationOutput", + "TextGenerationOutputBestOfSequence", + "TextGenerationOutputDetails", + "TextGenerationOutputFinishReason", + "TextGenerationOutputPrefillToken", + "TextGenerationOutputToken", + "TextGenerationStreamOutput", + "TextGenerationStreamOutputStreamDetails", + "TextGenerationStreamOutputToken", + "TextToAudioEarlyStoppingEnum", + "TextToAudioGenerationParameters", + "TextToAudioInput", + "TextToAudioOutput", + "TextToAudioParameters", + "TextToImageInput", + "TextToImageOutput", + "TextToImageParameters", + "TextToSpeechEarlyStoppingEnum", + "TextToSpeechGenerationParameters", + "TextToSpeechInput", + "TextToSpeechOutput", + "TextToSpeechParameters", + "TextToVideoInput", + "TextToVideoOutput", + "TextToVideoParameters", + "TokenClassificationAggregationStrategy", + "TokenClassificationInput", + "TokenClassificationOutputElement", + "TokenClassificationParameters", + "TranslationInput", + "TranslationOutput", + "TranslationParameters", + "TranslationTruncationStrategy", + "TypeEnum", + "User", + "UserLikes", + "VideoClassificationInput", + "VideoClassificationOutputElement", + "VideoClassificationOutputTransform", + "VideoClassificationParameters", + "VisualQuestionAnsweringInput", + "VisualQuestionAnsweringInputData", + "VisualQuestionAnsweringOutputElement", + "VisualQuestionAnsweringParameters", + "WebhookInfo", + "WebhookPayload", + "WebhookPayloadComment", + "WebhookPayloadDiscussion", + "WebhookPayloadDiscussionChanges", + "WebhookPayloadEvent", + "WebhookPayloadMovedTo", + "WebhookPayloadRepo", + "WebhookPayloadUrl", + "WebhookPayloadWebhook", + "WebhookWatchedItem", + "WebhooksServer", + "ZeroShotClassificationInput", + "ZeroShotClassificationOutputElement", + "ZeroShotClassificationParameters", + "ZeroShotImageClassificationInput", + "ZeroShotImageClassificationOutputElement", + "ZeroShotImageClassificationParameters", + "ZeroShotObjectDetectionBoundingBox", + "ZeroShotObjectDetectionInput", + "ZeroShotObjectDetectionOutputElement", + "ZeroShotObjectDetectionParameters", + "_CACHED_NO_EXIST", + "_save_pretrained_fastai", + "accept_access_request", + "add_collection_item", + "add_space_secret", + "add_space_variable", + "auth_check", + "auth_list", + "auth_switch", + "cached_assets_path", + "cancel_access_request", + "change_discussion_status", + "comment_discussion", + "configure_http_backend", + "create_branch", + "create_collection", + "create_commit", + "create_discussion", + "create_inference_endpoint", + "create_inference_endpoint_from_catalog", + "create_pull_request", + "create_repo", + "create_tag", + "create_webhook", + "dataset_info", + "delete_branch", + "delete_collection", + "delete_collection_item", + "delete_file", + "delete_folder", + "delete_inference_endpoint", + "delete_repo", + "delete_space_secret", + "delete_space_storage", + "delete_space_variable", + "delete_tag", + "delete_webhook", + "disable_webhook", + "dump_environment_info", + "duplicate_space", + "edit_discussion_comment", + "enable_webhook", + "export_entries_as_dduf", + "export_folder_as_dduf", + "file_exists", + "from_pretrained_fastai", + "from_pretrained_keras", + "get_collection", + "get_dataset_tags", + "get_discussion_details", + "get_full_repo_name", + "get_hf_file_metadata", + "get_inference_endpoint", + "get_model_tags", + "get_paths_info", + "get_repo_discussions", + "get_safetensors_metadata", + "get_session", + "get_space_runtime", + "get_space_variables", + "get_tf_storage_size", + "get_token", + "get_token_permission", + "get_torch_storage_id", + "get_torch_storage_size", + "get_user_overview", + "get_webhook", + "grant_access", + "hf_hub_download", + "hf_hub_url", + "interpreter_login", + "list_accepted_access_requests", + "list_collections", + "list_datasets", + "list_inference_catalog", + "list_inference_endpoints", + "list_lfs_files", + "list_liked_repos", + "list_models", + "list_organization_members", + "list_papers", + "list_pending_access_requests", + "list_rejected_access_requests", + "list_repo_commits", + "list_repo_files", + "list_repo_likers", + "list_repo_refs", + "list_repo_tree", + "list_spaces", + "list_user_followers", + "list_user_following", + "list_webhooks", + "load_state_dict_from_file", + "load_torch_model", + "logging", + "login", + "logout", + "merge_pull_request", + "metadata_eval_result", + "metadata_load", + "metadata_save", + "metadata_update", + "model_info", + "move_repo", + "notebook_login", + "paper_info", + "parse_safetensors_file_metadata", + "pause_inference_endpoint", + "pause_space", + "permanently_delete_lfs_files", + "preupload_lfs_files", + "push_to_hub_fastai", + "push_to_hub_keras", + "read_dduf_file", + "reject_access_request", + "rename_discussion", + "repo_exists", + "repo_info", + "repo_type_and_id_from_hf_id", + "request_space_hardware", + "request_space_storage", + "restart_space", + "resume_inference_endpoint", + "revision_exists", + "run_as_future", + "save_pretrained_keras", + "save_torch_model", + "save_torch_state_dict", + "scale_to_zero_inference_endpoint", + "scan_cache_dir", + "set_space_sleep_time", + "snapshot_download", + "space_info", + "split_state_dict_into_shards_factory", + "split_tf_state_dict_into_shards", + "split_torch_state_dict_into_shards", + "super_squash_history", + "try_to_load_from_cache", + "unlike", + "update_collection_item", + "update_collection_metadata", + "update_inference_endpoint", + "update_repo_settings", + "update_repo_visibility", + "update_webhook", + "upload_file", + "upload_folder", + "upload_large_folder", + "webhook_endpoint", + "whoami", +] + + +def _attach(package_name, submodules=None, submod_attrs=None): + """Attach lazily loaded submodules, functions, or other attributes. + + Typically, modules import submodules and attributes as follows: + + ```py + import mysubmodule + import anothersubmodule + + from .foo import someattr + ``` + + The idea is to replace a package's `__getattr__`, `__dir__`, such that all imports + work exactly the way they would with normal imports, except that the import occurs + upon first use. + + The typical way to call this function, replacing the above imports, is: + + ```python + __getattr__, __dir__ = lazy.attach( + __name__, + ['mysubmodule', 'anothersubmodule'], + {'foo': ['someattr']} + ) + ``` + This functionality requires Python 3.7 or higher. + + Args: + package_name (`str`): + Typically use `__name__`. + submodules (`set`): + List of submodules to attach. + submod_attrs (`dict`): + Dictionary of submodule -> list of attributes / functions. + These attributes are imported as they are used. + + Returns: + __getattr__, __dir__, __all__ + + """ + if submod_attrs is None: + submod_attrs = {} + + if submodules is None: + submodules = set() + else: + submodules = set(submodules) + + attr_to_modules = {attr: mod for mod, attrs in submod_attrs.items() for attr in attrs} + + def __getattr__(name): + if name in submodules: + try: + return importlib.import_module(f"{package_name}.{name}") + except Exception as e: + print(f"Error importing {package_name}.{name}: {e}") + raise + elif name in attr_to_modules: + submod_path = f"{package_name}.{attr_to_modules[name]}" + try: + submod = importlib.import_module(submod_path) + except Exception as e: + print(f"Error importing {submod_path}: {e}") + raise + attr = getattr(submod, name) + + # If the attribute lives in a file (module) with the same + # name as the attribute, ensure that the attribute and *not* + # the module is accessible on the package. + if name == attr_to_modules[name]: + pkg = sys.modules[package_name] + pkg.__dict__[name] = attr + + return attr + else: + raise AttributeError(f"No {package_name} attribute {name}") + + def __dir__(): + return __all__ + + return __getattr__, __dir__ + + +__getattr__, __dir__ = _attach(__name__, submodules=[], submod_attrs=_SUBMOD_ATTRS) + +if os.environ.get("EAGER_IMPORT", ""): + for attr in __all__: + __getattr__(attr) + +# WARNING: any content below this statement is generated automatically. Any manual edit +# will be lost when re-generating this file ! +# +# To update the static imports, please run the following command and commit the changes. +# ``` +# # Use script +# python utils/check_static_imports.py --update +# +# # Or run style on codebase +# make style +# ``` +if TYPE_CHECKING: # pragma: no cover + from ._commit_scheduler import CommitScheduler # noqa: F401 + from ._inference_endpoints import ( + InferenceEndpoint, # noqa: F401 + InferenceEndpointError, # noqa: F401 + InferenceEndpointStatus, # noqa: F401 + InferenceEndpointTimeoutError, # noqa: F401 + InferenceEndpointType, # noqa: F401 + ) + from ._login import ( + auth_list, # noqa: F401 + auth_switch, # noqa: F401 + interpreter_login, # noqa: F401 + login, # noqa: F401 + logout, # noqa: F401 + notebook_login, # noqa: F401 + ) + from ._snapshot_download import snapshot_download # noqa: F401 + from ._space_api import ( + SpaceHardware, # noqa: F401 + SpaceRuntime, # noqa: F401 + SpaceStage, # noqa: F401 + SpaceStorage, # noqa: F401 + SpaceVariable, # noqa: F401 + ) + from ._tensorboard_logger import HFSummaryWriter # noqa: F401 + from ._webhooks_payload import ( + WebhookPayload, # noqa: F401 + WebhookPayloadComment, # noqa: F401 + WebhookPayloadDiscussion, # noqa: F401 + WebhookPayloadDiscussionChanges, # noqa: F401 + WebhookPayloadEvent, # noqa: F401 + WebhookPayloadMovedTo, # noqa: F401 + WebhookPayloadRepo, # noqa: F401 + WebhookPayloadUrl, # noqa: F401 + WebhookPayloadWebhook, # noqa: F401 + ) + from ._webhooks_server import ( + WebhooksServer, # noqa: F401 + webhook_endpoint, # noqa: F401 + ) + from .community import ( + Discussion, # noqa: F401 + DiscussionComment, # noqa: F401 + DiscussionCommit, # noqa: F401 + DiscussionEvent, # noqa: F401 + DiscussionStatusChange, # noqa: F401 + DiscussionTitleChange, # noqa: F401 + DiscussionWithDetails, # noqa: F401 + ) + from .constants import ( + CONFIG_NAME, # noqa: F401 + FLAX_WEIGHTS_NAME, # noqa: F401 + HUGGINGFACE_CO_URL_HOME, # noqa: F401 + HUGGINGFACE_CO_URL_TEMPLATE, # noqa: F401 + PYTORCH_WEIGHTS_NAME, # noqa: F401 + REPO_TYPE_DATASET, # noqa: F401 + REPO_TYPE_MODEL, # noqa: F401 + REPO_TYPE_SPACE, # noqa: F401 + TF2_WEIGHTS_NAME, # noqa: F401 + TF_WEIGHTS_NAME, # noqa: F401 + ) + from .fastai_utils import ( + _save_pretrained_fastai, # noqa: F401 + from_pretrained_fastai, # noqa: F401 + push_to_hub_fastai, # noqa: F401 + ) + from .file_download import ( + _CACHED_NO_EXIST, # noqa: F401 + HfFileMetadata, # noqa: F401 + get_hf_file_metadata, # noqa: F401 + hf_hub_download, # noqa: F401 + hf_hub_url, # noqa: F401 + try_to_load_from_cache, # noqa: F401 + ) + from .hf_api import ( + Collection, # noqa: F401 + CollectionItem, # noqa: F401 + CommitInfo, # noqa: F401 + CommitOperation, # noqa: F401 + CommitOperationAdd, # noqa: F401 + CommitOperationCopy, # noqa: F401 + CommitOperationDelete, # noqa: F401 + DatasetInfo, # noqa: F401 + GitCommitInfo, # noqa: F401 + GitRefInfo, # noqa: F401 + GitRefs, # noqa: F401 + HfApi, # noqa: F401 + ModelInfo, # noqa: F401 + RepoUrl, # noqa: F401 + SpaceInfo, # noqa: F401 + User, # noqa: F401 + UserLikes, # noqa: F401 + WebhookInfo, # noqa: F401 + WebhookWatchedItem, # noqa: F401 + accept_access_request, # noqa: F401 + add_collection_item, # noqa: F401 + add_space_secret, # noqa: F401 + add_space_variable, # noqa: F401 + auth_check, # noqa: F401 + cancel_access_request, # noqa: F401 + change_discussion_status, # noqa: F401 + comment_discussion, # noqa: F401 + create_branch, # noqa: F401 + create_collection, # noqa: F401 + create_commit, # noqa: F401 + create_discussion, # noqa: F401 + create_inference_endpoint, # noqa: F401 + create_inference_endpoint_from_catalog, # noqa: F401 + create_pull_request, # noqa: F401 + create_repo, # noqa: F401 + create_tag, # noqa: F401 + create_webhook, # noqa: F401 + dataset_info, # noqa: F401 + delete_branch, # noqa: F401 + delete_collection, # noqa: F401 + delete_collection_item, # noqa: F401 + delete_file, # noqa: F401 + delete_folder, # noqa: F401 + delete_inference_endpoint, # noqa: F401 + delete_repo, # noqa: F401 + delete_space_secret, # noqa: F401 + delete_space_storage, # noqa: F401 + delete_space_variable, # noqa: F401 + delete_tag, # noqa: F401 + delete_webhook, # noqa: F401 + disable_webhook, # noqa: F401 + duplicate_space, # noqa: F401 + edit_discussion_comment, # noqa: F401 + enable_webhook, # noqa: F401 + file_exists, # noqa: F401 + get_collection, # noqa: F401 + get_dataset_tags, # noqa: F401 + get_discussion_details, # noqa: F401 + get_full_repo_name, # noqa: F401 + get_inference_endpoint, # noqa: F401 + get_model_tags, # noqa: F401 + get_paths_info, # noqa: F401 + get_repo_discussions, # noqa: F401 + get_safetensors_metadata, # noqa: F401 + get_space_runtime, # noqa: F401 + get_space_variables, # noqa: F401 + get_token_permission, # noqa: F401 + get_user_overview, # noqa: F401 + get_webhook, # noqa: F401 + grant_access, # noqa: F401 + list_accepted_access_requests, # noqa: F401 + list_collections, # noqa: F401 + list_datasets, # noqa: F401 + list_inference_catalog, # noqa: F401 + list_inference_endpoints, # noqa: F401 + list_lfs_files, # noqa: F401 + list_liked_repos, # noqa: F401 + list_models, # noqa: F401 + list_organization_members, # noqa: F401 + list_papers, # noqa: F401 + list_pending_access_requests, # noqa: F401 + list_rejected_access_requests, # noqa: F401 + list_repo_commits, # noqa: F401 + list_repo_files, # noqa: F401 + list_repo_likers, # noqa: F401 + list_repo_refs, # noqa: F401 + list_repo_tree, # noqa: F401 + list_spaces, # noqa: F401 + list_user_followers, # noqa: F401 + list_user_following, # noqa: F401 + list_webhooks, # noqa: F401 + merge_pull_request, # noqa: F401 + model_info, # noqa: F401 + move_repo, # noqa: F401 + paper_info, # noqa: F401 + parse_safetensors_file_metadata, # noqa: F401 + pause_inference_endpoint, # noqa: F401 + pause_space, # noqa: F401 + permanently_delete_lfs_files, # noqa: F401 + preupload_lfs_files, # noqa: F401 + reject_access_request, # noqa: F401 + rename_discussion, # noqa: F401 + repo_exists, # noqa: F401 + repo_info, # noqa: F401 + repo_type_and_id_from_hf_id, # noqa: F401 + request_space_hardware, # noqa: F401 + request_space_storage, # noqa: F401 + restart_space, # noqa: F401 + resume_inference_endpoint, # noqa: F401 + revision_exists, # noqa: F401 + run_as_future, # noqa: F401 + scale_to_zero_inference_endpoint, # noqa: F401 + set_space_sleep_time, # noqa: F401 + space_info, # noqa: F401 + super_squash_history, # noqa: F401 + unlike, # noqa: F401 + update_collection_item, # noqa: F401 + update_collection_metadata, # noqa: F401 + update_inference_endpoint, # noqa: F401 + update_repo_settings, # noqa: F401 + update_repo_visibility, # noqa: F401 + update_webhook, # noqa: F401 + upload_file, # noqa: F401 + upload_folder, # noqa: F401 + upload_large_folder, # noqa: F401 + whoami, # noqa: F401 + ) + from .hf_file_system import ( + HfFileSystem, # noqa: F401 + HfFileSystemFile, # noqa: F401 + HfFileSystemResolvedPath, # noqa: F401 + HfFileSystemStreamFile, # noqa: F401 + ) + from .hub_mixin import ( + ModelHubMixin, # noqa: F401 + PyTorchModelHubMixin, # noqa: F401 + ) + from .inference._client import ( + InferenceClient, # noqa: F401 + InferenceTimeoutError, # noqa: F401 + ) + from .inference._generated._async_client import AsyncInferenceClient # noqa: F401 + from .inference._generated.types import ( + AudioClassificationInput, # noqa: F401 + AudioClassificationOutputElement, # noqa: F401 + AudioClassificationOutputTransform, # noqa: F401 + AudioClassificationParameters, # noqa: F401 + AudioToAudioInput, # noqa: F401 + AudioToAudioOutputElement, # noqa: F401 + AutomaticSpeechRecognitionEarlyStoppingEnum, # noqa: F401 + AutomaticSpeechRecognitionGenerationParameters, # noqa: F401 + AutomaticSpeechRecognitionInput, # noqa: F401 + AutomaticSpeechRecognitionOutput, # noqa: F401 + AutomaticSpeechRecognitionOutputChunk, # noqa: F401 + AutomaticSpeechRecognitionParameters, # noqa: F401 + ChatCompletionInput, # noqa: F401 + ChatCompletionInputFunctionDefinition, # noqa: F401 + ChatCompletionInputFunctionName, # noqa: F401 + ChatCompletionInputGrammarType, # noqa: F401 + ChatCompletionInputGrammarTypeType, # noqa: F401 + ChatCompletionInputMessage, # noqa: F401 + ChatCompletionInputMessageChunk, # noqa: F401 + ChatCompletionInputMessageChunkType, # noqa: F401 + ChatCompletionInputStreamOptions, # noqa: F401 + ChatCompletionInputTool, # noqa: F401 + ChatCompletionInputToolCall, # noqa: F401 + ChatCompletionInputToolChoiceClass, # noqa: F401 + ChatCompletionInputToolChoiceEnum, # noqa: F401 + ChatCompletionInputURL, # noqa: F401 + ChatCompletionOutput, # noqa: F401 + ChatCompletionOutputComplete, # noqa: F401 + ChatCompletionOutputFunctionDefinition, # noqa: F401 + ChatCompletionOutputLogprob, # noqa: F401 + ChatCompletionOutputLogprobs, # noqa: F401 + ChatCompletionOutputMessage, # noqa: F401 + ChatCompletionOutputToolCall, # noqa: F401 + ChatCompletionOutputTopLogprob, # noqa: F401 + ChatCompletionOutputUsage, # noqa: F401 + ChatCompletionStreamOutput, # noqa: F401 + ChatCompletionStreamOutputChoice, # noqa: F401 + ChatCompletionStreamOutputDelta, # noqa: F401 + ChatCompletionStreamOutputDeltaToolCall, # noqa: F401 + ChatCompletionStreamOutputFunction, # noqa: F401 + ChatCompletionStreamOutputLogprob, # noqa: F401 + ChatCompletionStreamOutputLogprobs, # noqa: F401 + ChatCompletionStreamOutputTopLogprob, # noqa: F401 + ChatCompletionStreamOutputUsage, # noqa: F401 + DepthEstimationInput, # noqa: F401 + DepthEstimationOutput, # noqa: F401 + DocumentQuestionAnsweringInput, # noqa: F401 + DocumentQuestionAnsweringInputData, # noqa: F401 + DocumentQuestionAnsweringOutputElement, # noqa: F401 + DocumentQuestionAnsweringParameters, # noqa: F401 + FeatureExtractionInput, # noqa: F401 + FeatureExtractionInputTruncationDirection, # noqa: F401 + FillMaskInput, # noqa: F401 + FillMaskOutputElement, # noqa: F401 + FillMaskParameters, # noqa: F401 + ImageClassificationInput, # noqa: F401 + ImageClassificationOutputElement, # noqa: F401 + ImageClassificationOutputTransform, # noqa: F401 + ImageClassificationParameters, # noqa: F401 + ImageSegmentationInput, # noqa: F401 + ImageSegmentationOutputElement, # noqa: F401 + ImageSegmentationParameters, # noqa: F401 + ImageSegmentationSubtask, # noqa: F401 + ImageToImageInput, # noqa: F401 + ImageToImageOutput, # noqa: F401 + ImageToImageParameters, # noqa: F401 + ImageToImageTargetSize, # noqa: F401 + ImageToTextEarlyStoppingEnum, # noqa: F401 + ImageToTextGenerationParameters, # noqa: F401 + ImageToTextInput, # noqa: F401 + ImageToTextOutput, # noqa: F401 + ImageToTextParameters, # noqa: F401 + ObjectDetectionBoundingBox, # noqa: F401 + ObjectDetectionInput, # noqa: F401 + ObjectDetectionOutputElement, # noqa: F401 + ObjectDetectionParameters, # noqa: F401 + Padding, # noqa: F401 + QuestionAnsweringInput, # noqa: F401 + QuestionAnsweringInputData, # noqa: F401 + QuestionAnsweringOutputElement, # noqa: F401 + QuestionAnsweringParameters, # noqa: F401 + SentenceSimilarityInput, # noqa: F401 + SentenceSimilarityInputData, # noqa: F401 + SummarizationInput, # noqa: F401 + SummarizationOutput, # noqa: F401 + SummarizationParameters, # noqa: F401 + SummarizationTruncationStrategy, # noqa: F401 + TableQuestionAnsweringInput, # noqa: F401 + TableQuestionAnsweringInputData, # noqa: F401 + TableQuestionAnsweringOutputElement, # noqa: F401 + TableQuestionAnsweringParameters, # noqa: F401 + Text2TextGenerationInput, # noqa: F401 + Text2TextGenerationOutput, # noqa: F401 + Text2TextGenerationParameters, # noqa: F401 + Text2TextGenerationTruncationStrategy, # noqa: F401 + TextClassificationInput, # noqa: F401 + TextClassificationOutputElement, # noqa: F401 + TextClassificationOutputTransform, # noqa: F401 + TextClassificationParameters, # noqa: F401 + TextGenerationInput, # noqa: F401 + TextGenerationInputGenerateParameters, # noqa: F401 + TextGenerationInputGrammarType, # noqa: F401 + TextGenerationOutput, # noqa: F401 + TextGenerationOutputBestOfSequence, # noqa: F401 + TextGenerationOutputDetails, # noqa: F401 + TextGenerationOutputFinishReason, # noqa: F401 + TextGenerationOutputPrefillToken, # noqa: F401 + TextGenerationOutputToken, # noqa: F401 + TextGenerationStreamOutput, # noqa: F401 + TextGenerationStreamOutputStreamDetails, # noqa: F401 + TextGenerationStreamOutputToken, # noqa: F401 + TextToAudioEarlyStoppingEnum, # noqa: F401 + TextToAudioGenerationParameters, # noqa: F401 + TextToAudioInput, # noqa: F401 + TextToAudioOutput, # noqa: F401 + TextToAudioParameters, # noqa: F401 + TextToImageInput, # noqa: F401 + TextToImageOutput, # noqa: F401 + TextToImageParameters, # noqa: F401 + TextToSpeechEarlyStoppingEnum, # noqa: F401 + TextToSpeechGenerationParameters, # noqa: F401 + TextToSpeechInput, # noqa: F401 + TextToSpeechOutput, # noqa: F401 + TextToSpeechParameters, # noqa: F401 + TextToVideoInput, # noqa: F401 + TextToVideoOutput, # noqa: F401 + TextToVideoParameters, # noqa: F401 + TokenClassificationAggregationStrategy, # noqa: F401 + TokenClassificationInput, # noqa: F401 + TokenClassificationOutputElement, # noqa: F401 + TokenClassificationParameters, # noqa: F401 + TranslationInput, # noqa: F401 + TranslationOutput, # noqa: F401 + TranslationParameters, # noqa: F401 + TranslationTruncationStrategy, # noqa: F401 + TypeEnum, # noqa: F401 + VideoClassificationInput, # noqa: F401 + VideoClassificationOutputElement, # noqa: F401 + VideoClassificationOutputTransform, # noqa: F401 + VideoClassificationParameters, # noqa: F401 + VisualQuestionAnsweringInput, # noqa: F401 + VisualQuestionAnsweringInputData, # noqa: F401 + VisualQuestionAnsweringOutputElement, # noqa: F401 + VisualQuestionAnsweringParameters, # noqa: F401 + ZeroShotClassificationInput, # noqa: F401 + ZeroShotClassificationOutputElement, # noqa: F401 + ZeroShotClassificationParameters, # noqa: F401 + ZeroShotImageClassificationInput, # noqa: F401 + ZeroShotImageClassificationOutputElement, # noqa: F401 + ZeroShotImageClassificationParameters, # noqa: F401 + ZeroShotObjectDetectionBoundingBox, # noqa: F401 + ZeroShotObjectDetectionInput, # noqa: F401 + ZeroShotObjectDetectionOutputElement, # noqa: F401 + ZeroShotObjectDetectionParameters, # noqa: F401 + ) + from .inference_api import InferenceApi # noqa: F401 + from .keras_mixin import ( + KerasModelHubMixin, # noqa: F401 + from_pretrained_keras, # noqa: F401 + push_to_hub_keras, # noqa: F401 + save_pretrained_keras, # noqa: F401 + ) + from .repocard import ( + DatasetCard, # noqa: F401 + ModelCard, # noqa: F401 + RepoCard, # noqa: F401 + SpaceCard, # noqa: F401 + metadata_eval_result, # noqa: F401 + metadata_load, # noqa: F401 + metadata_save, # noqa: F401 + metadata_update, # noqa: F401 + ) + from .repocard_data import ( + CardData, # noqa: F401 + DatasetCardData, # noqa: F401 + EvalResult, # noqa: F401 + ModelCardData, # noqa: F401 + SpaceCardData, # noqa: F401 + ) + from .repository import Repository # noqa: F401 + from .serialization import ( + StateDictSplit, # noqa: F401 + get_tf_storage_size, # noqa: F401 + get_torch_storage_id, # noqa: F401 + get_torch_storage_size, # noqa: F401 + load_state_dict_from_file, # noqa: F401 + load_torch_model, # noqa: F401 + save_torch_model, # noqa: F401 + save_torch_state_dict, # noqa: F401 + split_state_dict_into_shards_factory, # noqa: F401 + split_tf_state_dict_into_shards, # noqa: F401 + split_torch_state_dict_into_shards, # noqa: F401 + ) + from .serialization._dduf import ( + DDUFEntry, # noqa: F401 + export_entries_as_dduf, # noqa: F401 + export_folder_as_dduf, # noqa: F401 + read_dduf_file, # noqa: F401 + ) + from .utils import ( + CachedFileInfo, # noqa: F401 + CachedRepoInfo, # noqa: F401 + CachedRevisionInfo, # noqa: F401 + CacheNotFound, # noqa: F401 + CorruptedCacheException, # noqa: F401 + DeleteCacheStrategy, # noqa: F401 + HFCacheInfo, # noqa: F401 + HfFolder, # noqa: F401 + cached_assets_path, # noqa: F401 + configure_http_backend, # noqa: F401 + dump_environment_info, # noqa: F401 + get_session, # noqa: F401 + get_token, # noqa: F401 + logging, # noqa: F401 + scan_cache_dir, # noqa: F401 + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38cb5e11be3e1d981a18459661af141606c2f5eb Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_api.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53314a8c3c1136b437158e163ca5d0cba5ad2b79 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e508ccc8859a182992edd55c2240cc1e378ac150 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_commit_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_inference_endpoints.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_inference_endpoints.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08635d59882f56409093ee2cbe7937ea34b492b6 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_inference_endpoints.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_local_folder.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_local_folder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07265feb4f9d4730c34367f469f2e07739546079 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_local_folder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_login.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_login.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46084c853e0bbaa88dcfd01268e894d33d4c8856 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_login.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_snapshot_download.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_snapshot_download.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..187fa6d6391ac5e87f1166397cf32da421379788 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_snapshot_download.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_space_api.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_space_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdddd7e412499dd3515d9f32efe6ef2121be5988 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_space_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_tensorboard_logger.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_tensorboard_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5405f03ac616fb585f6e2e293e9b029e62f5e044 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_tensorboard_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_upload_large_folder.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_upload_large_folder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9a036bb243aca9b2dd68de1a51033bcf58106b4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_upload_large_folder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_payload.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_payload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8778b0b7d4e9f76859d91f91f847b037fdeac89e Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_payload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_server.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_server.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c302aaac122f1135f4d4ac70532be82b6905ebf4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/_webhooks_server.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/community.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/community.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..faabfab27a13b79362ea5389b20e31ead1882381 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/community.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b2e577351c62a73e6fe982d8fa028bd2c04d53a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6f87712e3fde9ff6a9b8c6160a6f61d206499dc Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/errors.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/errors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2432f15458f294cd6ba6aa36f1d2a293812a4109 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/errors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/fastai_utils.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/fastai_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97a87a06a9284e37da91a9a06686b3f30f349760 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/fastai_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/file_download.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/file_download.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98136859d2ffc4b7ea41bb9fd88c81671cbe767a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/file_download.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/hf_file_system.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/hf_file_system.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3aeb17156a9ea44f44039e269f03a3e1520b5d32 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/hf_file_system.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/hub_mixin.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/hub_mixin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59dd99290df8341b47ce0cbb544ad44c10e7d59c Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/hub_mixin.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/inference_api.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/inference_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a89682c59857c990d172e90873cc20c30f315882 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/inference_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/keras_mixin.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/keras_mixin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8b4f36434c726a81bb84898b15ca7b09da4caa0 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/keras_mixin.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/lfs.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/lfs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c68aba08997076fd53597872e225dc1836811cf Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/lfs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51e2e35b98860725ac1fcb89269ab4b7737be2e4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard_data.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard_data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a5fd727cc8a43780246142485eb13a17714ee74b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repocard_data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/__pycache__/repository.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repository.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ce2179fabda783d03ed6197629c4d49804ef3af Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/__pycache__/repository.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/_commit_api.py b/lib/python3.12/site-packages/huggingface_hub/_commit_api.py new file mode 100644 index 0000000000000000000000000000000000000000..8fc82b8cb9feec212063141888a1b1e5dfacc70d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_commit_api.py @@ -0,0 +1,915 @@ +""" +Type definitions and utilities for the `create_commit` API +""" + +import base64 +import io +import math +import os +import warnings +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass, field +from itertools import groupby +from pathlib import Path, PurePosixPath +from typing import TYPE_CHECKING, Any, BinaryIO, Dict, Iterable, Iterator, List, Literal, Optional, Tuple, Union + +from tqdm.contrib.concurrent import thread_map + +from . import constants +from .errors import EntryNotFoundError, HfHubHTTPError, XetAuthorizationError, XetRefreshTokenError +from .file_download import hf_hub_url +from .lfs import UploadInfo, lfs_upload, post_lfs_batch_info +from .utils import ( + FORBIDDEN_FOLDERS, + XetTokenType, + chunk_iterable, + fetch_xet_connection_info_from_repo_info, + get_session, + hf_raise_for_status, + logging, + sha, + tqdm_stream_file, + validate_hf_hub_args, +) +from .utils import tqdm as hf_tqdm +from .utils.tqdm import _get_progress_bar_context + + +if TYPE_CHECKING: + from .hf_api import RepoFile + + +logger = logging.get_logger(__name__) + + +UploadMode = Literal["lfs", "regular"] + +# Max is 1,000 per request on the Hub for HfApi.get_paths_info +# Otherwise we get: +# HfHubHTTPError: 413 Client Error: Payload Too Large for url: https://huggingface.co/api/datasets/xxx (Request ID: xxx)\n\ntoo many parameters +# See https://github.com/huggingface/huggingface_hub/issues/1503 +FETCH_LFS_BATCH_SIZE = 500 + +UPLOAD_BATCH_MAX_NUM_FILES = 256 + + +@dataclass +class CommitOperationDelete: + """ + Data structure holding necessary info to delete a file or a folder from a repository + on the Hub. + + Args: + path_in_repo (`str`): + Relative filepath in the repo, for example: `"checkpoints/1fec34a/weights.bin"` + for a file or `"checkpoints/1fec34a/"` for a folder. + is_folder (`bool` or `Literal["auto"]`, *optional*) + Whether the Delete Operation applies to a folder or not. If "auto", the path + type (file or folder) is guessed automatically by looking if path ends with + a "/" (folder) or not (file). To explicitly set the path type, you can set + `is_folder=True` or `is_folder=False`. + """ + + path_in_repo: str + is_folder: Union[bool, Literal["auto"]] = "auto" + + def __post_init__(self): + self.path_in_repo = _validate_path_in_repo(self.path_in_repo) + + if self.is_folder == "auto": + self.is_folder = self.path_in_repo.endswith("/") + if not isinstance(self.is_folder, bool): + raise ValueError( + f"Wrong value for `is_folder`. Must be one of [`True`, `False`, `'auto'`]. Got '{self.is_folder}'." + ) + + +@dataclass +class CommitOperationCopy: + """ + Data structure holding necessary info to copy a file in a repository on the Hub. + + Limitations: + - Only LFS files can be copied. To copy a regular file, you need to download it locally and re-upload it + - Cross-repository copies are not supported. + + Note: you can combine a [`CommitOperationCopy`] and a [`CommitOperationDelete`] to rename an LFS file on the Hub. + + Args: + src_path_in_repo (`str`): + Relative filepath in the repo of the file to be copied, e.g. `"checkpoints/1fec34a/weights.bin"`. + path_in_repo (`str`): + Relative filepath in the repo where to copy the file, e.g. `"checkpoints/1fec34a/weights_copy.bin"`. + src_revision (`str`, *optional*): + The git revision of the file to be copied. Can be any valid git revision. + Default to the target commit revision. + """ + + src_path_in_repo: str + path_in_repo: str + src_revision: Optional[str] = None + # set to the OID of the file to be copied if it has already been uploaded + # useful to determine if a commit will be empty or not. + _src_oid: Optional[str] = None + # set to the OID of the file to copy to if it has already been uploaded + # useful to determine if a commit will be empty or not. + _dest_oid: Optional[str] = None + + def __post_init__(self): + self.src_path_in_repo = _validate_path_in_repo(self.src_path_in_repo) + self.path_in_repo = _validate_path_in_repo(self.path_in_repo) + + +@dataclass +class CommitOperationAdd: + """ + Data structure holding necessary info to upload a file to a repository on the Hub. + + Args: + path_in_repo (`str`): + Relative filepath in the repo, for example: `"checkpoints/1fec34a/weights.bin"` + path_or_fileobj (`str`, `Path`, `bytes`, or `BinaryIO`): + Either: + - a path to a local file (as `str` or `pathlib.Path`) to upload + - a buffer of bytes (`bytes`) holding the content of the file to upload + - a "file object" (subclass of `io.BufferedIOBase`), typically obtained + with `open(path, "rb")`. It must support `seek()` and `tell()` methods. + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `path_or_fileobj` is not one of `str`, `Path`, `bytes` or `io.BufferedIOBase`. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `path_or_fileobj` is a `str` or `Path` but not a path to an existing file. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `path_or_fileobj` is a `io.BufferedIOBase` but it doesn't support both + `seek()` and `tell()`. + """ + + path_in_repo: str + path_or_fileobj: Union[str, Path, bytes, BinaryIO] + upload_info: UploadInfo = field(init=False, repr=False) + + # Internal attributes + + # set to "lfs" or "regular" once known + _upload_mode: Optional[UploadMode] = field(init=False, repr=False, default=None) + + # set to True if .gitignore rules prevent the file from being uploaded as LFS + # (server-side check) + _should_ignore: Optional[bool] = field(init=False, repr=False, default=None) + + # set to the remote OID of the file if it has already been uploaded + # useful to determine if a commit will be empty or not + _remote_oid: Optional[str] = field(init=False, repr=False, default=None) + + # set to True once the file has been uploaded as LFS + _is_uploaded: bool = field(init=False, repr=False, default=False) + + # set to True once the file has been committed + _is_committed: bool = field(init=False, repr=False, default=False) + + def __post_init__(self) -> None: + """Validates `path_or_fileobj` and compute `upload_info`.""" + self.path_in_repo = _validate_path_in_repo(self.path_in_repo) + + # Validate `path_or_fileobj` value + if isinstance(self.path_or_fileobj, Path): + self.path_or_fileobj = str(self.path_or_fileobj) + if isinstance(self.path_or_fileobj, str): + path_or_fileobj = os.path.normpath(os.path.expanduser(self.path_or_fileobj)) + if not os.path.isfile(path_or_fileobj): + raise ValueError(f"Provided path: '{path_or_fileobj}' is not a file on the local file system") + elif not isinstance(self.path_or_fileobj, (io.BufferedIOBase, bytes)): + # ^^ Inspired from: https://stackoverflow.com/questions/44584829/how-to-determine-if-file-is-opened-in-binary-or-text-mode + raise ValueError( + "path_or_fileobj must be either an instance of str, bytes or" + " io.BufferedIOBase. If you passed a file-like object, make sure it is" + " in binary mode." + ) + if isinstance(self.path_or_fileobj, io.BufferedIOBase): + try: + self.path_or_fileobj.tell() + self.path_or_fileobj.seek(0, os.SEEK_CUR) + except (OSError, AttributeError) as exc: + raise ValueError( + "path_or_fileobj is a file-like object but does not implement seek() and tell()" + ) from exc + + # Compute "upload_info" attribute + if isinstance(self.path_or_fileobj, str): + self.upload_info = UploadInfo.from_path(self.path_or_fileobj) + elif isinstance(self.path_or_fileobj, bytes): + self.upload_info = UploadInfo.from_bytes(self.path_or_fileobj) + else: + self.upload_info = UploadInfo.from_fileobj(self.path_or_fileobj) + + @contextmanager + def as_file(self, with_tqdm: bool = False) -> Iterator[BinaryIO]: + """ + A context manager that yields a file-like object allowing to read the underlying + data behind `path_or_fileobj`. + + Args: + with_tqdm (`bool`, *optional*, defaults to `False`): + If True, iterating over the file object will display a progress bar. Only + works if the file-like object is a path to a file. Pure bytes and buffers + are not supported. + + Example: + + ```python + >>> operation = CommitOperationAdd( + ... path_in_repo="remote/dir/weights.h5", + ... path_or_fileobj="./local/weights.h5", + ... ) + CommitOperationAdd(path_in_repo='remote/dir/weights.h5', path_or_fileobj='./local/weights.h5') + + >>> with operation.as_file() as file: + ... content = file.read() + + >>> with operation.as_file(with_tqdm=True) as file: + ... while True: + ... data = file.read(1024) + ... if not data: + ... break + config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s] + + >>> with operation.as_file(with_tqdm=True) as file: + ... requests.put(..., data=file) + config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s] + ``` + """ + if isinstance(self.path_or_fileobj, str) or isinstance(self.path_or_fileobj, Path): + if with_tqdm: + with tqdm_stream_file(self.path_or_fileobj) as file: + yield file + else: + with open(self.path_or_fileobj, "rb") as file: + yield file + elif isinstance(self.path_or_fileobj, bytes): + yield io.BytesIO(self.path_or_fileobj) + elif isinstance(self.path_or_fileobj, io.BufferedIOBase): + prev_pos = self.path_or_fileobj.tell() + yield self.path_or_fileobj + self.path_or_fileobj.seek(prev_pos, io.SEEK_SET) + + def b64content(self) -> bytes: + """ + The base64-encoded content of `path_or_fileobj` + + Returns: `bytes` + """ + with self.as_file() as file: + return base64.b64encode(file.read()) + + @property + def _local_oid(self) -> Optional[str]: + """Return the OID of the local file. + + This OID is then compared to `self._remote_oid` to check if the file has changed compared to the remote one. + If the file did not change, we won't upload it again to prevent empty commits. + + For LFS files, the OID corresponds to the SHA256 of the file content (used a LFS ref). + For regular files, the OID corresponds to the SHA1 of the file content. + Note: this is slightly different to git OID computation since the oid of an LFS file is usually the git-SHA1 of the + pointer file content (not the actual file content). However, using the SHA256 is enough to detect changes + and more convenient client-side. + """ + if self._upload_mode is None: + return None + elif self._upload_mode == "lfs": + return self.upload_info.sha256.hex() + else: + # Regular file => compute sha1 + # => no need to read by chunk since the file is guaranteed to be <=5MB. + with self.as_file() as file: + return sha.git_hash(file.read()) + + +def _validate_path_in_repo(path_in_repo: str) -> str: + # Validate `path_in_repo` value to prevent a server-side issue + if path_in_repo.startswith("/"): + path_in_repo = path_in_repo[1:] + if path_in_repo == "." or path_in_repo == ".." or path_in_repo.startswith("../"): + raise ValueError(f"Invalid `path_in_repo` in CommitOperation: '{path_in_repo}'") + if path_in_repo.startswith("./"): + path_in_repo = path_in_repo[2:] + for forbidden in FORBIDDEN_FOLDERS: + if any(part == forbidden for part in path_in_repo.split("/")): + raise ValueError( + f"Invalid `path_in_repo` in CommitOperation: cannot update files under a '{forbidden}/' folder (path:" + f" '{path_in_repo}')." + ) + return path_in_repo + + +CommitOperation = Union[CommitOperationAdd, CommitOperationCopy, CommitOperationDelete] + + +def _warn_on_overwriting_operations(operations: List[CommitOperation]) -> None: + """ + Warn user when a list of operations is expected to overwrite itself in a single + commit. + + Rules: + - If a filepath is updated by multiple `CommitOperationAdd` operations, a warning + message is triggered. + - If a filepath is updated at least once by a `CommitOperationAdd` and then deleted + by a `CommitOperationDelete`, a warning is triggered. + - If a `CommitOperationDelete` deletes a filepath that is then updated by a + `CommitOperationAdd`, no warning is triggered. This is usually useless (no need to + delete before upload) but can happen if a user deletes an entire folder and then + add new files to it. + """ + nb_additions_per_path: Dict[str, int] = defaultdict(int) + for operation in operations: + path_in_repo = operation.path_in_repo + if isinstance(operation, CommitOperationAdd): + if nb_additions_per_path[path_in_repo] > 0: + warnings.warn( + "About to update multiple times the same file in the same commit:" + f" '{path_in_repo}'. This can cause undesired inconsistencies in" + " your repo." + ) + nb_additions_per_path[path_in_repo] += 1 + for parent in PurePosixPath(path_in_repo).parents: + # Also keep track of number of updated files per folder + # => warns if deleting a folder overwrite some contained files + nb_additions_per_path[str(parent)] += 1 + if isinstance(operation, CommitOperationDelete): + if nb_additions_per_path[str(PurePosixPath(path_in_repo))] > 0: + if operation.is_folder: + warnings.warn( + "About to delete a folder containing files that have just been" + f" updated within the same commit: '{path_in_repo}'. This can" + " cause undesired inconsistencies in your repo." + ) + else: + warnings.warn( + "About to delete a file that have just been updated within the" + f" same commit: '{path_in_repo}'. This can cause undesired" + " inconsistencies in your repo." + ) + + +@validate_hf_hub_args +def _upload_lfs_files( + *, + additions: List[CommitOperationAdd], + repo_type: str, + repo_id: str, + headers: Dict[str, str], + endpoint: Optional[str] = None, + num_threads: int = 5, + revision: Optional[str] = None, +): + """ + Uploads the content of `additions` to the Hub using the large file storage protocol. + + Relevant external documentation: + - LFS Batch API: https://github.com/git-lfs/git-lfs/blob/main/docs/api/batch.md + + Args: + additions (`List` of `CommitOperationAdd`): + The files to be uploaded + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + num_threads (`int`, *optional*): + The number of concurrent threads to use when uploading. Defaults to 5. + revision (`str`, *optional*): + The git revision to upload to. + + Raises: + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If an upload failed for any reason + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the server returns malformed responses + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + If the LFS batch endpoint returned an HTTP error. + """ + # Step 1: retrieve upload instructions from the LFS batch endpoint. + # Upload instructions are retrieved by chunk of 256 files to avoid reaching + # the payload limit. + batch_actions: List[Dict] = [] + for chunk in chunk_iterable(additions, chunk_size=UPLOAD_BATCH_MAX_NUM_FILES): + batch_actions_chunk, batch_errors_chunk = post_lfs_batch_info( + upload_infos=[op.upload_info for op in chunk], + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + endpoint=endpoint, + headers=headers, + token=None, # already passed in 'headers' + ) + + # If at least 1 error, we do not retrieve information for other chunks + if batch_errors_chunk: + message = "\n".join( + [ + f"Encountered error for file with OID {err.get('oid')}: `{err.get('error', {}).get('message')}" + for err in batch_errors_chunk + ] + ) + raise ValueError(f"LFS batch endpoint returned errors:\n{message}") + + batch_actions += batch_actions_chunk + oid2addop = {add_op.upload_info.sha256.hex(): add_op for add_op in additions} + + # Step 2: ignore files that have already been uploaded + filtered_actions = [] + for action in batch_actions: + if action.get("actions") is None: + logger.debug( + f"Content of file {oid2addop[action['oid']].path_in_repo} is already" + " present upstream - skipping upload." + ) + else: + filtered_actions.append(action) + + if len(filtered_actions) == 0: + logger.debug("No LFS files to upload.") + return + + # Step 3: upload files concurrently according to these instructions + def _wrapped_lfs_upload(batch_action) -> None: + try: + operation = oid2addop[batch_action["oid"]] + lfs_upload(operation=operation, lfs_batch_action=batch_action, headers=headers, endpoint=endpoint) + except Exception as exc: + raise RuntimeError(f"Error while uploading '{operation.path_in_repo}' to the Hub.") from exc + + if constants.HF_HUB_ENABLE_HF_TRANSFER: + logger.debug(f"Uploading {len(filtered_actions)} LFS files to the Hub using `hf_transfer`.") + for action in hf_tqdm(filtered_actions, name="huggingface_hub.lfs_upload"): + _wrapped_lfs_upload(action) + elif len(filtered_actions) == 1: + logger.debug("Uploading 1 LFS file to the Hub") + _wrapped_lfs_upload(filtered_actions[0]) + else: + logger.debug( + f"Uploading {len(filtered_actions)} LFS files to the Hub using up to {num_threads} threads concurrently" + ) + thread_map( + _wrapped_lfs_upload, + filtered_actions, + desc=f"Upload {len(filtered_actions)} LFS files", + max_workers=num_threads, + tqdm_class=hf_tqdm, + ) + + +@validate_hf_hub_args +def _upload_xet_files( + *, + additions: List[CommitOperationAdd], + repo_type: str, + repo_id: str, + headers: Dict[str, str], + endpoint: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, +): + """ + Uploads the content of `additions` to the Hub using the xet storage protocol. + This chunks the files and deduplicates the chunks before uploading them to xetcas storage. + + Args: + additions (`List` of `CommitOperationAdd`): + The files to be uploaded. + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + endpoint: (`str`, *optional*): + The endpoint to use for the xetcas service. Defaults to `constants.ENDPOINT`. + revision (`str`, *optional*): + The git revision to upload to. + create_pr (`bool`, *optional*): + Whether or not to create a Pull Request with that commit. + + Raises: + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If an upload failed for any reason. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the server returns malformed responses or if the user is unauthorized to upload to xet storage. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + If the LFS batch endpoint returned an HTTP error. + + **How it works:** + The file download system uses Xet storage, which is a content-addressable storage system that breaks files into chunks + for efficient storage and transfer. + + `hf_xet.upload_files` manages uploading files by: + - Taking a list of file paths to upload + - Breaking files into smaller chunks for efficient storage + - Avoiding duplicate storage by recognizing identical chunks across files + - Connecting to a storage server (CAS server) that manages these chunks + + The upload process works like this: + 1. Create a local folder at ~/.cache/huggingface/xet/chunk-cache to store file chunks for reuse. + 2. Process files in parallel (up to 8 files at once): + 2.1. Read the file content. + 2.2. Split the file content into smaller chunks based on content patterns: each chunk gets a unique ID based on what's in it. + 2.3. For each chunk: + - Check if it already exists in storage. + - Skip uploading chunks that already exist. + 2.4. Group chunks into larger blocks for efficient transfer. + 2.5. Upload these blocks to the storage server. + 2.6. Create and upload information about how the file is structured. + 3. Return reference files that contain information about the uploaded files, which can be used later to download them. + """ + if len(additions) == 0: + return + # at this point, we know that hf_xet is installed + from hf_xet import upload_bytes, upload_files + + try: + xet_connection_info = fetch_xet_connection_info_from_repo_info( + token_type=XetTokenType.WRITE, + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + headers=headers, + endpoint=endpoint, + params={"create_pr": "1"} if create_pr else None, + ) + except HfHubHTTPError as e: + if e.response.status_code == 401: + raise XetAuthorizationError( + f"You are unauthorized to upload to xet storage for {repo_type}/{repo_id}. " + f"Please check that you have configured your access token with write access to the repo." + ) from e + raise + + xet_endpoint = xet_connection_info.endpoint + access_token_info = (xet_connection_info.access_token, xet_connection_info.expiration_unix_epoch) + + def token_refresher() -> Tuple[str, int]: + new_xet_connection = fetch_xet_connection_info_from_repo_info( + token_type=XetTokenType.WRITE, + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + headers=headers, + endpoint=endpoint, + params={"create_pr": "1"} if create_pr else None, + ) + if new_xet_connection is None: + raise XetRefreshTokenError("Failed to refresh xet token") + return new_xet_connection.access_token, new_xet_connection.expiration_unix_epoch + + num_chunks = math.ceil(len(additions) / UPLOAD_BATCH_MAX_NUM_FILES) + num_chunks_num_digits = int(math.log10(num_chunks)) + 1 + for i, chunk in enumerate(chunk_iterable(additions, chunk_size=UPLOAD_BATCH_MAX_NUM_FILES)): + _chunk = [op for op in chunk] + + bytes_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, bytes)] + paths_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, (str, Path))] + expected_size = sum(op.upload_info.size for op in bytes_ops + paths_ops) + + if num_chunks > 1: + description = f"Uploading Batch [{str(i + 1).zfill(num_chunks_num_digits)}/{num_chunks}]..." + else: + description = "Uploading..." + progress_cm = _get_progress_bar_context( + desc=description, + total=expected_size, + initial=0, + unit="B", + unit_scale=True, + name="huggingface_hub.xet_put", + log_level=logger.getEffectiveLevel(), + ) + with progress_cm as progress: + + def update_progress(increment: int): + progress.update(increment) + + if len(paths_ops) > 0: + upload_files( + [str(op.path_or_fileobj) for op in paths_ops], + xet_endpoint, + access_token_info, + token_refresher, + update_progress, + repo_type, + ) + if len(bytes_ops) > 0: + upload_bytes( + [op.path_or_fileobj for op in bytes_ops], + xet_endpoint, + access_token_info, + token_refresher, + update_progress, + repo_type, + ) + return + + +def _validate_preupload_info(preupload_info: dict): + files = preupload_info.get("files") + if not isinstance(files, list): + raise ValueError("preupload_info is improperly formatted") + for file_info in files: + if not ( + isinstance(file_info, dict) + and isinstance(file_info.get("path"), str) + and isinstance(file_info.get("uploadMode"), str) + and (file_info["uploadMode"] in ("lfs", "regular")) + ): + raise ValueError("preupload_info is improperly formatted:") + return preupload_info + + +@validate_hf_hub_args +def _fetch_upload_modes( + additions: Iterable[CommitOperationAdd], + repo_type: str, + repo_id: str, + headers: Dict[str, str], + revision: str, + endpoint: Optional[str] = None, + create_pr: bool = False, + gitignore_content: Optional[str] = None, +) -> None: + """ + Requests the Hub "preupload" endpoint to determine whether each input file should be uploaded as a regular git blob, + as a git LFS blob, or as a XET file. Input `additions` are mutated in-place with the upload mode. + + Args: + additions (`Iterable` of :class:`CommitOperationAdd`): + Iterable of :class:`CommitOperationAdd` describing the files to + upload to the Hub. + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + revision (`str`): + The git revision to upload the files to. Can be any valid git revision. + gitignore_content (`str`, *optional*): + The content of the `.gitignore` file to know which files should be ignored. The order of priority + is to first check if `gitignore_content` is passed, then check if the `.gitignore` file is present + in the list of files to commit and finally default to the `.gitignore` file already hosted on the Hub + (if any). + Raises: + [`~utils.HfHubHTTPError`] + If the Hub API returned an error. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the Hub API response is improperly formatted. + """ + endpoint = endpoint if endpoint is not None else constants.ENDPOINT + + # Fetch upload mode (LFS or regular) chunk by chunk. + upload_modes: Dict[str, UploadMode] = {} + should_ignore_info: Dict[str, bool] = {} + oid_info: Dict[str, Optional[str]] = {} + + for chunk in chunk_iterable(additions, 256): + payload: Dict = { + "files": [ + { + "path": op.path_in_repo, + "sample": base64.b64encode(op.upload_info.sample).decode("ascii"), + "size": op.upload_info.size, + } + for op in chunk + ] + } + if gitignore_content is not None: + payload["gitIgnore"] = gitignore_content + + resp = get_session().post( + f"{endpoint}/api/{repo_type}s/{repo_id}/preupload/{revision}", + json=payload, + headers=headers, + params={"create_pr": "1"} if create_pr else None, + ) + hf_raise_for_status(resp) + preupload_info = _validate_preupload_info(resp.json()) + upload_modes.update(**{file["path"]: file["uploadMode"] for file in preupload_info["files"]}) + should_ignore_info.update(**{file["path"]: file["shouldIgnore"] for file in preupload_info["files"]}) + oid_info.update(**{file["path"]: file.get("oid") for file in preupload_info["files"]}) + + # Set upload mode for each addition operation + for addition in additions: + addition._upload_mode = upload_modes[addition.path_in_repo] + addition._should_ignore = should_ignore_info[addition.path_in_repo] + addition._remote_oid = oid_info[addition.path_in_repo] + + # Empty files cannot be uploaded as LFS (S3 would fail with a 501 Not Implemented) + # => empty files are uploaded as "regular" to still allow users to commit them. + for addition in additions: + if addition.upload_info.size == 0: + addition._upload_mode = "regular" + + +@validate_hf_hub_args +def _fetch_files_to_copy( + copies: Iterable[CommitOperationCopy], + repo_type: str, + repo_id: str, + headers: Dict[str, str], + revision: str, + endpoint: Optional[str] = None, +) -> Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]]: + """ + Fetch information about the files to copy. + + For LFS files, we only need their metadata (file size and sha256) while for regular files + we need to download the raw content from the Hub. + + Args: + copies (`Iterable` of :class:`CommitOperationCopy`): + Iterable of :class:`CommitOperationCopy` describing the files to + copy on the Hub. + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + revision (`str`): + The git revision to upload the files to. Can be any valid git revision. + + Returns: `Dict[Tuple[str, Optional[str]], Union[RepoFile, bytes]]]` + Key is the file path and revision of the file to copy. + Value is the raw content as bytes (for regular files) or the file information as a RepoFile (for LFS files). + + Raises: + [`~utils.HfHubHTTPError`] + If the Hub API returned an error. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the Hub API response is improperly formatted. + """ + from .hf_api import HfApi, RepoFolder + + hf_api = HfApi(endpoint=endpoint, headers=headers) + files_to_copy: Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]] = {} + # Store (path, revision) -> oid mapping + oid_info: Dict[Tuple[str, Optional[str]], Optional[str]] = {} + # 1. Fetch OIDs for destination paths in batches. + dest_paths = [op.path_in_repo for op in copies] + for offset in range(0, len(dest_paths), FETCH_LFS_BATCH_SIZE): + dest_repo_files = hf_api.get_paths_info( + repo_id=repo_id, + paths=dest_paths[offset : offset + FETCH_LFS_BATCH_SIZE], + revision=revision, + repo_type=repo_type, + ) + for file in dest_repo_files: + if not isinstance(file, RepoFolder): + oid_info[(file.path, revision)] = file.blob_id + + # 2. Group by source revision and fetch source file info in batches. + for src_revision, operations in groupby(copies, key=lambda op: op.src_revision): + operations = list(operations) # type: ignore + src_paths = [op.src_path_in_repo for op in operations] + for offset in range(0, len(src_paths), FETCH_LFS_BATCH_SIZE): + src_repo_files = hf_api.get_paths_info( + repo_id=repo_id, + paths=src_paths[offset : offset + FETCH_LFS_BATCH_SIZE], + revision=src_revision or revision, + repo_type=repo_type, + ) + + for src_repo_file in src_repo_files: + if isinstance(src_repo_file, RepoFolder): + raise NotImplementedError("Copying a folder is not implemented.") + oid_info[(src_repo_file.path, src_revision)] = src_repo_file.blob_id + # If it's an LFS file, store the RepoFile object. Otherwise, download raw bytes. + if src_repo_file.lfs: + files_to_copy[(src_repo_file.path, src_revision)] = src_repo_file + else: + # TODO: (optimization) download regular files to copy concurrently + url = hf_hub_url( + endpoint=endpoint, + repo_type=repo_type, + repo_id=repo_id, + revision=src_revision or revision, + filename=src_repo_file.path, + ) + response = get_session().get(url, headers=headers) + hf_raise_for_status(response) + files_to_copy[(src_repo_file.path, src_revision)] = response.content + # 3. Ensure all operations found a corresponding file in the Hub + # and track src/dest OIDs for each operation. + for operation in operations: + if (operation.src_path_in_repo, src_revision) not in files_to_copy: + raise EntryNotFoundError( + f"Cannot copy {operation.src_path_in_repo} at revision " + f"{src_revision or revision}: file is missing on repo." + ) + operation._src_oid = oid_info.get((operation.src_path_in_repo, operation.src_revision)) + operation._dest_oid = oid_info.get((operation.path_in_repo, revision)) + return files_to_copy + + +def _prepare_commit_payload( + operations: Iterable[CommitOperation], + files_to_copy: Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]], + commit_message: str, + commit_description: Optional[str] = None, + parent_commit: Optional[str] = None, +) -> Iterable[Dict[str, Any]]: + """ + Builds the payload to POST to the `/commit` API of the Hub. + + Payload is returned as an iterator so that it can be streamed as a ndjson in the + POST request. + + For more information, see: + - https://github.com/huggingface/huggingface_hub/issues/1085#issuecomment-1265208073 + - http://ndjson.org/ + """ + commit_description = commit_description if commit_description is not None else "" + + # 1. Send a header item with the commit metadata + header_value = {"summary": commit_message, "description": commit_description} + if parent_commit is not None: + header_value["parentCommit"] = parent_commit + yield {"key": "header", "value": header_value} + + nb_ignored_files = 0 + + # 2. Send operations, one per line + for operation in operations: + # Skip ignored files + if isinstance(operation, CommitOperationAdd) and operation._should_ignore: + logger.debug(f"Skipping file '{operation.path_in_repo}' in commit (ignored by gitignore file).") + nb_ignored_files += 1 + continue + + # 2.a. Case adding a regular file + if isinstance(operation, CommitOperationAdd) and operation._upload_mode == "regular": + yield { + "key": "file", + "value": { + "content": operation.b64content().decode(), + "path": operation.path_in_repo, + "encoding": "base64", + }, + } + # 2.b. Case adding an LFS file + elif isinstance(operation, CommitOperationAdd) and operation._upload_mode == "lfs": + yield { + "key": "lfsFile", + "value": { + "path": operation.path_in_repo, + "algo": "sha256", + "oid": operation.upload_info.sha256.hex(), + "size": operation.upload_info.size, + }, + } + # 2.c. Case deleting a file or folder + elif isinstance(operation, CommitOperationDelete): + yield { + "key": "deletedFolder" if operation.is_folder else "deletedFile", + "value": {"path": operation.path_in_repo}, + } + # 2.d. Case copying a file or folder + elif isinstance(operation, CommitOperationCopy): + file_to_copy = files_to_copy[(operation.src_path_in_repo, operation.src_revision)] + if isinstance(file_to_copy, bytes): + yield { + "key": "file", + "value": { + "content": base64.b64encode(file_to_copy).decode(), + "path": operation.path_in_repo, + "encoding": "base64", + }, + } + elif file_to_copy.lfs: + yield { + "key": "lfsFile", + "value": { + "path": operation.path_in_repo, + "algo": "sha256", + "oid": file_to_copy.lfs.sha256, + }, + } + else: + raise ValueError( + "Malformed files_to_copy (should be raw file content as bytes or RepoFile objects with LFS info." + ) + # 2.e. Never expected to happen + else: + raise ValueError( + f"Unknown operation to commit. Operation: {operation}. Upload mode:" + f" {getattr(operation, '_upload_mode', None)}" + ) + + if nb_ignored_files > 0: + logger.info(f"Skipped {nb_ignored_files} file(s) in commit (ignored by gitignore file).") diff --git a/lib/python3.12/site-packages/huggingface_hub/_commit_scheduler.py b/lib/python3.12/site-packages/huggingface_hub/_commit_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f20339e7df2d17588623dc13bb3c6be6a46b53 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_commit_scheduler.py @@ -0,0 +1,353 @@ +import atexit +import logging +import os +import time +from concurrent.futures import Future +from dataclasses import dataclass +from io import SEEK_END, SEEK_SET, BytesIO +from pathlib import Path +from threading import Lock, Thread +from typing import Dict, List, Optional, Union + +from .hf_api import DEFAULT_IGNORE_PATTERNS, CommitInfo, CommitOperationAdd, HfApi +from .utils import filter_repo_objects + + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class _FileToUpload: + """Temporary dataclass to store info about files to upload. Not meant to be used directly.""" + + local_path: Path + path_in_repo: str + size_limit: int + last_modified: float + + +class CommitScheduler: + """ + Scheduler to upload a local folder to the Hub at regular intervals (e.g. push to hub every 5 minutes). + + The recommended way to use the scheduler is to use it as a context manager. This ensures that the scheduler is + properly stopped and the last commit is triggered when the script ends. The scheduler can also be stopped manually + with the `stop` method. Checkout the [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#scheduled-uploads) + to learn more about how to use it. + + Args: + repo_id (`str`): + The id of the repo to commit to. + folder_path (`str` or `Path`): + Path to the local folder to upload regularly. + every (`int` or `float`, *optional*): + The number of minutes between each commit. Defaults to 5 minutes. + path_in_repo (`str`, *optional*): + Relative path of the directory in the repo, for example: `"checkpoints/"`. Defaults to the root folder + of the repository. + repo_type (`str`, *optional*): + The type of the repo to commit to. Defaults to `model`. + revision (`str`, *optional*): + The revision of the repo to commit to. Defaults to `main`. + private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + token (`str`, *optional*): + The token to use to commit to the repo. Defaults to the token saved on the machine. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are uploaded. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not uploaded. + squash_history (`bool`, *optional*): + Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is + useful to avoid degraded performances on the repo when it grows too large. + hf_api (`HfApi`, *optional*): + The [`HfApi`] client to use to commit to the Hub. Can be set with custom settings (user agent, token,...). + + Example: + ```py + >>> from pathlib import Path + >>> from huggingface_hub import CommitScheduler + + # Scheduler uploads every 10 minutes + >>> csv_path = Path("watched_folder/data.csv") + >>> CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path=csv_path.parent, every=10) + + >>> with csv_path.open("a") as f: + ... f.write("first line") + + # Some time later (...) + >>> with csv_path.open("a") as f: + ... f.write("second line") + ``` + + Example using a context manager: + ```py + >>> from pathlib import Path + >>> from huggingface_hub import CommitScheduler + + >>> with CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path="watched_folder", every=10) as scheduler: + ... csv_path = Path("watched_folder/data.csv") + ... with csv_path.open("a") as f: + ... f.write("first line") + ... (...) + ... with csv_path.open("a") as f: + ... f.write("second line") + + # Scheduler is now stopped and last commit have been triggered + ``` + """ + + def __init__( + self, + *, + repo_id: str, + folder_path: Union[str, Path], + every: Union[int, float] = 5, + path_in_repo: Optional[str] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + private: Optional[bool] = None, + token: Optional[str] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + squash_history: bool = False, + hf_api: Optional["HfApi"] = None, + ) -> None: + self.api = hf_api or HfApi(token=token) + + # Folder + self.folder_path = Path(folder_path).expanduser().resolve() + self.path_in_repo = path_in_repo or "" + self.allow_patterns = allow_patterns + + if ignore_patterns is None: + ignore_patterns = [] + elif isinstance(ignore_patterns, str): + ignore_patterns = [ignore_patterns] + self.ignore_patterns = ignore_patterns + DEFAULT_IGNORE_PATTERNS + + if self.folder_path.is_file(): + raise ValueError(f"'folder_path' must be a directory, not a file: '{self.folder_path}'.") + self.folder_path.mkdir(parents=True, exist_ok=True) + + # Repository + repo_url = self.api.create_repo(repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True) + self.repo_id = repo_url.repo_id + self.repo_type = repo_type + self.revision = revision + self.token = token + + # Keep track of already uploaded files + self.last_uploaded: Dict[Path, float] = {} # key is local path, value is timestamp + + # Scheduler + if not every > 0: + raise ValueError(f"'every' must be a positive integer, not '{every}'.") + self.lock = Lock() + self.every = every + self.squash_history = squash_history + + logger.info(f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes.") + self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True) + self._scheduler_thread.start() + atexit.register(self._push_to_hub) + + self.__stopped = False + + def stop(self) -> None: + """Stop the scheduler. + + A stopped scheduler cannot be restarted. Mostly for tests purposes. + """ + self.__stopped = True + + def __enter__(self) -> "CommitScheduler": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + # Upload last changes before exiting + self.trigger().result() + self.stop() + return + + def _run_scheduler(self) -> None: + """Dumb thread waiting between each scheduled push to Hub.""" + while True: + self.last_future = self.trigger() + time.sleep(self.every * 60) + if self.__stopped: + break + + def trigger(self) -> Future: + """Trigger a `push_to_hub` and return a future. + + This method is automatically called every `every` minutes. You can also call it manually to trigger a commit + immediately, without waiting for the next scheduled commit. + """ + return self.api.run_as_future(self._push_to_hub) + + def _push_to_hub(self) -> Optional[CommitInfo]: + if self.__stopped: # If stopped, already scheduled commits are ignored + return None + + logger.info("(Background) scheduled commit triggered.") + try: + value = self.push_to_hub() + if self.squash_history: + logger.info("(Background) squashing repo history.") + self.api.super_squash_history(repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision) + return value + except Exception as e: + logger.error(f"Error while pushing to Hub: {e}") # Depending on the setup, error might be silenced + raise + + def push_to_hub(self) -> Optional[CommitInfo]: + """ + Push folder to the Hub and return the commit info. + + + + This method is not meant to be called directly. It is run in the background by the scheduler, respecting a + queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency + issues. + + + + The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and + uploads only changed files. If no changes are found, the method returns without committing anything. If you want + to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful + for example to compress data together in a single file before committing. For more details and examples, check + out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads). + """ + # Check files to upload (with lock) + with self.lock: + logger.debug("Listing files to upload for scheduled commit.") + + # List files from folder (taken from `_prepare_upload_folder_additions`) + relpath_to_abspath = { + path.relative_to(self.folder_path).as_posix(): path + for path in sorted(self.folder_path.glob("**/*")) # sorted to be deterministic + if path.is_file() + } + prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else "" + + # Filter with pattern + filter out unchanged files + retrieve current file size + files_to_upload: List[_FileToUpload] = [] + for relpath in filter_repo_objects( + relpath_to_abspath.keys(), allow_patterns=self.allow_patterns, ignore_patterns=self.ignore_patterns + ): + local_path = relpath_to_abspath[relpath] + stat = local_path.stat() + if self.last_uploaded.get(local_path) is None or self.last_uploaded[local_path] != stat.st_mtime: + files_to_upload.append( + _FileToUpload( + local_path=local_path, + path_in_repo=prefix + relpath, + size_limit=stat.st_size, + last_modified=stat.st_mtime, + ) + ) + + # Return if nothing to upload + if len(files_to_upload) == 0: + logger.debug("Dropping schedule commit: no changed file to upload.") + return None + + # Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size) + logger.debug("Removing unchanged files since previous scheduled commit.") + add_operations = [ + CommitOperationAdd( + # Cap the file to its current size, even if the user append data to it while a scheduled commit is happening + path_or_fileobj=PartialFileIO(file_to_upload.local_path, size_limit=file_to_upload.size_limit), + path_in_repo=file_to_upload.path_in_repo, + ) + for file_to_upload in files_to_upload + ] + + # Upload files (append mode expected - no need for lock) + logger.debug("Uploading files for scheduled commit.") + commit_info = self.api.create_commit( + repo_id=self.repo_id, + repo_type=self.repo_type, + operations=add_operations, + commit_message="Scheduled Commit", + revision=self.revision, + ) + + # Successful commit: keep track of the latest "last_modified" for each file + for file in files_to_upload: + self.last_uploaded[file.local_path] = file.last_modified + return commit_info + + +class PartialFileIO(BytesIO): + """A file-like object that reads only the first part of a file. + + Useful to upload a file to the Hub when the user might still be appending data to it. Only the first part of the + file is uploaded (i.e. the part that was available when the filesystem was first scanned). + + In practice, only used internally by the CommitScheduler to regularly push a folder to the Hub with minimal + disturbance for the user. The object is passed to `CommitOperationAdd`. + + Only supports `read`, `tell` and `seek` methods. + + Args: + file_path (`str` or `Path`): + Path to the file to read. + size_limit (`int`): + The maximum number of bytes to read from the file. If the file is larger than this, only the first part + will be read (and uploaded). + """ + + def __init__(self, file_path: Union[str, Path], size_limit: int) -> None: + self._file_path = Path(file_path) + self._file = self._file_path.open("rb") + self._size_limit = min(size_limit, os.fstat(self._file.fileno()).st_size) + + def __del__(self) -> None: + self._file.close() + return super().__del__() + + def __repr__(self) -> str: + return f"" + + def __len__(self) -> int: + return self._size_limit + + def __getattribute__(self, name: str): + if name.startswith("_") or name in ("read", "tell", "seek"): # only 3 public methods supported + return super().__getattribute__(name) + raise NotImplementedError(f"PartialFileIO does not support '{name}'.") + + def tell(self) -> int: + """Return the current file position.""" + return self._file.tell() + + def seek(self, __offset: int, __whence: int = SEEK_SET) -> int: + """Change the stream position to the given offset. + + Behavior is the same as a regular file, except that the position is capped to the size limit. + """ + if __whence == SEEK_END: + # SEEK_END => set from the truncated end + __offset = len(self) + __offset + __whence = SEEK_SET + + pos = self._file.seek(__offset, __whence) + if pos > self._size_limit: + return self._file.seek(self._size_limit) + return pos + + def read(self, __size: Optional[int] = -1) -> bytes: + """Read at most `__size` bytes from the file. + + Behavior is the same as a regular file, except that it is capped to the size limit. + """ + current = self._file.tell() + if __size is None or __size < 0: + # Read until file limit + truncated_size = self._size_limit - current + else: + # Read until file limit or __size + truncated_size = min(__size, self._size_limit - current) + return self._file.read(truncated_size) diff --git a/lib/python3.12/site-packages/huggingface_hub/_inference_endpoints.py b/lib/python3.12/site-packages/huggingface_hub/_inference_endpoints.py new file mode 100644 index 0000000000000000000000000000000000000000..52a31361b43fb8bad273e9144abe1eb8b0e59f88 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_inference_endpoints.py @@ -0,0 +1,410 @@ +import time +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING, Dict, Optional, Union + +from huggingface_hub.errors import InferenceEndpointError, InferenceEndpointTimeoutError + +from .utils import get_session, logging, parse_datetime + + +if TYPE_CHECKING: + from .hf_api import HfApi + from .inference._client import InferenceClient + from .inference._generated._async_client import AsyncInferenceClient + +logger = logging.get_logger(__name__) + + +class InferenceEndpointStatus(str, Enum): + PENDING = "pending" + INITIALIZING = "initializing" + UPDATING = "updating" + UPDATE_FAILED = "updateFailed" + RUNNING = "running" + PAUSED = "paused" + FAILED = "failed" + SCALED_TO_ZERO = "scaledToZero" + + +class InferenceEndpointType(str, Enum): + PUBlIC = "public" + PROTECTED = "protected" + PRIVATE = "private" + + +@dataclass +class InferenceEndpoint: + """ + Contains information about a deployed Inference Endpoint. + + Args: + name (`str`): + The unique name of the Inference Endpoint. + namespace (`str`): + The namespace where the Inference Endpoint is located. + repository (`str`): + The name of the model repository deployed on this Inference Endpoint. + status ([`InferenceEndpointStatus`]): + The current status of the Inference Endpoint. + url (`str`, *optional*): + The URL of the Inference Endpoint, if available. Only a deployed Inference Endpoint will have a URL. + framework (`str`): + The machine learning framework used for the model. + revision (`str`): + The specific model revision deployed on the Inference Endpoint. + task (`str`): + The task associated with the deployed model. + created_at (`datetime.datetime`): + The timestamp when the Inference Endpoint was created. + updated_at (`datetime.datetime`): + The timestamp of the last update of the Inference Endpoint. + type ([`InferenceEndpointType`]): + The type of the Inference Endpoint (public, protected, private). + raw (`Dict`): + The raw dictionary data returned from the API. + token (`str` or `bool`, *optional*): + Authentication token for the Inference Endpoint, if set when requesting the API. Will default to the + locally saved token if not provided. Pass `token=False` if you don't want to send your token to the server. + + Example: + ```python + >>> from huggingface_hub import get_inference_endpoint + >>> endpoint = get_inference_endpoint("my-text-to-image") + >>> endpoint + InferenceEndpoint(name='my-text-to-image', ...) + + # Get status + >>> endpoint.status + 'running' + >>> endpoint.url + 'https://my-text-to-image.region.vendor.endpoints.huggingface.cloud' + + # Run inference + >>> endpoint.client.text_to_image(...) + + # Pause endpoint to save $$$ + >>> endpoint.pause() + + # ... + # Resume and wait for deployment + >>> endpoint.resume() + >>> endpoint.wait() + >>> endpoint.client.text_to_image(...) + ``` + """ + + # Field in __repr__ + name: str = field(init=False) + namespace: str + repository: str = field(init=False) + status: InferenceEndpointStatus = field(init=False) + url: Optional[str] = field(init=False) + + # Other fields + framework: str = field(repr=False, init=False) + revision: str = field(repr=False, init=False) + task: str = field(repr=False, init=False) + created_at: datetime = field(repr=False, init=False) + updated_at: datetime = field(repr=False, init=False) + type: InferenceEndpointType = field(repr=False, init=False) + + # Raw dict from the API + raw: Dict = field(repr=False) + + # Internal fields + _token: Union[str, bool, None] = field(repr=False, compare=False) + _api: "HfApi" = field(repr=False, compare=False) + + @classmethod + def from_raw( + cls, raw: Dict, namespace: str, token: Union[str, bool, None] = None, api: Optional["HfApi"] = None + ) -> "InferenceEndpoint": + """Initialize object from raw dictionary.""" + if api is None: + from .hf_api import HfApi + + api = HfApi() + if token is None: + token = api.token + + # All other fields are populated in __post_init__ + return cls(raw=raw, namespace=namespace, _token=token, _api=api) + + def __post_init__(self) -> None: + """Populate fields from raw dictionary.""" + self._populate_from_raw() + + @property + def client(self) -> "InferenceClient": + """Returns a client to make predictions on this Inference Endpoint. + + Returns: + [`InferenceClient`]: an inference client pointing to the deployed endpoint. + + Raises: + [`InferenceEndpointError`]: If the Inference Endpoint is not yet deployed. + """ + if self.url is None: + raise InferenceEndpointError( + "Cannot create a client for this Inference Endpoint as it is not yet deployed. " + "Please wait for the Inference Endpoint to be deployed using `endpoint.wait()` and try again." + ) + from .inference._client import InferenceClient + + return InferenceClient( + model=self.url, + token=self._token, # type: ignore[arg-type] # boolean token shouldn't be possible. In practice it's ok. + ) + + @property + def async_client(self) -> "AsyncInferenceClient": + """Returns a client to make predictions on this Inference Endpoint. + + Returns: + [`AsyncInferenceClient`]: an asyncio-compatible inference client pointing to the deployed endpoint. + + Raises: + [`InferenceEndpointError`]: If the Inference Endpoint is not yet deployed. + """ + if self.url is None: + raise InferenceEndpointError( + "Cannot create a client for this Inference Endpoint as it is not yet deployed. " + "Please wait for the Inference Endpoint to be deployed using `endpoint.wait()` and try again." + ) + from .inference._generated._async_client import AsyncInferenceClient + + return AsyncInferenceClient( + model=self.url, + token=self._token, # type: ignore[arg-type] # boolean token shouldn't be possible. In practice it's ok. + ) + + def wait(self, timeout: Optional[int] = None, refresh_every: int = 5) -> "InferenceEndpoint": + """Wait for the Inference Endpoint to be deployed. + + Information from the server will be fetched every 1s. If the Inference Endpoint is not deployed after `timeout` + seconds, a [`InferenceEndpointTimeoutError`] will be raised. The [`InferenceEndpoint`] will be mutated in place with the latest + data. + + Args: + timeout (`int`, *optional*): + The maximum time to wait for the Inference Endpoint to be deployed, in seconds. If `None`, will wait + indefinitely. + refresh_every (`int`, *optional*): + The time to wait between each fetch of the Inference Endpoint status, in seconds. Defaults to 5s. + + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + + Raises: + [`InferenceEndpointError`] + If the Inference Endpoint ended up in a failed state. + [`InferenceEndpointTimeoutError`] + If the Inference Endpoint is not deployed after `timeout` seconds. + """ + if timeout is not None and timeout < 0: + raise ValueError("`timeout` cannot be negative.") + if refresh_every <= 0: + raise ValueError("`refresh_every` must be positive.") + + start = time.time() + while True: + if self.status == InferenceEndpointStatus.FAILED: + raise InferenceEndpointError( + f"Inference Endpoint {self.name} failed to deploy. Please check the logs for more information." + ) + if self.status == InferenceEndpointStatus.UPDATE_FAILED: + raise InferenceEndpointError( + f"Inference Endpoint {self.name} failed to update. Please check the logs for more information." + ) + if self.status == InferenceEndpointStatus.RUNNING and self.url is not None: + # Verify the endpoint is actually reachable + response = get_session().get(self.url, headers=self._api._build_hf_headers(token=self._token)) + if response.status_code == 200: + logger.info("Inference Endpoint is ready to be used.") + return self + + if timeout is not None: + if time.time() - start > timeout: + raise InferenceEndpointTimeoutError("Timeout while waiting for Inference Endpoint to be deployed.") + logger.info(f"Inference Endpoint is not deployed yet ({self.status}). Waiting {refresh_every}s...") + time.sleep(refresh_every) + self.fetch() + + def fetch(self) -> "InferenceEndpoint": + """Fetch latest information about the Inference Endpoint. + + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + """ + obj = self._api.get_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type] + self.raw = obj.raw + self._populate_from_raw() + return self + + def update( + self, + *, + # Compute update + accelerator: Optional[str] = None, + instance_size: Optional[str] = None, + instance_type: Optional[str] = None, + min_replica: Optional[int] = None, + max_replica: Optional[int] = None, + scale_to_zero_timeout: Optional[int] = None, + # Model update + repository: Optional[str] = None, + framework: Optional[str] = None, + revision: Optional[str] = None, + task: Optional[str] = None, + custom_image: Optional[Dict] = None, + secrets: Optional[Dict[str, str]] = None, + ) -> "InferenceEndpoint": + """Update the Inference Endpoint. + + This method allows the update of either the compute configuration, the deployed model, or both. All arguments are + optional but at least one must be provided. + + This is an alias for [`HfApi.update_inference_endpoint`]. The current object is mutated in place with the + latest data from the server. + + Args: + accelerator (`str`, *optional*): + The hardware accelerator to be used for inference (e.g. `"cpu"`). + instance_size (`str`, *optional*): + The size or type of the instance to be used for hosting the model (e.g. `"x4"`). + instance_type (`str`, *optional*): + The cloud instance type where the Inference Endpoint will be deployed (e.g. `"intel-icl"`). + min_replica (`int`, *optional*): + The minimum number of replicas (instances) to keep running for the Inference Endpoint. + max_replica (`int`, *optional*): + The maximum number of replicas (instances) to scale to for the Inference Endpoint. + scale_to_zero_timeout (`int`, *optional*): + The duration in minutes before an inactive endpoint is scaled to zero. + + repository (`str`, *optional*): + The name of the model repository associated with the Inference Endpoint (e.g. `"gpt2"`). + framework (`str`, *optional*): + The machine learning framework used for the model (e.g. `"custom"`). + revision (`str`, *optional*): + The specific model revision to deploy on the Inference Endpoint (e.g. `"6c0e6080953db56375760c0471a8c5f2929baf11"`). + task (`str`, *optional*): + The task on which to deploy the model (e.g. `"text-classification"`). + custom_image (`Dict`, *optional*): + A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an + Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + """ + # Make API call + obj = self._api.update_inference_endpoint( + name=self.name, + namespace=self.namespace, + accelerator=accelerator, + instance_size=instance_size, + instance_type=instance_type, + min_replica=min_replica, + max_replica=max_replica, + scale_to_zero_timeout=scale_to_zero_timeout, + repository=repository, + framework=framework, + revision=revision, + task=task, + custom_image=custom_image, + secrets=secrets, + token=self._token, # type: ignore [arg-type] + ) + + # Mutate current object + self.raw = obj.raw + self._populate_from_raw() + return self + + def pause(self) -> "InferenceEndpoint": + """Pause the Inference Endpoint. + + A paused Inference Endpoint will not be charged. It can be resumed at any time using [`InferenceEndpoint.resume`]. + This is different than scaling the Inference Endpoint to zero with [`InferenceEndpoint.scale_to_zero`], which + would be automatically restarted when a request is made to it. + + This is an alias for [`HfApi.pause_inference_endpoint`]. The current object is mutated in place with the + latest data from the server. + + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + """ + obj = self._api.pause_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type] + self.raw = obj.raw + self._populate_from_raw() + return self + + def resume(self, running_ok: bool = True) -> "InferenceEndpoint": + """Resume the Inference Endpoint. + + This is an alias for [`HfApi.resume_inference_endpoint`]. The current object is mutated in place with the + latest data from the server. + + Args: + running_ok (`bool`, *optional*): + If `True`, the method will not raise an error if the Inference Endpoint is already running. Defaults to + `True`. + + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + """ + obj = self._api.resume_inference_endpoint( + name=self.name, namespace=self.namespace, running_ok=running_ok, token=self._token + ) # type: ignore [arg-type] + self.raw = obj.raw + self._populate_from_raw() + return self + + def scale_to_zero(self) -> "InferenceEndpoint": + """Scale Inference Endpoint to zero. + + An Inference Endpoint scaled to zero will not be charged. It will be resume on the next request to it, with a + cold start delay. This is different than pausing the Inference Endpoint with [`InferenceEndpoint.pause`], which + would require a manual resume with [`InferenceEndpoint.resume`]. + + This is an alias for [`HfApi.scale_to_zero_inference_endpoint`]. The current object is mutated in place with the + latest data from the server. + + Returns: + [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. + """ + obj = self._api.scale_to_zero_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type] + self.raw = obj.raw + self._populate_from_raw() + return self + + def delete(self) -> None: + """Delete the Inference Endpoint. + + This operation is not reversible. If you don't want to be charged for an Inference Endpoint, it is preferable + to pause it with [`InferenceEndpoint.pause`] or scale it to zero with [`InferenceEndpoint.scale_to_zero`]. + + This is an alias for [`HfApi.delete_inference_endpoint`]. + """ + self._api.delete_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type] + + def _populate_from_raw(self) -> None: + """Populate fields from raw dictionary. + + Called in __post_init__ + each time the Inference Endpoint is updated. + """ + # Repr fields + self.name = self.raw["name"] + self.repository = self.raw["model"]["repository"] + self.status = self.raw["status"]["state"] + self.url = self.raw["status"].get("url") + + # Other fields + self.framework = self.raw["model"]["framework"] + self.revision = self.raw["model"]["revision"] + self.task = self.raw["model"]["task"] + self.created_at = parse_datetime(self.raw["status"]["createdAt"]) + self.updated_at = parse_datetime(self.raw["status"]["updatedAt"]) + self.type = self.raw["type"] diff --git a/lib/python3.12/site-packages/huggingface_hub/_local_folder.py b/lib/python3.12/site-packages/huggingface_hub/_local_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..264d51c58e890bcfda1903ebb5fb22cc68f9516d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_local_folder.py @@ -0,0 +1,432 @@ +# coding=utf-8 +# Copyright 2024-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle the `../.cache/huggingface` folder in local directories. + +First discussed in https://github.com/huggingface/huggingface_hub/issues/1738 to store +download metadata when downloading files from the hub to a local directory (without +using the cache). + +./.cache/huggingface folder structure: +[4.0K] data +├── [4.0K] .cache +│ └── [4.0K] huggingface +│ └── [4.0K] download +│ ├── [ 16] file.parquet.metadata +│ ├── [ 16] file.txt.metadata +│ └── [4.0K] folder +│ └── [ 16] file.parquet.metadata +│ +├── [6.5G] file.parquet +├── [1.5K] file.txt +└── [4.0K] folder + └── [ 16] file.parquet + + +Download metadata file structure: +``` +# file.txt.metadata +11c5a3d5811f50298f278a704980280950aedb10 +a16a55fda99d2f2e7b69cce5cf93ff4ad3049930 +1712656091.123 + +# file.parquet.metadata +11c5a3d5811f50298f278a704980280950aedb10 +7c5d3f4b8b76583b422fcb9189ad6c89d5d97a094541ce8932dce3ecabde1421 +1712656091.123 +} +``` +""" + +import base64 +import hashlib +import logging +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from .utils import WeakFileLock + + +logger = logging.getLogger(__name__) + + +@dataclass +class LocalDownloadFilePaths: + """ + Paths to the files related to a download process in a local dir. + + Returned by [`get_local_download_paths`]. + + Attributes: + file_path (`Path`): + Path where the file will be saved. + lock_path (`Path`): + Path to the lock file used to ensure atomicity when reading/writing metadata. + metadata_path (`Path`): + Path to the metadata file. + """ + + file_path: Path + lock_path: Path + metadata_path: Path + + def incomplete_path(self, etag: str) -> Path: + """Return the path where a file will be temporarily downloaded before being moved to `file_path`.""" + return self.metadata_path.parent / f"{_short_hash(self.metadata_path.name)}.{etag}.incomplete" + + +@dataclass(frozen=True) +class LocalUploadFilePaths: + """ + Paths to the files related to an upload process in a local dir. + + Returned by [`get_local_upload_paths`]. + + Attributes: + path_in_repo (`str`): + Path of the file in the repo. + file_path (`Path`): + Path where the file will be saved. + lock_path (`Path`): + Path to the lock file used to ensure atomicity when reading/writing metadata. + metadata_path (`Path`): + Path to the metadata file. + """ + + path_in_repo: str + file_path: Path + lock_path: Path + metadata_path: Path + + +@dataclass +class LocalDownloadFileMetadata: + """ + Metadata about a file in the local directory related to a download process. + + Attributes: + filename (`str`): + Path of the file in the repo. + commit_hash (`str`): + Commit hash of the file in the repo. + etag (`str`): + ETag of the file in the repo. Used to check if the file has changed. + For LFS files, this is the sha256 of the file. For regular files, it corresponds to the git hash. + timestamp (`int`): + Unix timestamp of when the metadata was saved i.e. when the metadata was accurate. + """ + + filename: str + commit_hash: str + etag: str + timestamp: float + + +@dataclass +class LocalUploadFileMetadata: + """ + Metadata about a file in the local directory related to an upload process. + """ + + size: int + + # Default values correspond to "we don't know yet" + timestamp: Optional[float] = None + should_ignore: Optional[bool] = None + sha256: Optional[str] = None + upload_mode: Optional[str] = None + is_uploaded: bool = False + is_committed: bool = False + + def save(self, paths: LocalUploadFilePaths) -> None: + """Save the metadata to disk.""" + with WeakFileLock(paths.lock_path): + with paths.metadata_path.open("w") as f: + new_timestamp = time.time() + f.write(str(new_timestamp) + "\n") + + f.write(str(self.size)) # never None + f.write("\n") + + if self.should_ignore is not None: + f.write(str(int(self.should_ignore))) + f.write("\n") + + if self.sha256 is not None: + f.write(self.sha256) + f.write("\n") + + if self.upload_mode is not None: + f.write(self.upload_mode) + f.write("\n") + + f.write(str(int(self.is_uploaded)) + "\n") + f.write(str(int(self.is_committed)) + "\n") + + self.timestamp = new_timestamp + + +def get_local_download_paths(local_dir: Path, filename: str) -> LocalDownloadFilePaths: + """Compute paths to the files related to a download process. + + Folders containing the paths are all guaranteed to exist. + + Args: + local_dir (`Path`): + Path to the local directory in which files are downloaded. + filename (`str`): + Path of the file in the repo. + + Return: + [`LocalDownloadFilePaths`]: the paths to the files (file_path, lock_path, metadata_path, incomplete_path). + """ + # filename is the path in the Hub repository (separated by '/') + # make sure to have a cross platform transcription + sanitized_filename = os.path.join(*filename.split("/")) + if os.name == "nt": + if sanitized_filename.startswith("..\\") or "\\..\\" in sanitized_filename: + raise ValueError( + f"Invalid filename: cannot handle filename '{sanitized_filename}' on Windows. Please ask the repository" + " owner to rename this file." + ) + file_path = local_dir / sanitized_filename + metadata_path = _huggingface_dir(local_dir) / "download" / f"{sanitized_filename}.metadata" + lock_path = metadata_path.with_suffix(".lock") + + # Some Windows versions do not allow for paths longer than 255 characters. + # In this case, we must specify it as an extended path by using the "\\?\" prefix + if os.name == "nt": + if not str(local_dir).startswith("\\\\?\\") and len(os.path.abspath(lock_path)) > 255: + file_path = Path("\\\\?\\" + os.path.abspath(file_path)) + lock_path = Path("\\\\?\\" + os.path.abspath(lock_path)) + metadata_path = Path("\\\\?\\" + os.path.abspath(metadata_path)) + + file_path.parent.mkdir(parents=True, exist_ok=True) + metadata_path.parent.mkdir(parents=True, exist_ok=True) + return LocalDownloadFilePaths(file_path=file_path, lock_path=lock_path, metadata_path=metadata_path) + + +def get_local_upload_paths(local_dir: Path, filename: str) -> LocalUploadFilePaths: + """Compute paths to the files related to an upload process. + + Folders containing the paths are all guaranteed to exist. + + Args: + local_dir (`Path`): + Path to the local directory that is uploaded. + filename (`str`): + Path of the file in the repo. + + Return: + [`LocalUploadFilePaths`]: the paths to the files (file_path, lock_path, metadata_path). + """ + # filename is the path in the Hub repository (separated by '/') + # make sure to have a cross platform transcription + sanitized_filename = os.path.join(*filename.split("/")) + if os.name == "nt": + if sanitized_filename.startswith("..\\") or "\\..\\" in sanitized_filename: + raise ValueError( + f"Invalid filename: cannot handle filename '{sanitized_filename}' on Windows. Please ask the repository" + " owner to rename this file." + ) + file_path = local_dir / sanitized_filename + metadata_path = _huggingface_dir(local_dir) / "upload" / f"{sanitized_filename}.metadata" + lock_path = metadata_path.with_suffix(".lock") + + # Some Windows versions do not allow for paths longer than 255 characters. + # In this case, we must specify it as an extended path by using the "\\?\" prefix + if os.name == "nt": + if not str(local_dir).startswith("\\\\?\\") and len(os.path.abspath(lock_path)) > 255: + file_path = Path("\\\\?\\" + os.path.abspath(file_path)) + lock_path = Path("\\\\?\\" + os.path.abspath(lock_path)) + metadata_path = Path("\\\\?\\" + os.path.abspath(metadata_path)) + + file_path.parent.mkdir(parents=True, exist_ok=True) + metadata_path.parent.mkdir(parents=True, exist_ok=True) + return LocalUploadFilePaths( + path_in_repo=filename, file_path=file_path, lock_path=lock_path, metadata_path=metadata_path + ) + + +def read_download_metadata(local_dir: Path, filename: str) -> Optional[LocalDownloadFileMetadata]: + """Read metadata about a file in the local directory related to a download process. + + Args: + local_dir (`Path`): + Path to the local directory in which files are downloaded. + filename (`str`): + Path of the file in the repo. + + Return: + `[LocalDownloadFileMetadata]` or `None`: the metadata if it exists, `None` otherwise. + """ + paths = get_local_download_paths(local_dir, filename) + with WeakFileLock(paths.lock_path): + if paths.metadata_path.exists(): + try: + with paths.metadata_path.open() as f: + commit_hash = f.readline().strip() + etag = f.readline().strip() + timestamp = float(f.readline().strip()) + metadata = LocalDownloadFileMetadata( + filename=filename, + commit_hash=commit_hash, + etag=etag, + timestamp=timestamp, + ) + except Exception as e: + # remove the metadata file if it is corrupted / not the right format + logger.warning( + f"Invalid metadata file {paths.metadata_path}: {e}. Removing it from disk and continue." + ) + try: + paths.metadata_path.unlink() + except Exception as e: + logger.warning(f"Could not remove corrupted metadata file {paths.metadata_path}: {e}") + + try: + # check if the file exists and hasn't been modified since the metadata was saved + stat = paths.file_path.stat() + if ( + stat.st_mtime - 1 <= metadata.timestamp + ): # allow 1s difference as stat.st_mtime might not be precise + return metadata + logger.info(f"Ignored metadata for '{filename}' (outdated). Will re-compute hash.") + except FileNotFoundError: + # file does not exist => metadata is outdated + return None + return None + + +def read_upload_metadata(local_dir: Path, filename: str) -> LocalUploadFileMetadata: + """Read metadata about a file in the local directory related to an upload process. + + TODO: factorize logic with `read_download_metadata`. + + Args: + local_dir (`Path`): + Path to the local directory in which files are downloaded. + filename (`str`): + Path of the file in the repo. + + Return: + `[LocalUploadFileMetadata]` or `None`: the metadata if it exists, `None` otherwise. + """ + paths = get_local_upload_paths(local_dir, filename) + with WeakFileLock(paths.lock_path): + if paths.metadata_path.exists(): + try: + with paths.metadata_path.open() as f: + timestamp = float(f.readline().strip()) + + size = int(f.readline().strip()) # never None + + _should_ignore = f.readline().strip() + should_ignore = None if _should_ignore == "" else bool(int(_should_ignore)) + + _sha256 = f.readline().strip() + sha256 = None if _sha256 == "" else _sha256 + + _upload_mode = f.readline().strip() + upload_mode = None if _upload_mode == "" else _upload_mode + if upload_mode not in (None, "regular", "lfs"): + raise ValueError(f"Invalid upload mode in metadata {paths.path_in_repo}: {upload_mode}") + + is_uploaded = bool(int(f.readline().strip())) + is_committed = bool(int(f.readline().strip())) + + metadata = LocalUploadFileMetadata( + timestamp=timestamp, + size=size, + should_ignore=should_ignore, + sha256=sha256, + upload_mode=upload_mode, + is_uploaded=is_uploaded, + is_committed=is_committed, + ) + except Exception as e: + # remove the metadata file if it is corrupted / not the right format + logger.warning( + f"Invalid metadata file {paths.metadata_path}: {e}. Removing it from disk and continue." + ) + try: + paths.metadata_path.unlink() + except Exception as e: + logger.warning(f"Could not remove corrupted metadata file {paths.metadata_path}: {e}") + + # TODO: can we do better? + if ( + metadata.timestamp is not None + and metadata.is_uploaded # file was uploaded + and not metadata.is_committed # but not committed + and time.time() - metadata.timestamp > 20 * 3600 # and it's been more than 20 hours + ): # => we consider it as garbage-collected by S3 + metadata.is_uploaded = False + + # check if the file exists and hasn't been modified since the metadata was saved + try: + if metadata.timestamp is not None and paths.file_path.stat().st_mtime <= metadata.timestamp: + return metadata + logger.info(f"Ignored metadata for '{filename}' (outdated). Will re-compute hash.") + except FileNotFoundError: + # file does not exist => metadata is outdated + pass + + # empty metadata => we don't know anything expect its size + return LocalUploadFileMetadata(size=paths.file_path.stat().st_size) + + +def write_download_metadata(local_dir: Path, filename: str, commit_hash: str, etag: str) -> None: + """Write metadata about a file in the local directory related to a download process. + + Args: + local_dir (`Path`): + Path to the local directory in which files are downloaded. + """ + paths = get_local_download_paths(local_dir, filename) + with WeakFileLock(paths.lock_path): + with paths.metadata_path.open("w") as f: + f.write(f"{commit_hash}\n{etag}\n{time.time()}\n") + + +def _huggingface_dir(local_dir: Path) -> Path: + """Return the path to the `.cache/huggingface` directory in a local directory.""" + # Wrap in lru_cache to avoid overwriting the .gitignore file if called multiple times + path = local_dir / ".cache" / "huggingface" + path.mkdir(exist_ok=True, parents=True) + + # Create a .gitignore file in the .cache/huggingface directory if it doesn't exist + # Should be thread-safe enough like this. + gitignore = path / ".gitignore" + gitignore_lock = path / ".gitignore.lock" + if not gitignore.exists(): + try: + with WeakFileLock(gitignore_lock, timeout=0.1): + gitignore.write_text("*") + except IndexError: + pass + except OSError: # TimeoutError, FileNotFoundError, PermissionError, etc. + pass + try: + gitignore_lock.unlink() + except OSError: + pass + return path + + +def _short_hash(filename: str) -> str: + return base64.urlsafe_b64encode(hashlib.sha1(filename.encode()).digest()).decode() diff --git a/lib/python3.12/site-packages/huggingface_hub/_login.py b/lib/python3.12/site-packages/huggingface_hub/_login.py new file mode 100644 index 0000000000000000000000000000000000000000..b14702201d45bebde47f95a5bc7fc85c9e93c84b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_login.py @@ -0,0 +1,520 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains methods to log in to the Hub.""" + +import os +import subprocess +from getpass import getpass +from pathlib import Path +from typing import Optional + +from . import constants +from .commands._cli_utils import ANSI +from .utils import ( + capture_output, + get_token, + is_google_colab, + is_notebook, + list_credential_helpers, + logging, + run_subprocess, + set_git_credential, + unset_git_credential, +) +from .utils._auth import ( + _get_token_by_name, + _get_token_from_environment, + _get_token_from_file, + _get_token_from_google_colab, + _save_stored_tokens, + _save_token, + get_stored_tokens, +) +from .utils._deprecation import _deprecate_arguments, _deprecate_positional_args + + +logger = logging.get_logger(__name__) + +_HF_LOGO_ASCII = """ + _| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_| + _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _| + _|_|_|_| _| _| _| _|_| _| _|_| _| _| _| _| _| _|_| _|_|_| _|_|_|_| _| _|_|_| + _| _| _| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _| + _| _| _|_| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _| _| _| _|_|_| _|_|_|_| +""" + + +@_deprecate_arguments( + version="1.0", + deprecated_args="write_permission", + custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.", +) +@_deprecate_positional_args(version="1.0") +def login( + token: Optional[str] = None, + *, + add_to_git_credential: bool = False, + new_session: bool = True, + write_permission: bool = False, +) -> None: + """Login the machine to access the Hub. + + The `token` is persisted in cache and set as a git credential. Once done, the machine + is logged in and the access token will be available across all `huggingface_hub` + components. If `token` is not provided, it will be prompted to the user either with + a widget (in a notebook) or via the terminal. + + To log in from outside of a script, one can also use `huggingface-cli login` which is + a cli command that wraps [`login`]. + + + + [`login`] is a drop-in replacement method for [`notebook_login`] as it wraps and + extends its capabilities. + + + + + + When the token is not passed, [`login`] will automatically detect if the script runs + in a notebook or not. However, this detection might not be accurate due to the + variety of notebooks that exists nowadays. If that is the case, you can always force + the UI by using [`notebook_login`] or [`interpreter_login`]. + + + + Args: + token (`str`, *optional*): + User access token to generate from https://huggingface.co/settings/token. + add_to_git_credential (`bool`, defaults to `False`): + If `True`, token will be set as git credential. If no git credential helper + is configured, a warning will be displayed to the user. If `token` is `None`, + the value of `add_to_git_credential` is ignored and will be prompted again + to the end user. + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`): + Ignored and deprecated argument. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If an organization token is passed. Only personal account tokens are valid + to log in. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If token is invalid. + [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + If running in a notebook but `ipywidgets` is not installed. + """ + if token is not None: + if not add_to_git_credential: + logger.info( + "The token has not been saved to the git credentials helper. Pass " + "`add_to_git_credential=True` in this function directly or " + "`--add-to-git-credential` if using via `huggingface-cli` if " + "you want to set the git credential as well." + ) + _login(token, add_to_git_credential=add_to_git_credential) + elif is_notebook(): + notebook_login(new_session=new_session) + else: + interpreter_login(new_session=new_session) + + +def logout(token_name: Optional[str] = None) -> None: + """Logout the machine from the Hub. + + Token is deleted from the machine and removed from git credential. + + Args: + token_name (`str`, *optional*): + Name of the access token to logout from. If `None`, will logout from all saved access tokens. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. + """ + if get_token() is None and not get_stored_tokens(): # No active token and no saved access tokens + logger.warning("Not logged in!") + return + if not token_name: + # Delete all saved access tokens and token + for file_path in (constants.HF_TOKEN_PATH, constants.HF_STORED_TOKENS_PATH): + try: + Path(file_path).unlink() + except FileNotFoundError: + pass + logger.info("Successfully logged out from all access tokens.") + else: + _logout_from_token(token_name) + logger.info(f"Successfully logged out from access token: {token_name}.") + + unset_git_credential() + + # Check if still logged in + if _get_token_from_google_colab() is not None: + raise EnvironmentError( + "You are automatically logged in using a Google Colab secret.\n" + "To log out, you must unset the `HF_TOKEN` secret in your Colab settings." + ) + if _get_token_from_environment() is not None: + raise EnvironmentError( + "Token has been deleted from your machine but you are still logged in.\n" + "To log out, you must clear out both `HF_TOKEN` and `HUGGING_FACE_HUB_TOKEN` environment variables." + ) + + +def auth_switch(token_name: str, add_to_git_credential: bool = False) -> None: + """Switch to a different access token. + + Args: + token_name (`str`): + Name of the access token to switch to. + add_to_git_credential (`bool`, defaults to `False`): + If `True`, token will be set as git credential. If no git credential helper + is configured, a warning will be displayed to the user. If `token` is `None`, + the value of `add_to_git_credential` is ignored and will be prompted again + to the end user. + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. + """ + token = _get_token_by_name(token_name) + if not token: + raise ValueError(f"Access token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}") + # Write token to HF_TOKEN_PATH + _set_active_token(token_name, add_to_git_credential) + logger.info(f"The current active token is: {token_name}") + token_from_environment = _get_token_from_environment() + if token_from_environment is not None and token_from_environment != token: + logger.warning( + "The environment variable `HF_TOKEN` is set and will override the access token you've just switched to." + ) + + +def auth_list() -> None: + """List all stored access tokens.""" + tokens = get_stored_tokens() + + if not tokens: + logger.info("No access tokens found.") + return + # Find current token + current_token = get_token() + current_token_name = None + for token_name in tokens: + if tokens.get(token_name) == current_token: + current_token_name = token_name + # Print header + max_offset = max(len("token"), max(len(token) for token in tokens)) + 2 + print(f" {{:<{max_offset}}}| {{:<15}}".format("name", "token")) + print("-" * (max_offset + 2) + "|" + "-" * 15) + + # Print saved access tokens + for token_name in tokens: + token = tokens.get(token_name, "") + masked_token = f"{token[:3]}****{token[-4:]}" if token != "" else token + is_current = "*" if token == current_token else " " + + print(f"{is_current} {{:<{max_offset}}}| {{:<15}}".format(token_name, masked_token)) + + if _get_token_from_environment(): + logger.warning( + "\nNote: Environment variable `HF_TOKEN` is set and is the current active token independently from the stored tokens listed above." + ) + elif current_token_name is None: + logger.warning( + "\nNote: No active token is set and no environment variable `HF_TOKEN` is found. Use `huggingface-cli login` to log in." + ) + + +### +# Interpreter-based login (text) +### + + +@_deprecate_arguments( + version="1.0", + deprecated_args="write_permission", + custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.", +) +@_deprecate_positional_args(version="1.0") +def interpreter_login(*, new_session: bool = True, write_permission: bool = False) -> None: + """ + Displays a prompt to log in to the HF website and store the token. + + This is equivalent to [`login`] without passing a token when not run in a notebook. + [`interpreter_login`] is useful if you want to force the use of the terminal prompt + instead of a notebook widget. + + For more details, see [`login`]. + + Args: + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`): + Ignored and deprecated argument. + """ + if not new_session and get_token() is not None: + logger.info("User is already logged in.") + return + + from .commands.delete_cache import _ask_for_confirmation_no_tui + + print(_HF_LOGO_ASCII) + if get_token() is not None: + logger.info( + " A token is already saved on your machine. Run `huggingface-cli" + " whoami` to get more information or `huggingface-cli logout` if you want" + " to log out." + ) + logger.info(" Setting a new token will erase the existing one.") + + logger.info( + " To log in, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens ." + ) + if os.name == "nt": + logger.info("Token can be pasted using 'Right-Click'.") + token = getpass("Enter your token (input will not be visible): ") + add_to_git_credential = _ask_for_confirmation_no_tui("Add token as git credential?") + + _login(token=token, add_to_git_credential=add_to_git_credential) + + +### +# Notebook-based login (widget) +### + +NOTEBOOK_LOGIN_PASSWORD_HTML = """

Immediately click login after typing your password or +it might be stored in plain text in this notebook file.
""" + + +NOTEBOOK_LOGIN_TOKEN_HTML_START = """

Copy a token from your Hugging Face +tokens page and paste it below.
Immediately click login after copying +your token or it might be stored in plain text in this notebook file.
""" + + +NOTEBOOK_LOGIN_TOKEN_HTML_END = """ +Pro Tip: If you don't already have one, you can create a dedicated +'notebooks' token with 'write' access, that you can then easily reuse for all +notebooks. """ + + +@_deprecate_arguments( + version="1.0", + deprecated_args="write_permission", + custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.", +) +@_deprecate_positional_args(version="1.0") +def notebook_login(*, new_session: bool = True, write_permission: bool = False) -> None: + """ + Displays a widget to log in to the HF website and store the token. + + This is equivalent to [`login`] without passing a token when run in a notebook. + [`notebook_login`] is useful if you want to force the use of the notebook widget + instead of a prompt in the terminal. + + For more details, see [`login`]. + + Args: + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`): + Ignored and deprecated argument. + """ + try: + import ipywidgets.widgets as widgets # type: ignore + from IPython.display import display # type: ignore + except ImportError: + raise ImportError( + "The `notebook_login` function can only be used in a notebook (Jupyter or" + " Colab) and you need the `ipywidgets` module: `pip install ipywidgets`." + ) + if not new_session and get_token() is not None: + logger.info("User is already logged in.") + return + + box_layout = widgets.Layout(display="flex", flex_flow="column", align_items="center", width="50%") + + token_widget = widgets.Password(description="Token:") + git_checkbox_widget = widgets.Checkbox(value=True, description="Add token as git credential?") + token_finish_button = widgets.Button(description="Login") + + login_token_widget = widgets.VBox( + [ + widgets.HTML(NOTEBOOK_LOGIN_TOKEN_HTML_START), + token_widget, + git_checkbox_widget, + token_finish_button, + widgets.HTML(NOTEBOOK_LOGIN_TOKEN_HTML_END), + ], + layout=box_layout, + ) + display(login_token_widget) + + # On click events + def login_token_event(t): + """Event handler for the login button.""" + token = token_widget.value + add_to_git_credential = git_checkbox_widget.value + # Erase token and clear value to make sure it's not saved in the notebook. + token_widget.value = "" + # Hide inputs + login_token_widget.children = [widgets.Label("Connecting...")] + try: + with capture_output() as captured: + _login(token, add_to_git_credential=add_to_git_credential) + message = captured.getvalue() + except Exception as error: + message = str(error) + # Print result (success message or error) + login_token_widget.children = [widgets.Label(line) for line in message.split("\n") if line.strip()] + + token_finish_button.on_click(login_token_event) + + +### +# Login private helpers +### + + +def _login( + token: str, + add_to_git_credential: bool, +) -> None: + from .hf_api import whoami # avoid circular import + + if token.startswith("api_org"): + raise ValueError("You must use your personal account token, not an organization token.") + + token_info = whoami(token) + permission = token_info["auth"]["accessToken"]["role"] + logger.info(f"Token is valid (permission: {permission}).") + + token_name = token_info["auth"]["accessToken"]["displayName"] + # Store token locally + _save_token(token=token, token_name=token_name) + # Set active token + _set_active_token(token_name=token_name, add_to_git_credential=add_to_git_credential) + logger.info("Login successful.") + if _get_token_from_environment(): + logger.warning( + "Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured." + ) + else: + logger.info(f"The current active token is: `{token_name}`") + + +def _logout_from_token(token_name: str) -> None: + """Logout from a specific access token. + + Args: + token_name (`str`): + The name of the access token to logout from. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. + """ + stored_tokens = get_stored_tokens() + # If there is no access tokens saved or the access token name is not found, do nothing + if not stored_tokens or token_name not in stored_tokens: + return + + token = stored_tokens.pop(token_name) + _save_stored_tokens(stored_tokens) + + if token == _get_token_from_file(): + logger.warning(f"Active token '{token_name}' has been deleted.") + Path(constants.HF_TOKEN_PATH).unlink(missing_ok=True) + + +def _set_active_token( + token_name: str, + add_to_git_credential: bool, +) -> None: + """Set the active access token. + + Args: + token_name (`str`): + The name of the token to set as active. + """ + token = _get_token_by_name(token_name) + if not token: + raise ValueError(f"Token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}") + if add_to_git_credential: + if _is_git_credential_helper_configured(): + set_git_credential(token) + logger.info( + "Your token has been saved in your configured git credential helpers" + + f" ({','.join(list_credential_helpers())})." + ) + else: + logger.warning("Token has not been saved to git credential helper.") + # Write token to HF_TOKEN_PATH + path = Path(constants.HF_TOKEN_PATH) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(token) + logger.info(f"Your token has been saved to {constants.HF_TOKEN_PATH}") + + +def _is_git_credential_helper_configured() -> bool: + """Check if a git credential helper is configured. + + Warns user if not the case (except for Google Colab where "store" is set by default + by `huggingface_hub`). + """ + helpers = list_credential_helpers() + if len(helpers) > 0: + return True # Do not warn: at least 1 helper is set + + # Only in Google Colab to avoid the warning message + # See https://github.com/huggingface/huggingface_hub/issues/1043#issuecomment-1247010710 + if is_google_colab(): + _set_store_as_git_credential_helper_globally() + return True # Do not warn: "store" is used by default in Google Colab + + # Otherwise, warn user + print( + ANSI.red( + "Cannot authenticate through git-credential as no helper is defined on your" + " machine.\nYou might have to re-authenticate when pushing to the Hugging" + " Face Hub.\nRun the following command in your terminal in case you want to" + " set the 'store' credential helper as default.\n\ngit config --global" + " credential.helper store\n\nRead" + " https://git-scm.com/book/en/v2/Git-Tools-Credential-Storage for more" + " details." + ) + ) + return False + + +def _set_store_as_git_credential_helper_globally() -> None: + """Set globally the credential.helper to `store`. + + To be used only in Google Colab as we assume the user doesn't care about the git + credential config. It is the only particular case where we don't want to display the + warning message in [`notebook_login()`]. + + Related: + - https://github.com/huggingface/huggingface_hub/issues/1043 + - https://github.com/huggingface/huggingface_hub/issues/1051 + - https://git-scm.com/docs/git-credential-store + """ + try: + run_subprocess("git config --global credential.helper store") + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) diff --git a/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py b/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5ae38dab850ae4f4d616292bc21591ac849464 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py @@ -0,0 +1,308 @@ +import os +from pathlib import Path +from typing import Dict, List, Literal, Optional, Union + +import requests +from tqdm.auto import tqdm as base_tqdm +from tqdm.contrib.concurrent import thread_map + +from . import constants +from .errors import GatedRepoError, LocalEntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError +from .file_download import REGEX_COMMIT_HASH, hf_hub_download, repo_folder_name +from .hf_api import DatasetInfo, HfApi, ModelInfo, SpaceInfo +from .utils import OfflineModeIsEnabled, filter_repo_objects, logging, validate_hf_hub_args +from .utils import tqdm as hf_tqdm + + +logger = logging.get_logger(__name__) + + +@validate_hf_hub_args +def snapshot_download( + repo_id: str, + *, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + cache_dir: Union[str, Path, None] = None, + local_dir: Union[str, Path, None] = None, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Optional[Union[Dict, str]] = None, + proxies: Optional[Dict] = None, + etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT, + force_download: bool = False, + token: Optional[Union[bool, str]] = None, + local_files_only: bool = False, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + max_workers: int = 8, + tqdm_class: Optional[base_tqdm] = None, + headers: Optional[Dict[str, str]] = None, + endpoint: Optional[str] = None, + # Deprecated args + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + resume_download: Optional[bool] = None, +) -> str: + """Download repo files. + + Download a whole snapshot of a repo's files at the specified revision. This is useful when you want all files from + a repo, because you don't know which ones you will need a priori. All files are nested inside a folder in order + to keep their actual filename relative to that folder. You can also filter which files to download using + `allow_patterns` and `ignore_patterns`. + + If `local_dir` is provided, the file structure from the repo will be replicated in this location. When using this + option, the `cache_dir` will not be used and a `.cache/huggingface/` folder will be created at the root of `local_dir` + to store some metadata related to the downloaded files. While this mechanism is not as robust as the main + cache-system, it's optimized for regularly pulling the latest version of a repository. + + An alternative would be to clone the repo but this requires git and git-lfs to be installed and properly + configured. It is also not possible to filter which files to download when cloning a repository using git. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if downloading from a dataset or space, + `None` or `"model"` if downloading from a model. Default is `None`. + revision (`str`, *optional*): + An optional Git revision id which can be a branch name, a tag, or a + commit hash. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_dir (`str` or `Path`, *optional*): + If provided, the downloaded files will be placed under this directory. + library_name (`str`, *optional*): + The name of the library to which the object corresponds. + library_version (`str`, *optional*): + The version of the library. + user_agent (`str`, `dict`, *optional*): + The user-agent info in the form of a dictionary or a string. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to + `requests.request`. + etag_timeout (`float`, *optional*, defaults to `10`): + When fetching ETag, how many seconds to wait for the server to send + data before giving up which is passed to `requests.request`. + force_download (`bool`, *optional*, defaults to `False`): + Whether the file should be downloaded even if it already exists in the local cache. + token (`str`, `bool`, *optional*): + A token to be used for the download. + - If `True`, the token is read from the HuggingFace config + folder. + - If a string, it's used as the authentication token. + headers (`dict`, *optional*): + Additional headers to include in the request. Those headers take precedence over the others. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the + local cached file if it exists. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are downloaded. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not downloaded. + max_workers (`int`, *optional*): + Number of concurrent threads to download files (1 thread = 1 file download). + Defaults to 8. + tqdm_class (`tqdm`, *optional*): + If provided, overwrites the default behavior for the progress bar. Passed + argument must inherit from `tqdm.auto.tqdm` or at least mimic its behavior. + Note that the `tqdm_class` is not passed to each individual download. + Defaults to the custom HF progress bar that can be disabled by setting + `HF_HUB_DISABLE_PROGRESS_BARS` environment variable. + + Returns: + `str`: folder path of the repo snapshot. + + Raises: + [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `token=True` and the token cannot be found. + [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if + ETag cannot be determined. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid. + """ + if cache_dir is None: + cache_dir = constants.HF_HUB_CACHE + if revision is None: + revision = constants.DEFAULT_REVISION + if isinstance(cache_dir, Path): + cache_dir = str(cache_dir) + + if repo_type is None: + repo_type = "model" + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type: {repo_type}. Accepted repo types are: {str(constants.REPO_TYPES)}") + + storage_folder = os.path.join(cache_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type)) + + repo_info: Union[ModelInfo, DatasetInfo, SpaceInfo, None] = None + api_call_error: Optional[Exception] = None + if not local_files_only: + # try/except logic to handle different errors => taken from `hf_hub_download` + try: + # if we have internet connection we want to list files to download + api = HfApi( + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + endpoint=endpoint, + headers=headers, + ) + repo_info = api.repo_info(repo_id=repo_id, repo_type=repo_type, revision=revision, token=token) + except (requests.exceptions.SSLError, requests.exceptions.ProxyError): + # Actually raise for those subclasses of ConnectionError + raise + except ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + OfflineModeIsEnabled, + ) as error: + # Internet connection is down + # => will try to use local files only + api_call_error = error + pass + except RevisionNotFoundError: + # The repo was found but the revision doesn't exist on the Hub (never existed or got deleted) + raise + except requests.HTTPError as error: + # Multiple reasons for an http error: + # - Repository is private and invalid/missing token sent + # - Repository is gated and invalid/missing token sent + # - Hub is down (error 500 or 504) + # => let's switch to 'local_files_only=True' to check if the files are already cached. + # (if it's not the case, the error will be re-raised) + api_call_error = error + pass + + # At this stage, if `repo_info` is None it means either: + # - internet connection is down + # - internet connection is deactivated (local_files_only=True or HF_HUB_OFFLINE=True) + # - repo is private/gated and invalid/missing token sent + # - Hub is down + # => let's look if we can find the appropriate folder in the cache: + # - if the specified revision is a commit hash, look inside "snapshots". + # - f the specified revision is a branch or tag, look inside "refs". + # => if local_dir is not None, we will return the path to the local folder if it exists. + if repo_info is None: + # Try to get which commit hash corresponds to the specified revision + commit_hash = None + if REGEX_COMMIT_HASH.match(revision): + commit_hash = revision + else: + ref_path = os.path.join(storage_folder, "refs", revision) + if os.path.exists(ref_path): + # retrieve commit_hash from refs file + with open(ref_path) as f: + commit_hash = f.read() + + # Try to locate snapshot folder for this commit hash + if commit_hash is not None and local_dir is None: + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) + if os.path.exists(snapshot_folder): + # Snapshot folder exists => let's return it + # (but we can't check if all the files are actually there) + return snapshot_folder + + # If local_dir is not None, return it if it exists and is not empty + if local_dir is not None: + local_dir = Path(local_dir) + if local_dir.is_dir() and any(local_dir.iterdir()): + logger.warning( + f"Returning existing local_dir `{local_dir}` as remote repo cannot be accessed in `snapshot_download` ({api_call_error})." + ) + return str(local_dir.resolve()) + # If we couldn't find the appropriate folder on disk, raise an error. + if local_files_only: + raise LocalEntryNotFoundError( + "Cannot find an appropriate cached snapshot folder for the specified revision on the local disk and " + "outgoing traffic has been disabled. To enable repo look-ups and downloads online, pass " + "'local_files_only=False' as input." + ) + elif isinstance(api_call_error, OfflineModeIsEnabled): + raise LocalEntryNotFoundError( + "Cannot find an appropriate cached snapshot folder for the specified revision on the local disk and " + "outgoing traffic has been disabled. To enable repo look-ups and downloads online, set " + "'HF_HUB_OFFLINE=0' as environment variable." + ) from api_call_error + elif isinstance(api_call_error, RepositoryNotFoundError) or isinstance(api_call_error, GatedRepoError): + # Repo not found => let's raise the actual error + raise api_call_error + else: + # Otherwise: most likely a connection issue or Hub downtime => let's warn the user + raise LocalEntryNotFoundError( + "An error happened while trying to locate the files on the Hub and we cannot find the appropriate" + " snapshot folder for the specified revision on the local disk. Please check your internet connection" + " and try again." + ) from api_call_error + + # At this stage, internet connection is up and running + # => let's download the files! + assert repo_info.sha is not None, "Repo info returned from server must have a revision sha." + assert repo_info.siblings is not None, "Repo info returned from server must have a siblings list." + filtered_repo_files = list( + filter_repo_objects( + items=[f.rfilename for f in repo_info.siblings], + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + ) + ) + commit_hash = repo_info.sha + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) + # if passed revision is not identical to commit_hash + # then revision has to be a branch name or tag name. + # In that case store a ref. + if revision != commit_hash: + ref_path = os.path.join(storage_folder, "refs", revision) + try: + os.makedirs(os.path.dirname(ref_path), exist_ok=True) + with open(ref_path, "w") as f: + f.write(commit_hash) + except OSError as e: + logger.warning(f"Ignored error while writing commit hash to {ref_path}: {e}.") + + # we pass the commit_hash to hf_hub_download + # so no network call happens if we already + # have the file locally. + def _inner_hf_hub_download(repo_file: str): + return hf_hub_download( + repo_id, + filename=repo_file, + repo_type=repo_type, + revision=commit_hash, + endpoint=endpoint, + cache_dir=cache_dir, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + proxies=proxies, + etag_timeout=etag_timeout, + resume_download=resume_download, + force_download=force_download, + token=token, + headers=headers, + ) + + if constants.HF_HUB_ENABLE_HF_TRANSFER: + # when using hf_transfer we don't want extra parallelism + # from the one hf_transfer provides + for file in filtered_repo_files: + _inner_hf_hub_download(file) + else: + thread_map( + _inner_hf_hub_download, + filtered_repo_files, + desc=f"Fetching {len(filtered_repo_files)} files", + max_workers=max_workers, + # User can use its own tqdm class or the default one from `huggingface_hub.utils` + tqdm_class=tqdm_class or hf_tqdm, + ) + + if local_dir is not None: + return str(os.path.realpath(local_dir)) + return snapshot_folder diff --git a/lib/python3.12/site-packages/huggingface_hub/_space_api.py b/lib/python3.12/site-packages/huggingface_hub/_space_api.py new file mode 100644 index 0000000000000000000000000000000000000000..05fccfbc1ebdfc14840a88751914b8fc0d1a498d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_space_api.py @@ -0,0 +1,168 @@ +# coding=utf-8 +# Copyright 2019-present, the HuggingFace Inc. team. +# +# 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 dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import Dict, Optional + +from huggingface_hub.utils import parse_datetime + + +class SpaceStage(str, Enum): + """ + Enumeration of possible stage of a Space on the Hub. + + Value can be compared to a string: + ```py + assert SpaceStage.BUILDING == "BUILDING" + ``` + + Taken from https://github.com/huggingface/moon-landing/blob/main/server/repo_types/SpaceInfo.ts#L61 (private url). + """ + + # Copied from moon-landing > server > repo_types > SpaceInfo.ts (private repo) + NO_APP_FILE = "NO_APP_FILE" + CONFIG_ERROR = "CONFIG_ERROR" + BUILDING = "BUILDING" + BUILD_ERROR = "BUILD_ERROR" + RUNNING = "RUNNING" + RUNNING_BUILDING = "RUNNING_BUILDING" + RUNTIME_ERROR = "RUNTIME_ERROR" + DELETING = "DELETING" + STOPPED = "STOPPED" + PAUSED = "PAUSED" + + +class SpaceHardware(str, Enum): + """ + Enumeration of hardwares available to run your Space on the Hub. + + Value can be compared to a string: + ```py + assert SpaceHardware.CPU_BASIC == "cpu-basic" + ``` + + Taken from https://github.com/huggingface-internal/moon-landing/blob/main/server/repo_types/SpaceHardwareFlavor.ts (private url). + """ + + # CPU + CPU_BASIC = "cpu-basic" + CPU_UPGRADE = "cpu-upgrade" + CPU_XL = "cpu-xl" + + # ZeroGPU + ZERO_A10G = "zero-a10g" + + # GPU + T4_SMALL = "t4-small" + T4_MEDIUM = "t4-medium" + L4X1 = "l4x1" + L4X4 = "l4x4" + L40SX1 = "l40sx1" + L40SX4 = "l40sx4" + L40SX8 = "l40sx8" + A10G_SMALL = "a10g-small" + A10G_LARGE = "a10g-large" + A10G_LARGEX2 = "a10g-largex2" + A10G_LARGEX4 = "a10g-largex4" + A100_LARGE = "a100-large" + H100 = "h100" + H100X8 = "h100x8" + + +class SpaceStorage(str, Enum): + """ + Enumeration of persistent storage available for your Space on the Hub. + + Value can be compared to a string: + ```py + assert SpaceStorage.SMALL == "small" + ``` + + Taken from https://github.com/huggingface/moon-landing/blob/main/server/repo_types/SpaceHardwareFlavor.ts#L24 (private url). + """ + + SMALL = "small" + MEDIUM = "medium" + LARGE = "large" + + +@dataclass +class SpaceRuntime: + """ + Contains information about the current runtime of a Space. + + Args: + stage (`str`): + Current stage of the space. Example: RUNNING. + hardware (`str` or `None`): + Current hardware of the space. Example: "cpu-basic". Can be `None` if Space + is `BUILDING` for the first time. + requested_hardware (`str` or `None`): + Requested hardware. Can be different than `hardware` especially if the request + has just been made. Example: "t4-medium". Can be `None` if no hardware has + been requested yet. + sleep_time (`int` or `None`): + Number of seconds the Space will be kept alive after the last request. By default (if value is `None`), the + Space will never go to sleep if it's running on an upgraded hardware, while it will go to sleep after 48 + hours on a free 'cpu-basic' hardware. For more details, see https://huggingface.co/docs/hub/spaces-gpus#sleep-time. + raw (`dict`): + Raw response from the server. Contains more information about the Space + runtime like number of replicas, number of cpu, memory size,... + """ + + stage: SpaceStage + hardware: Optional[SpaceHardware] + requested_hardware: Optional[SpaceHardware] + sleep_time: Optional[int] + storage: Optional[SpaceStorage] + raw: Dict + + def __init__(self, data: Dict) -> None: + self.stage = data["stage"] + self.hardware = data.get("hardware", {}).get("current") + self.requested_hardware = data.get("hardware", {}).get("requested") + self.sleep_time = data.get("gcTimeout") + self.storage = data.get("storage") + self.raw = data + + +@dataclass +class SpaceVariable: + """ + Contains information about the current variables of a Space. + + Args: + key (`str`): + Variable key. Example: `"MODEL_REPO_ID"` + value (`str`): + Variable value. Example: `"the_model_repo_id"`. + description (`str` or None): + Description of the variable. Example: `"Model Repo ID of the implemented model"`. + updatedAt (`datetime` or None): + datetime of the last update of the variable (if the variable has been updated at least once). + """ + + key: str + value: str + description: Optional[str] + updated_at: Optional[datetime] + + def __init__(self, key: str, values: Dict) -> None: + self.key = key + self.value = values["value"] + self.description = values.get("description") + updated_at = values.get("updatedAt") + self.updated_at = parse_datetime(updated_at) if updated_at is not None else None diff --git a/lib/python3.12/site-packages/huggingface_hub/_tensorboard_logger.py b/lib/python3.12/site-packages/huggingface_hub/_tensorboard_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..5e910972463d3e6bc8b8796c95fde5696d999952 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_tensorboard_logger.py @@ -0,0 +1,194 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains a logger to push training logs to the Hub, using Tensorboard.""" + +from pathlib import Path +from typing import TYPE_CHECKING, List, Optional, Union + +from ._commit_scheduler import CommitScheduler +from .errors import EntryNotFoundError +from .repocard import ModelCard +from .utils import experimental + + +# Depending on user's setup, SummaryWriter can come either from 'tensorboardX' +# or from 'torch.utils.tensorboard'. Both are compatible so let's try to load +# from either of them. +try: + from tensorboardX import SummaryWriter + + is_summary_writer_available = True + +except ImportError: + try: + from torch.utils.tensorboard import SummaryWriter + + is_summary_writer_available = False + except ImportError: + # Dummy class to avoid failing at import. Will raise on instance creation. + SummaryWriter = object + is_summary_writer_available = False + +if TYPE_CHECKING: + from tensorboardX import SummaryWriter + + +class HFSummaryWriter(SummaryWriter): + """ + Wrapper around the tensorboard's `SummaryWriter` to push training logs to the Hub. + + Data is logged locally and then pushed to the Hub asynchronously. Pushing data to the Hub is done in a separate + thread to avoid blocking the training script. In particular, if the upload fails for any reason (e.g. a connection + issue), the main script will not be interrupted. Data is automatically pushed to the Hub every `commit_every` + minutes (default to every 5 minutes). + + + + `HFSummaryWriter` is experimental. Its API is subject to change in the future without prior notice. + + + + Args: + repo_id (`str`): + The id of the repo to which the logs will be pushed. + logdir (`str`, *optional*): + The directory where the logs will be written. If not specified, a local directory will be created by the + underlying `SummaryWriter` object. + commit_every (`int` or `float`, *optional*): + The frequency (in minutes) at which the logs will be pushed to the Hub. Defaults to 5 minutes. + squash_history (`bool`, *optional*): + Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is + useful to avoid degraded performances on the repo when it grows too large. + repo_type (`str`, *optional*): + The type of the repo to which the logs will be pushed. Defaults to "model". + repo_revision (`str`, *optional*): + The revision of the repo to which the logs will be pushed. Defaults to "main". + repo_private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + path_in_repo (`str`, *optional*): + The path to the folder in the repo where the logs will be pushed. Defaults to "tensorboard/". + repo_allow_patterns (`List[str]` or `str`, *optional*): + A list of patterns to include in the upload. Defaults to `"*.tfevents.*"`. Check out the + [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details. + repo_ignore_patterns (`List[str]` or `str`, *optional*): + A list of patterns to exclude in the upload. Check out the + [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details. + token (`str`, *optional*): + Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more + details + kwargs: + Additional keyword arguments passed to `SummaryWriter`. + + Examples: + ```diff + # Taken from https://pytorch.org/docs/stable/tensorboard.html + - from torch.utils.tensorboard import SummaryWriter + + from huggingface_hub import HFSummaryWriter + + import numpy as np + + - writer = SummaryWriter() + + writer = HFSummaryWriter(repo_id="username/my-trained-model") + + for n_iter in range(100): + writer.add_scalar('Loss/train', np.random.random(), n_iter) + writer.add_scalar('Loss/test', np.random.random(), n_iter) + writer.add_scalar('Accuracy/train', np.random.random(), n_iter) + writer.add_scalar('Accuracy/test', np.random.random(), n_iter) + ``` + + ```py + >>> from huggingface_hub import HFSummaryWriter + + # Logs are automatically pushed every 15 minutes (5 by default) + when exiting the context manager + >>> with HFSummaryWriter(repo_id="test_hf_logger", commit_every=15) as logger: + ... logger.add_scalar("a", 1) + ... logger.add_scalar("b", 2) + ``` + """ + + @experimental + def __new__(cls, *args, **kwargs) -> "HFSummaryWriter": + if not is_summary_writer_available: + raise ImportError( + "You must have `tensorboard` installed to use `HFSummaryWriter`. Please run `pip install --upgrade" + " tensorboardX` first." + ) + return super().__new__(cls) + + def __init__( + self, + repo_id: str, + *, + logdir: Optional[str] = None, + commit_every: Union[int, float] = 5, + squash_history: bool = False, + repo_type: Optional[str] = None, + repo_revision: Optional[str] = None, + repo_private: Optional[bool] = None, + path_in_repo: Optional[str] = "tensorboard", + repo_allow_patterns: Optional[Union[List[str], str]] = "*.tfevents.*", + repo_ignore_patterns: Optional[Union[List[str], str]] = None, + token: Optional[str] = None, + **kwargs, + ): + # Initialize SummaryWriter + super().__init__(logdir=logdir, **kwargs) + + # Check logdir has been correctly initialized and fail early otherwise. In practice, SummaryWriter takes care of it. + if not isinstance(self.logdir, str): + raise ValueError(f"`self.logdir` must be a string. Got '{self.logdir}' of type {type(self.logdir)}.") + + # Append logdir name to `path_in_repo` + if path_in_repo is None or path_in_repo == "": + path_in_repo = Path(self.logdir).name + else: + path_in_repo = path_in_repo.strip("/") + "/" + Path(self.logdir).name + + # Initialize scheduler + self.scheduler = CommitScheduler( + folder_path=self.logdir, + path_in_repo=path_in_repo, + repo_id=repo_id, + repo_type=repo_type, + revision=repo_revision, + private=repo_private, + token=token, + allow_patterns=repo_allow_patterns, + ignore_patterns=repo_ignore_patterns, + every=commit_every, + squash_history=squash_history, + ) + + # Exposing some high-level info at root level + self.repo_id = self.scheduler.repo_id + self.repo_type = self.scheduler.repo_type + self.repo_revision = self.scheduler.revision + + # Add `hf-summary-writer` tag to the model card metadata + try: + card = ModelCard.load(repo_id_or_path=self.repo_id, repo_type=self.repo_type) + except EntryNotFoundError: + card = ModelCard("") + tags = card.data.get("tags", []) + if "hf-summary-writer" not in tags: + tags.append("hf-summary-writer") + card.data["tags"] = tags + card.push_to_hub(repo_id=self.repo_id, repo_type=self.repo_type) + + def __exit__(self, exc_type, exc_val, exc_tb): + """Push to hub in a non-blocking way when exiting the logger's context manager.""" + super().__exit__(exc_type, exc_val, exc_tb) + future = self.scheduler.trigger() + future.result() diff --git a/lib/python3.12/site-packages/huggingface_hub/_upload_large_folder.py b/lib/python3.12/site-packages/huggingface_hub/_upload_large_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..e1dc132c14f18c91f4566d23991b30c5b4713a97 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_upload_large_folder.py @@ -0,0 +1,645 @@ +# coding=utf-8 +# Copyright 2024-present, the HuggingFace Inc. team. +# +# 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 enum +import logging +import os +import queue +import shutil +import sys +import threading +import time +import traceback +from datetime import datetime +from pathlib import Path +from threading import Lock +from typing import TYPE_CHECKING, List, Optional, Tuple, Union +from urllib.parse import quote + +from . import constants +from ._commit_api import CommitOperationAdd, UploadInfo, _fetch_upload_modes +from ._local_folder import LocalUploadFileMetadata, LocalUploadFilePaths, get_local_upload_paths, read_upload_metadata +from .constants import DEFAULT_REVISION, REPO_TYPES +from .utils import DEFAULT_IGNORE_PATTERNS, filter_repo_objects, tqdm +from .utils._cache_manager import _format_size +from .utils.sha import sha_fileobj + + +if TYPE_CHECKING: + from .hf_api import HfApi + +logger = logging.getLogger(__name__) + +WAITING_TIME_IF_NO_TASKS = 10 # seconds +MAX_NB_REGULAR_FILES_PER_COMMIT = 75 +MAX_NB_LFS_FILES_PER_COMMIT = 150 +COMMIT_SIZE_SCALE: List[int] = [20, 50, 75, 100, 125, 200, 250, 400, 600, 1000] + + +def upload_large_folder_internal( + api: "HfApi", + repo_id: str, + folder_path: Union[str, Path], + *, + repo_type: str, # Repo type is required! + revision: Optional[str] = None, + private: Optional[bool] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + num_workers: Optional[int] = None, + print_report: bool = True, + print_report_every: int = 60, +): + """Upload a large folder to the Hub in the most resilient way possible. + + See [`HfApi.upload_large_folder`] for the full documentation. + """ + # 1. Check args and setup + if repo_type is None: + raise ValueError( + "For large uploads, `repo_type` is explicitly required. Please set it to `model`, `dataset` or `space`." + " If you are using the CLI, pass it as `--repo-type=model`." + ) + if repo_type not in REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {REPO_TYPES}") + if revision is None: + revision = DEFAULT_REVISION + + folder_path = Path(folder_path).expanduser().resolve() + if not folder_path.is_dir(): + raise ValueError(f"Provided path: '{folder_path}' is not a directory") + + if ignore_patterns is None: + ignore_patterns = [] + elif isinstance(ignore_patterns, str): + ignore_patterns = [ignore_patterns] + ignore_patterns += DEFAULT_IGNORE_PATTERNS + + if num_workers is None: + nb_cores = os.cpu_count() or 1 + num_workers = max(nb_cores - 2, 2) # Use all but 2 cores, or at least 2 cores + + # 2. Create repo if missing + repo_url = api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private, exist_ok=True) + logger.info(f"Repo created: {repo_url}") + repo_id = repo_url.repo_id + + # 3. List files to upload + filtered_paths_list = filter_repo_objects( + (path.relative_to(folder_path).as_posix() for path in folder_path.glob("**/*") if path.is_file()), + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + ) + paths_list = [get_local_upload_paths(folder_path, relpath) for relpath in filtered_paths_list] + logger.info(f"Found {len(paths_list)} candidate files to upload") + + # Read metadata for each file + items = [ + (paths, read_upload_metadata(folder_path, paths.path_in_repo)) + for paths in tqdm(paths_list, desc="Recovering from metadata files") + ] + + # 4. Start workers + status = LargeUploadStatus(items) + threads = [ + threading.Thread( + target=_worker_job, + kwargs={ + "status": status, + "api": api, + "repo_id": repo_id, + "repo_type": repo_type, + "revision": revision, + }, + ) + for _ in range(num_workers) + ] + + for thread in threads: + thread.start() + + # 5. Print regular reports + if print_report: + print("\n\n" + status.current_report()) + last_report_ts = time.time() + while True: + time.sleep(1) + if time.time() - last_report_ts >= print_report_every: + if print_report: + _print_overwrite(status.current_report()) + last_report_ts = time.time() + if status.is_done(): + logging.info("Is done: exiting main loop") + break + + for thread in threads: + thread.join() + + logger.info(status.current_report()) + logging.info("Upload is complete!") + + +#################### +# Logic to manage workers and synchronize tasks +#################### + + +class WorkerJob(enum.Enum): + SHA256 = enum.auto() + GET_UPLOAD_MODE = enum.auto() + PREUPLOAD_LFS = enum.auto() + COMMIT = enum.auto() + WAIT = enum.auto() # if no tasks are available but we don't want to exit + + +JOB_ITEM_T = Tuple[LocalUploadFilePaths, LocalUploadFileMetadata] + + +class LargeUploadStatus: + """Contains information, queues and tasks for a large upload process.""" + + def __init__(self, items: List[JOB_ITEM_T]): + self.items = items + self.queue_sha256: "queue.Queue[JOB_ITEM_T]" = queue.Queue() + self.queue_get_upload_mode: "queue.Queue[JOB_ITEM_T]" = queue.Queue() + self.queue_preupload_lfs: "queue.Queue[JOB_ITEM_T]" = queue.Queue() + self.queue_commit: "queue.Queue[JOB_ITEM_T]" = queue.Queue() + self.lock = Lock() + + self.nb_workers_sha256: int = 0 + self.nb_workers_get_upload_mode: int = 0 + self.nb_workers_preupload_lfs: int = 0 + self.nb_workers_commit: int = 0 + self.nb_workers_waiting: int = 0 + self.last_commit_attempt: Optional[float] = None + + self._started_at = datetime.now() + self._chunk_idx: int = 1 + self._chunk_lock: Lock = Lock() + + # Setup queues + for item in self.items: + paths, metadata = item + if metadata.sha256 is None: + self.queue_sha256.put(item) + elif metadata.upload_mode is None: + self.queue_get_upload_mode.put(item) + elif metadata.upload_mode == "lfs" and not metadata.is_uploaded: + self.queue_preupload_lfs.put(item) + elif not metadata.is_committed: + self.queue_commit.put(item) + else: + logger.debug(f"Skipping file {paths.path_in_repo} (already uploaded and committed)") + + def target_chunk(self) -> int: + with self._chunk_lock: + return COMMIT_SIZE_SCALE[self._chunk_idx] + + def update_chunk(self, success: bool, nb_items: int, duration: float) -> None: + with self._chunk_lock: + if not success: + logger.warning(f"Failed to commit {nb_items} files at once. Will retry with less files in next batch.") + self._chunk_idx -= 1 + elif nb_items >= COMMIT_SIZE_SCALE[self._chunk_idx] and duration < 40: + logger.info(f"Successfully committed {nb_items} at once. Increasing the limit for next batch.") + self._chunk_idx += 1 + + self._chunk_idx = max(0, min(self._chunk_idx, len(COMMIT_SIZE_SCALE) - 1)) + + def current_report(self) -> str: + """Generate a report of the current status of the large upload.""" + nb_hashed = 0 + size_hashed = 0 + nb_preuploaded = 0 + nb_lfs = 0 + nb_lfs_unsure = 0 + size_preuploaded = 0 + nb_committed = 0 + size_committed = 0 + total_size = 0 + ignored_files = 0 + total_files = 0 + + with self.lock: + for _, metadata in self.items: + if metadata.should_ignore: + ignored_files += 1 + continue + total_size += metadata.size + total_files += 1 + if metadata.sha256 is not None: + nb_hashed += 1 + size_hashed += metadata.size + if metadata.upload_mode == "lfs": + nb_lfs += 1 + if metadata.upload_mode is None: + nb_lfs_unsure += 1 + if metadata.is_uploaded: + nb_preuploaded += 1 + size_preuploaded += metadata.size + if metadata.is_committed: + nb_committed += 1 + size_committed += metadata.size + total_size_str = _format_size(total_size) + + now = datetime.now() + now_str = now.strftime("%Y-%m-%d %H:%M:%S") + elapsed = now - self._started_at + elapsed_str = str(elapsed).split(".")[0] # remove milliseconds + + message = "\n" + "-" * 10 + message += f" {now_str} ({elapsed_str}) " + message += "-" * 10 + "\n" + + message += "Files: " + message += f"hashed {nb_hashed}/{total_files} ({_format_size(size_hashed)}/{total_size_str}) | " + message += f"pre-uploaded: {nb_preuploaded}/{nb_lfs} ({_format_size(size_preuploaded)}/{total_size_str})" + if nb_lfs_unsure > 0: + message += f" (+{nb_lfs_unsure} unsure)" + message += f" | committed: {nb_committed}/{total_files} ({_format_size(size_committed)}/{total_size_str})" + message += f" | ignored: {ignored_files}\n" + + message += "Workers: " + message += f"hashing: {self.nb_workers_sha256} | " + message += f"get upload mode: {self.nb_workers_get_upload_mode} | " + message += f"pre-uploading: {self.nb_workers_preupload_lfs} | " + message += f"committing: {self.nb_workers_commit} | " + message += f"waiting: {self.nb_workers_waiting}\n" + message += "-" * 51 + + return message + + def is_done(self) -> bool: + with self.lock: + return all(metadata.is_committed or metadata.should_ignore for _, metadata in self.items) + + +def _worker_job( + status: LargeUploadStatus, + api: "HfApi", + repo_id: str, + repo_type: str, + revision: str, +): + """ + Main process for a worker. The worker will perform tasks based on the priority list until all files are uploaded + and committed. If no tasks are available, the worker will wait for 10 seconds before checking again. + + If a task fails for any reason, the item(s) are put back in the queue for another worker to pick up. + + Read `upload_large_folder` docstring for more information on how tasks are prioritized. + """ + while True: + next_job: Optional[Tuple[WorkerJob, List[JOB_ITEM_T]]] = None + + # Determine next task + next_job = _determine_next_job(status) + if next_job is None: + return + job, items = next_job + + # Perform task + if job == WorkerJob.SHA256: + item = items[0] # single item + try: + _compute_sha256(item) + status.queue_get_upload_mode.put(item) + except KeyboardInterrupt: + raise + except Exception as e: + logger.error(f"Failed to compute sha256: {e}") + traceback.format_exc() + status.queue_sha256.put(item) + + with status.lock: + status.nb_workers_sha256 -= 1 + + elif job == WorkerJob.GET_UPLOAD_MODE: + try: + _get_upload_mode(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision) + except KeyboardInterrupt: + raise + except Exception as e: + logger.error(f"Failed to get upload mode: {e}") + traceback.format_exc() + + # Items are either: + # - dropped (if should_ignore) + # - put in LFS queue (if LFS) + # - put in commit queue (if regular) + # - or put back (if error occurred). + for item in items: + _, metadata = item + if metadata.should_ignore: + continue + if metadata.upload_mode == "lfs": + status.queue_preupload_lfs.put(item) + elif metadata.upload_mode == "regular": + status.queue_commit.put(item) + else: + status.queue_get_upload_mode.put(item) + + with status.lock: + status.nb_workers_get_upload_mode -= 1 + + elif job == WorkerJob.PREUPLOAD_LFS: + item = items[0] # single item + try: + _preupload_lfs(item, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision) + status.queue_commit.put(item) + except KeyboardInterrupt: + raise + except Exception as e: + logger.error(f"Failed to preupload LFS: {e}") + traceback.format_exc() + status.queue_preupload_lfs.put(item) + + with status.lock: + status.nb_workers_preupload_lfs -= 1 + + elif job == WorkerJob.COMMIT: + start_ts = time.time() + success = True + try: + _commit(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision) + except KeyboardInterrupt: + raise + except Exception as e: + logger.error(f"Failed to commit: {e}") + traceback.format_exc() + for item in items: + status.queue_commit.put(item) + success = False + duration = time.time() - start_ts + status.update_chunk(success, len(items), duration) + with status.lock: + status.last_commit_attempt = time.time() + status.nb_workers_commit -= 1 + + elif job == WorkerJob.WAIT: + time.sleep(WAITING_TIME_IF_NO_TASKS) + with status.lock: + status.nb_workers_waiting -= 1 + + +def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob, List[JOB_ITEM_T]]]: + with status.lock: + # 1. Commit if more than 5 minutes since last commit attempt (and at least 1 file) + if ( + status.nb_workers_commit == 0 + and status.queue_commit.qsize() > 0 + and status.last_commit_attempt is not None + and time.time() - status.last_commit_attempt > 5 * 60 + ): + status.nb_workers_commit += 1 + logger.debug("Job: commit (more than 5 minutes since last commit attempt)") + return (WorkerJob.COMMIT, _get_items_to_commit(status.queue_commit)) + + # 2. Commit if at least 100 files are ready to commit + elif status.nb_workers_commit == 0 and status.queue_commit.qsize() >= 150: + status.nb_workers_commit += 1 + logger.debug("Job: commit (>100 files ready)") + return (WorkerJob.COMMIT, _get_items_to_commit(status.queue_commit)) + + # 3. Get upload mode if at least 10 files + elif status.queue_get_upload_mode.qsize() >= 10: + status.nb_workers_get_upload_mode += 1 + logger.debug("Job: get upload mode (>10 files ready)") + return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, status.target_chunk())) + + # 4. Preupload LFS file if at least 1 file and no worker is preuploading LFS + elif status.queue_preupload_lfs.qsize() > 0 and status.nb_workers_preupload_lfs == 0: + status.nb_workers_preupload_lfs += 1 + logger.debug("Job: preupload LFS (no other worker preuploading LFS)") + return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs)) + + # 5. Compute sha256 if at least 1 file and no worker is computing sha256 + elif status.queue_sha256.qsize() > 0 and status.nb_workers_sha256 == 0: + status.nb_workers_sha256 += 1 + logger.debug("Job: sha256 (no other worker computing sha256)") + return (WorkerJob.SHA256, _get_one(status.queue_sha256)) + + # 6. Get upload mode if at least 1 file and no worker is getting upload mode + elif status.queue_get_upload_mode.qsize() > 0 and status.nb_workers_get_upload_mode == 0: + status.nb_workers_get_upload_mode += 1 + logger.debug("Job: get upload mode (no other worker getting upload mode)") + return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, status.target_chunk())) + + # 7. Preupload LFS file if at least 1 file + # Skip if hf_transfer is enabled and there is already a worker preuploading LFS + elif status.queue_preupload_lfs.qsize() > 0 and ( + status.nb_workers_preupload_lfs == 0 or not constants.HF_HUB_ENABLE_HF_TRANSFER + ): + status.nb_workers_preupload_lfs += 1 + logger.debug("Job: preupload LFS") + return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs)) + + # 8. Compute sha256 if at least 1 file + elif status.queue_sha256.qsize() > 0: + status.nb_workers_sha256 += 1 + logger.debug("Job: sha256") + return (WorkerJob.SHA256, _get_one(status.queue_sha256)) + + # 9. Get upload mode if at least 1 file + elif status.queue_get_upload_mode.qsize() > 0: + status.nb_workers_get_upload_mode += 1 + logger.debug("Job: get upload mode") + return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, status.target_chunk())) + + # 10. Commit if at least 1 file and 1 min since last commit attempt + elif ( + status.nb_workers_commit == 0 + and status.queue_commit.qsize() > 0 + and status.last_commit_attempt is not None + and time.time() - status.last_commit_attempt > 1 * 60 + ): + status.nb_workers_commit += 1 + logger.debug("Job: commit (1 min since last commit attempt)") + return (WorkerJob.COMMIT, _get_items_to_commit(status.queue_commit)) + + # 11. Commit if at least 1 file all other queues are empty and all workers are waiting + # e.g. when it's the last commit + elif ( + status.nb_workers_commit == 0 + and status.queue_commit.qsize() > 0 + and status.queue_sha256.qsize() == 0 + and status.queue_get_upload_mode.qsize() == 0 + and status.queue_preupload_lfs.qsize() == 0 + and status.nb_workers_sha256 == 0 + and status.nb_workers_get_upload_mode == 0 + and status.nb_workers_preupload_lfs == 0 + ): + status.nb_workers_commit += 1 + logger.debug("Job: commit") + return (WorkerJob.COMMIT, _get_items_to_commit(status.queue_commit)) + + # 12. If all queues are empty, exit + elif all(metadata.is_committed or metadata.should_ignore for _, metadata in status.items): + logger.info("All files have been processed! Exiting worker.") + return None + + # 13. If no task is available, wait + else: + status.nb_workers_waiting += 1 + logger.debug(f"No task available, waiting... ({WAITING_TIME_IF_NO_TASKS}s)") + return (WorkerJob.WAIT, []) + + +#################### +# Atomic jobs (sha256, get_upload_mode, preupload_lfs, commit) +#################### + + +def _compute_sha256(item: JOB_ITEM_T) -> None: + """Compute sha256 of a file and save it in metadata.""" + paths, metadata = item + if metadata.sha256 is None: + with paths.file_path.open("rb") as f: + metadata.sha256 = sha_fileobj(f).hex() + metadata.save(paths) + + +def _get_upload_mode(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None: + """Get upload mode for each file and update metadata. + + Also receive info if the file should be ignored. + """ + additions = [_build_hacky_operation(item) for item in items] + _fetch_upload_modes( + additions=additions, + repo_type=repo_type, + repo_id=repo_id, + headers=api._build_hf_headers(), + revision=quote(revision, safe=""), + ) + for item, addition in zip(items, additions): + paths, metadata = item + metadata.upload_mode = addition._upload_mode + metadata.should_ignore = addition._should_ignore + metadata.save(paths) + + +def _preupload_lfs(item: JOB_ITEM_T, api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None: + """Preupload LFS file and update metadata.""" + paths, metadata = item + addition = _build_hacky_operation(item) + api.preupload_lfs_files( + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + additions=[addition], + ) + + metadata.is_uploaded = True + metadata.save(paths) + + +def _commit(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None: + """Commit files to the repo.""" + additions = [_build_hacky_operation(item) for item in items] + api.create_commit( + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + operations=additions, + commit_message="Add files using upload-large-folder tool", + ) + for paths, metadata in items: + metadata.is_committed = True + metadata.save(paths) + + +#################### +# Hacks with CommitOperationAdd to bypass checks/sha256 calculation +#################### + + +class HackyCommitOperationAdd(CommitOperationAdd): + def __post_init__(self) -> None: + if isinstance(self.path_or_fileobj, Path): + self.path_or_fileobj = str(self.path_or_fileobj) + + +def _build_hacky_operation(item: JOB_ITEM_T) -> HackyCommitOperationAdd: + paths, metadata = item + operation = HackyCommitOperationAdd(path_in_repo=paths.path_in_repo, path_or_fileobj=paths.file_path) + with paths.file_path.open("rb") as file: + sample = file.peek(512)[:512] + if metadata.sha256 is None: + raise ValueError("sha256 must have been computed by now!") + operation.upload_info = UploadInfo(sha256=bytes.fromhex(metadata.sha256), size=metadata.size, sample=sample) + return operation + + +#################### +# Misc helpers +#################### + + +def _get_one(queue: "queue.Queue[JOB_ITEM_T]") -> List[JOB_ITEM_T]: + return [queue.get()] + + +def _get_n(queue: "queue.Queue[JOB_ITEM_T]", n: int) -> List[JOB_ITEM_T]: + return [queue.get() for _ in range(min(queue.qsize(), n))] + + +def _get_items_to_commit(queue: "queue.Queue[JOB_ITEM_T]") -> List[JOB_ITEM_T]: + """Special case for commit job: the number of items to commit depends on the type of files.""" + # Can take at most 50 regular files and/or 100 LFS files in a single commit + items: List[JOB_ITEM_T] = [] + nb_lfs, nb_regular = 0, 0 + while True: + # If empty queue => commit everything + if queue.qsize() == 0: + return items + + # If we have enough items => commit them + if nb_lfs >= MAX_NB_LFS_FILES_PER_COMMIT or nb_regular >= MAX_NB_REGULAR_FILES_PER_COMMIT: + return items + + # Else, get a new item and increase counter + item = queue.get() + items.append(item) + _, metadata = item + if metadata.upload_mode == "lfs": + nb_lfs += 1 + else: + nb_regular += 1 + + +def _print_overwrite(report: str) -> None: + """Print a report, overwriting the previous lines. + + Since tqdm in using `sys.stderr` to (re-)write progress bars, we need to use `sys.stdout` + to print the report. + + Note: works well only if no other process is writing to `sys.stdout`! + """ + report += "\n" + # Get terminal width + terminal_width = shutil.get_terminal_size().columns + + # Count number of lines that should be cleared + nb_lines = sum(len(line) // terminal_width + 1 for line in report.splitlines()) + + # Clear previous lines based on the number of lines in the report + for _ in range(nb_lines): + sys.stdout.write("\r\033[K") # Clear line + sys.stdout.write("\033[F") # Move cursor up one line + + # Print the new report, filling remaining space with whitespace + sys.stdout.write(report) + sys.stdout.write(" " * (terminal_width - len(report.splitlines()[-1]))) + sys.stdout.flush() diff --git a/lib/python3.12/site-packages/huggingface_hub/_webhooks_payload.py b/lib/python3.12/site-packages/huggingface_hub/_webhooks_payload.py new file mode 100644 index 0000000000000000000000000000000000000000..288f4b08b9428980e99ca06703442eab62fad277 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_webhooks_payload.py @@ -0,0 +1,137 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains data structures to parse the webhooks payload.""" + +from typing import List, Literal, Optional + +from .utils import is_pydantic_available + + +if is_pydantic_available(): + from pydantic import BaseModel +else: + # Define a dummy BaseModel to avoid import errors when pydantic is not installed + # Import error will be raised when trying to use the class + + class BaseModel: # type: ignore [no-redef] + def __init__(self, *args, **kwargs) -> None: + raise ImportError( + "You must have `pydantic` installed to use `WebhookPayload`. This is an optional dependency that" + " should be installed separately. Please run `pip install --upgrade pydantic` and retry." + ) + + +# This is an adaptation of the ReportV3 interface implemented in moon-landing. V0, V1 and V2 have been ignored as they +# are not in used anymore. To keep in sync when format is updated in +# https://github.com/huggingface/moon-landing/blob/main/server/lib/HFWebhooks.ts (internal link). + + +WebhookEvent_T = Literal[ + "create", + "delete", + "move", + "update", +] +RepoChangeEvent_T = Literal[ + "add", + "move", + "remove", + "update", +] +RepoType_T = Literal[ + "dataset", + "model", + "space", +] +DiscussionStatus_T = Literal[ + "closed", + "draft", + "open", + "merged", +] +SupportedWebhookVersion = Literal[3] + + +class ObjectId(BaseModel): + id: str + + +class WebhookPayloadUrl(BaseModel): + web: str + api: Optional[str] = None + + +class WebhookPayloadMovedTo(BaseModel): + name: str + owner: ObjectId + + +class WebhookPayloadWebhook(ObjectId): + version: SupportedWebhookVersion + + +class WebhookPayloadEvent(BaseModel): + action: WebhookEvent_T + scope: str + + +class WebhookPayloadDiscussionChanges(BaseModel): + base: str + mergeCommitId: Optional[str] = None + + +class WebhookPayloadComment(ObjectId): + author: ObjectId + hidden: bool + content: Optional[str] = None + url: WebhookPayloadUrl + + +class WebhookPayloadDiscussion(ObjectId): + num: int + author: ObjectId + url: WebhookPayloadUrl + title: str + isPullRequest: bool + status: DiscussionStatus_T + changes: Optional[WebhookPayloadDiscussionChanges] = None + pinned: Optional[bool] = None + + +class WebhookPayloadRepo(ObjectId): + owner: ObjectId + head_sha: Optional[str] = None + name: str + private: bool + subdomain: Optional[str] = None + tags: Optional[List[str]] = None + type: Literal["dataset", "model", "space"] + url: WebhookPayloadUrl + + +class WebhookPayloadUpdatedRef(BaseModel): + ref: str + oldSha: Optional[str] = None + newSha: Optional[str] = None + + +class WebhookPayload(BaseModel): + event: WebhookPayloadEvent + repo: WebhookPayloadRepo + discussion: Optional[WebhookPayloadDiscussion] = None + comment: Optional[WebhookPayloadComment] = None + webhook: WebhookPayloadWebhook + movedTo: Optional[WebhookPayloadMovedTo] = None + updatedRefs: Optional[List[WebhookPayloadUpdatedRef]] = None diff --git a/lib/python3.12/site-packages/huggingface_hub/_webhooks_server.py b/lib/python3.12/site-packages/huggingface_hub/_webhooks_server.py new file mode 100644 index 0000000000000000000000000000000000000000..a7bd6c86261b1cc26dfcfe3a65f5aec9851a1162 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/_webhooks_server.py @@ -0,0 +1,388 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains `WebhooksServer` and `webhook_endpoint` to create a webhook server easily.""" + +import atexit +import inspect +import os +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional + +from .utils import experimental, is_fastapi_available, is_gradio_available + + +if TYPE_CHECKING: + import gradio as gr + from fastapi import Request + +if is_fastapi_available(): + from fastapi import FastAPI, Request + from fastapi.responses import JSONResponse +else: + # Will fail at runtime if FastAPI is not available + FastAPI = Request = JSONResponse = None # type: ignore [misc, assignment] + + +_global_app: Optional["WebhooksServer"] = None +_is_local = os.environ.get("SPACE_ID") is None + + +@experimental +class WebhooksServer: + """ + The [`WebhooksServer`] class lets you create an instance of a Gradio app that can receive Huggingface webhooks. + These webhooks can be registered using the [`~WebhooksServer.add_webhook`] decorator. Webhook endpoints are added to + the app as a POST endpoint to the FastAPI router. Once all the webhooks are registered, the `launch` method has to be + called to start the app. + + It is recommended to accept [`WebhookPayload`] as the first argument of the webhook function. It is a Pydantic + model that contains all the information about the webhook event. The data will be parsed automatically for you. + + Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to setup your + WebhooksServer and deploy it on a Space. + + + + `WebhooksServer` is experimental. Its API is subject to change in the future. + + + + + + You must have `gradio` installed to use `WebhooksServer` (`pip install --upgrade gradio`). + + + + Args: + ui (`gradio.Blocks`, optional): + A Gradio UI instance to be used as the Space landing page. If `None`, a UI displaying instructions + about the configured webhooks is created. + webhook_secret (`str`, optional): + A secret key to verify incoming webhook requests. You can set this value to any secret you want as long as + you also configure it in your [webhooks settings panel](https://huggingface.co/settings/webhooks). You + can also set this value as the `WEBHOOK_SECRET` environment variable. If no secret is provided, the + webhook endpoints are opened without any security. + + Example: + + ```python + import gradio as gr + from huggingface_hub import WebhooksServer, WebhookPayload + + with gr.Blocks() as ui: + ... + + app = WebhooksServer(ui=ui, webhook_secret="my_secret_key") + + @app.add_webhook("/say_hello") + async def hello(payload: WebhookPayload): + return {"message": "hello"} + + app.launch() + ``` + """ + + def __new__(cls, *args, **kwargs) -> "WebhooksServer": + if not is_gradio_available(): + raise ImportError( + "You must have `gradio` installed to use `WebhooksServer`. Please run `pip install --upgrade gradio`" + " first." + ) + if not is_fastapi_available(): + raise ImportError( + "You must have `fastapi` installed to use `WebhooksServer`. Please run `pip install --upgrade fastapi`" + " first." + ) + return super().__new__(cls) + + def __init__( + self, + ui: Optional["gr.Blocks"] = None, + webhook_secret: Optional[str] = None, + ) -> None: + self._ui = ui + + self.webhook_secret = webhook_secret or os.getenv("WEBHOOK_SECRET") + self.registered_webhooks: Dict[str, Callable] = {} + _warn_on_empty_secret(self.webhook_secret) + + def add_webhook(self, path: Optional[str] = None) -> Callable: + """ + Decorator to add a webhook to the [`WebhooksServer`] server. + + Args: + path (`str`, optional): + The URL path to register the webhook function. If not provided, the function name will be used as the + path. In any case, all webhooks are registered under `/webhooks`. + + Raises: + ValueError: If the provided path is already registered as a webhook. + + Example: + ```python + from huggingface_hub import WebhooksServer, WebhookPayload + + app = WebhooksServer() + + @app.add_webhook + async def trigger_training(payload: WebhookPayload): + if payload.repo.type == "dataset" and payload.event.action == "update": + # Trigger a training job if a dataset is updated + ... + + app.launch() + ``` + """ + # Usage: directly as decorator. Example: `@app.add_webhook` + if callable(path): + # If path is a function, it means it was used as a decorator without arguments + return self.add_webhook()(path) + + # Usage: provide a path. Example: `@app.add_webhook(...)` + @wraps(FastAPI.post) + def _inner_post(*args, **kwargs): + func = args[0] + abs_path = f"/webhooks/{(path or func.__name__).strip('/')}" + if abs_path in self.registered_webhooks: + raise ValueError(f"Webhook {abs_path} already exists.") + self.registered_webhooks[abs_path] = func + + return _inner_post + + def launch(self, prevent_thread_lock: bool = False, **launch_kwargs: Any) -> None: + """Launch the Gradio app and register webhooks to the underlying FastAPI server. + + Input parameters are forwarded to Gradio when launching the app. + """ + ui = self._ui or self._get_default_ui() + + # Start Gradio App + # - as non-blocking so that webhooks can be added afterwards + # - as shared if launch locally (to debug webhooks) + launch_kwargs.setdefault("share", _is_local) + self.fastapi_app, _, _ = ui.launch(prevent_thread_lock=True, **launch_kwargs) + + # Register webhooks to FastAPI app + for path, func in self.registered_webhooks.items(): + # Add secret check if required + if self.webhook_secret is not None: + func = _wrap_webhook_to_check_secret(func, webhook_secret=self.webhook_secret) + + # Add route to FastAPI app + self.fastapi_app.post(path)(func) + + # Print instructions and block main thread + space_host = os.environ.get("SPACE_HOST") + url = "https://" + space_host if space_host is not None else (ui.share_url or ui.local_url) + if url is None: + raise ValueError("Cannot find the URL of the app. Please provide a valid `ui` or update `gradio` version.") + url = url.strip("/") + message = "\nWebhooks are correctly setup and ready to use:" + message += "\n" + "\n".join(f" - POST {url}{webhook}" for webhook in self.registered_webhooks) + message += "\nGo to https://huggingface.co/settings/webhooks to setup your webhooks." + print(message) + + if not prevent_thread_lock: + ui.block_thread() + + def _get_default_ui(self) -> "gr.Blocks": + """Default UI if not provided (lists webhooks and provides basic instructions).""" + import gradio as gr + + with gr.Blocks() as ui: + gr.Markdown("# This is an app to process 🤗 Webhooks") + gr.Markdown( + "Webhooks are a foundation for MLOps-related features. They allow you to listen for new changes on" + " specific repos or to all repos belonging to particular set of users/organizations (not just your" + " repos, but any repo). Check out this [guide](https://huggingface.co/docs/hub/webhooks) to get to" + " know more about webhooks on the Huggingface Hub." + ) + gr.Markdown( + f"{len(self.registered_webhooks)} webhook(s) are registered:" + + "\n\n" + + "\n ".join( + f"- [{webhook_path}]({_get_webhook_doc_url(webhook.__name__, webhook_path)})" + for webhook_path, webhook in self.registered_webhooks.items() + ) + ) + gr.Markdown( + "Go to https://huggingface.co/settings/webhooks to setup your webhooks." + + "\nYou app is running locally. Please look at the logs to check the full URL you need to set." + if _is_local + else ( + "\nThis app is running on a Space. You can find the corresponding URL in the options menu" + " (top-right) > 'Embed the Space'. The URL looks like 'https://{username}-{repo_name}.hf.space'." + ) + ) + return ui + + +@experimental +def webhook_endpoint(path: Optional[str] = None) -> Callable: + """Decorator to start a [`WebhooksServer`] and register the decorated function as a webhook endpoint. + + This is a helper to get started quickly. If you need more flexibility (custom landing page or webhook secret), + you can use [`WebhooksServer`] directly. You can register multiple webhook endpoints (to the same server) by using + this decorator multiple times. + + Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to setup your + server and deploy it on a Space. + + + + `webhook_endpoint` is experimental. Its API is subject to change in the future. + + + + + + You must have `gradio` installed to use `webhook_endpoint` (`pip install --upgrade gradio`). + + + + Args: + path (`str`, optional): + The URL path to register the webhook function. If not provided, the function name will be used as the path. + In any case, all webhooks are registered under `/webhooks`. + + Examples: + The default usage is to register a function as a webhook endpoint. The function name will be used as the path. + The server will be started automatically at exit (i.e. at the end of the script). + + ```python + from huggingface_hub import webhook_endpoint, WebhookPayload + + @webhook_endpoint + async def trigger_training(payload: WebhookPayload): + if payload.repo.type == "dataset" and payload.event.action == "update": + # Trigger a training job if a dataset is updated + ... + + # Server is automatically started at the end of the script. + ``` + + Advanced usage: register a function as a webhook endpoint and start the server manually. This is useful if you + are running it in a notebook. + + ```python + from huggingface_hub import webhook_endpoint, WebhookPayload + + @webhook_endpoint + async def trigger_training(payload: WebhookPayload): + if payload.repo.type == "dataset" and payload.event.action == "update": + # Trigger a training job if a dataset is updated + ... + + # Start the server manually + trigger_training.launch() + ``` + """ + if callable(path): + # If path is a function, it means it was used as a decorator without arguments + return webhook_endpoint()(path) + + @wraps(WebhooksServer.add_webhook) + def _inner(func: Callable) -> Callable: + app = _get_global_app() + app.add_webhook(path)(func) + if len(app.registered_webhooks) == 1: + # Register `app.launch` to run at exit (only once) + atexit.register(app.launch) + + @wraps(app.launch) + def _launch_now(): + # Run the app directly (without waiting atexit) + atexit.unregister(app.launch) + app.launch() + + func.launch = _launch_now # type: ignore + return func + + return _inner + + +def _get_global_app() -> WebhooksServer: + global _global_app + if _global_app is None: + _global_app = WebhooksServer() + return _global_app + + +def _warn_on_empty_secret(webhook_secret: Optional[str]) -> None: + if webhook_secret is None: + print("Webhook secret is not defined. This means your webhook endpoints will be open to everyone.") + print( + "To add a secret, set `WEBHOOK_SECRET` as environment variable or pass it at initialization: " + "\n\t`app = WebhooksServer(webhook_secret='my_secret', ...)`" + ) + print( + "For more details about webhook secrets, please refer to" + " https://huggingface.co/docs/hub/webhooks#webhook-secret." + ) + else: + print("Webhook secret is correctly defined.") + + +def _get_webhook_doc_url(webhook_name: str, webhook_path: str) -> str: + """Returns the anchor to a given webhook in the docs (experimental)""" + return "/docs#/default/" + webhook_name + webhook_path.replace("/", "_") + "_post" + + +def _wrap_webhook_to_check_secret(func: Callable, webhook_secret: str) -> Callable: + """Wraps a webhook function to check the webhook secret before calling the function. + + This is a hacky way to add the `request` parameter to the function signature. Since FastAPI based itself on route + parameters to inject the values to the function, we need to hack the function signature to retrieve the `Request` + object (and hence the headers). A far cleaner solution would be to use a middleware. However, since + `fastapi==0.90.1`, a middleware cannot be added once the app has started. And since the FastAPI app is started by + Gradio internals (and not by us), we cannot add a middleware. + + This method is called only when a secret has been defined by the user. If a request is sent without the + "x-webhook-secret", the function will return a 401 error (unauthorized). If the header is sent but is incorrect, + the function will return a 403 error (forbidden). + + Inspired by https://stackoverflow.com/a/33112180. + """ + initial_sig = inspect.signature(func) + + @wraps(func) + async def _protected_func(request: Request, **kwargs): + request_secret = request.headers.get("x-webhook-secret") + if request_secret is None: + return JSONResponse({"error": "x-webhook-secret header not set."}, status_code=401) + if request_secret != webhook_secret: + return JSONResponse({"error": "Invalid webhook secret."}, status_code=403) + + # Inject `request` in kwargs if required + if "request" in initial_sig.parameters: + kwargs["request"] = request + + # Handle both sync and async routes + if inspect.iscoroutinefunction(func): + return await func(**kwargs) + else: + return func(**kwargs) + + # Update signature to include request + if "request" not in initial_sig.parameters: + _protected_func.__signature__ = initial_sig.replace( # type: ignore + parameters=( + inspect.Parameter(name="request", kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Request), + ) + + tuple(initial_sig.parameters.values()) + ) + + # Return protected route + return _protected_func diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__init__.py b/lib/python3.12/site-packages/huggingface_hub/commands/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..49d088214505b9604964ab142e7f8a5b38ccd5ef --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC, abstractmethod +from argparse import _SubParsersAction + + +class BaseHuggingfaceCLICommand(ABC): + @staticmethod + @abstractmethod + def register_subcommand(parser: _SubParsersAction): + raise NotImplementedError() + + @abstractmethod + def run(self): + raise NotImplementedError() diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41741661b00b2b89517ffc8c2c7ae2a9a9a1d72f Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/_cli_utils.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/_cli_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a000126b0a118f707be4042e75616ac9331476ce Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/_cli_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/delete_cache.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/delete_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da23f1185ac71ce4574090bdaa7e405358e37063 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/delete_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/download.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/download.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d81548c46545b21190f2316bde9fe533691fe0b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/download.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/env.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..693d5b483811a9c73415de418cfd120a222e031a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/huggingface_cli.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/huggingface_cli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..248dd2743f0f01ae6722e8220c67e6da7a315949 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/huggingface_cli.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/lfs.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/lfs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..092cc87759e1476f6d9aa2e620548b6540e1c945 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/lfs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/repo_files.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/repo_files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc00f192eab1e1ed583e9aafd22f172c940212e1 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/repo_files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/scan_cache.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/scan_cache.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e37c4e19c15e0bff598c49a64d2cc0dd722ab876 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/scan_cache.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/tag.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/tag.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9796f8dbaab4ec6ebea490d3e4d4552569bad74 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/tag.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efd58384edc23da330ebb6d41a52d180d1d68710 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload_large_folder.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload_large_folder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ccac1da9b05abeb450c3632765b20ccde0431de Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/upload_large_folder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/user.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/user.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6003cb8a02ecc625fb1ba92366483b474e6a41da Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/user.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95833b28394d1264157d412668ce865d08ded566 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/commands/__pycache__/version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/_cli_utils.py b/lib/python3.12/site-packages/huggingface_hub/commands/_cli_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bd56ad6896db2a257323e022896940c0ba0d68d3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/_cli_utils.py @@ -0,0 +1,69 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains a utility for good-looking prints.""" + +import os +from typing import List, Union + + +class ANSI: + """ + Helper for en.wikipedia.org/wiki/ANSI_escape_code + """ + + _bold = "\u001b[1m" + _gray = "\u001b[90m" + _red = "\u001b[31m" + _reset = "\u001b[0m" + _yellow = "\u001b[33m" + + @classmethod + def bold(cls, s: str) -> str: + return cls._format(s, cls._bold) + + @classmethod + def gray(cls, s: str) -> str: + return cls._format(s, cls._gray) + + @classmethod + def red(cls, s: str) -> str: + return cls._format(s, cls._bold + cls._red) + + @classmethod + def yellow(cls, s: str) -> str: + return cls._format(s, cls._yellow) + + @classmethod + def _format(cls, s: str, code: str) -> str: + if os.environ.get("NO_COLOR"): + # See https://no-color.org/ + return s + return f"{code}{s}{cls._reset}" + + +def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str: + """ + Inspired by: + + - stackoverflow.com/a/8356620/593036 + - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data + """ + col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)] + row_format = ("{{:{}}} " * len(headers)).format(*col_widths) + lines = [] + lines.append(row_format.format(*headers)) + lines.append(row_format.format(*["-" * w for w in col_widths])) + for row in rows: + lines.append(row_format.format(*row)) + return "\n".join(lines) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/delete_cache.py b/lib/python3.12/site-packages/huggingface_hub/commands/delete_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..fc9eecf46977e631730c4d985bc8c4bd3c5286db --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/delete_cache.py @@ -0,0 +1,474 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to delete some revisions from the HF cache directory. + +Usage: + huggingface-cli delete-cache + huggingface-cli delete-cache --disable-tui + huggingface-cli delete-cache --dir ~/.cache/huggingface/hub + huggingface-cli delete-cache --sort=size + +NOTE: + This command is based on `InquirerPy` to build the multiselect menu in the terminal. + This dependency has to be installed with `pip install huggingface_hub[cli]`. Since + we want to avoid as much as possible cross-platform issues, I chose a library that + is built on top of `python-prompt-toolkit` which seems to be a reference in terminal + GUI (actively maintained on both Unix and Windows, 7.9k stars). + + For the moment, the TUI feature is in beta. + + See: + - https://github.com/kazhala/InquirerPy + - https://inquirerpy.readthedocs.io/en/latest/ + - https://github.com/prompt-toolkit/python-prompt-toolkit + + Other solutions could have been: + - `simple_term_menu`: would be good as well for our use case but some issues suggest + that Windows is less supported. + See: https://github.com/IngoMeyer441/simple-term-menu + - `PyInquirer`: very similar to `InquirerPy` but older and not maintained anymore. + In particular, no support of Python3.10. + See: https://github.com/CITGuru/PyInquirer + - `pick` (or `pickpack`): easy to use and flexible but built on top of Python's + standard library `curses` that is specific to Unix (not implemented on Windows). + See https://github.com/wong2/pick and https://github.com/anafvana/pickpack. + - `inquirer`: lot of traction (700 stars) but explicitly states "experimental + support of Windows". Not built on top of `python-prompt-toolkit`. + See https://github.com/magmax/python-inquirer + +TODO: add support for `huggingface-cli delete-cache aaaaaa bbbbbb cccccc (...)` ? +TODO: add "--keep-last" arg to delete revisions that are not on `main` ref +TODO: add "--filter" arg to filter repositories by name ? +TODO: add "--limit" arg to limit to X repos ? +TODO: add "-y" arg for immediate deletion ? +See discussions in https://github.com/huggingface/huggingface_hub/issues/1025. +""" + +import os +from argparse import Namespace, _SubParsersAction +from functools import wraps +from tempfile import mkstemp +from typing import Any, Callable, Iterable, List, Literal, Optional, Union + +from ..utils import CachedRepoInfo, CachedRevisionInfo, HFCacheInfo, scan_cache_dir +from . import BaseHuggingfaceCLICommand +from ._cli_utils import ANSI + + +try: + from InquirerPy import inquirer + from InquirerPy.base.control import Choice + from InquirerPy.separator import Separator + + _inquirer_py_available = True +except ImportError: + _inquirer_py_available = False + +SortingOption_T = Literal["alphabetical", "lastUpdated", "lastUsed", "size"] + + +def require_inquirer_py(fn: Callable) -> Callable: + """Decorator to flag methods that require `InquirerPy`.""" + + # TODO: refactor this + imports in a unified pattern across codebase + @wraps(fn) + def _inner(*args, **kwargs): + if not _inquirer_py_available: + raise ImportError( + "The `delete-cache` command requires extra dependencies to work with" + " the TUI.\nPlease run `pip install huggingface_hub[cli]` to install" + " them.\nOtherwise, disable TUI using the `--disable-tui` flag." + ) + + return fn(*args, **kwargs) + + return _inner + + +# Possibility for the user to cancel deletion +_CANCEL_DELETION_STR = "CANCEL_DELETION" + + +class DeleteCacheCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + delete_cache_parser = parser.add_parser("delete-cache", help="Delete revisions from the cache directory.") + + delete_cache_parser.add_argument( + "--dir", + type=str, + default=None, + help="cache directory (optional). Default to the default HuggingFace cache.", + ) + + delete_cache_parser.add_argument( + "--disable-tui", + action="store_true", + help=( + "Disable Terminal User Interface (TUI) mode. Useful if your" + " platform/terminal doesn't support the multiselect menu." + ), + ) + + delete_cache_parser.add_argument( + "--sort", + nargs="?", + choices=["alphabetical", "lastUpdated", "lastUsed", "size"], + help=( + "Sort repositories by the specified criteria. Options: " + "'alphabetical' (A-Z), " + "'lastUpdated' (newest first), " + "'lastUsed' (most recent first), " + "'size' (largest first)." + ), + ) + + delete_cache_parser.set_defaults(func=DeleteCacheCommand) + + def __init__(self, args: Namespace) -> None: + self.cache_dir: Optional[str] = args.dir + self.disable_tui: bool = args.disable_tui + self.sort_by: Optional[SortingOption_T] = args.sort + + def run(self): + """Run `delete-cache` command with or without TUI.""" + # Scan cache directory + hf_cache_info = scan_cache_dir(self.cache_dir) + + # Manual review from the user + if self.disable_tui: + selected_hashes = _manual_review_no_tui(hf_cache_info, preselected=[], sort_by=self.sort_by) + else: + selected_hashes = _manual_review_tui(hf_cache_info, preselected=[], sort_by=self.sort_by) + + # If deletion is not cancelled + if len(selected_hashes) > 0 and _CANCEL_DELETION_STR not in selected_hashes: + confirm_message = _get_expectations_str(hf_cache_info, selected_hashes) + " Confirm deletion ?" + + # Confirm deletion + if self.disable_tui: + confirmed = _ask_for_confirmation_no_tui(confirm_message) + else: + confirmed = _ask_for_confirmation_tui(confirm_message) + + # Deletion is confirmed + if confirmed: + strategy = hf_cache_info.delete_revisions(*selected_hashes) + print("Start deletion.") + strategy.execute() + print( + f"Done. Deleted {len(strategy.repos)} repo(s) and" + f" {len(strategy.snapshots)} revision(s) for a total of" + f" {strategy.expected_freed_size_str}." + ) + return + + # Deletion is cancelled + print("Deletion is cancelled. Do nothing.") + + +def _get_repo_sorting_key(repo: CachedRepoInfo, sort_by: Optional[SortingOption_T] = None): + if sort_by == "alphabetical": + return (repo.repo_type, repo.repo_id.lower()) # by type then name + elif sort_by == "lastUpdated": + return -max(rev.last_modified for rev in repo.revisions) # newest first + elif sort_by == "lastUsed": + return -repo.last_accessed # most recently used first + elif sort_by == "size": + return -repo.size_on_disk # largest first + else: + return (repo.repo_type, repo.repo_id) # default stable order + + +@require_inquirer_py +def _manual_review_tui( + hf_cache_info: HFCacheInfo, + preselected: List[str], + sort_by: Optional[SortingOption_T] = None, +) -> List[str]: + """Ask the user for a manual review of the revisions to delete. + + Displays a multi-select menu in the terminal (TUI). + """ + # Define multiselect list + choices = _get_tui_choices_from_scan( + repos=hf_cache_info.repos, + preselected=preselected, + sort_by=sort_by, + ) + checkbox = inquirer.checkbox( + message="Select revisions to delete:", + choices=choices, # List of revisions with some pre-selection + cycle=False, # No loop between top and bottom + height=100, # Large list if possible + # We use the instruction to display to the user the expected effect of the + # deletion. + instruction=_get_expectations_str( + hf_cache_info, + selected_hashes=[c.value for c in choices if isinstance(c, Choice) and c.enabled], + ), + # We use the long instruction to should keybindings instructions to the user + long_instruction="Press to select, to validate and to quit without modification.", + # Message that is displayed once the user validates its selection. + transformer=lambda result: f"{len(result)} revision(s) selected.", + ) + + # Add a callback to update the information line when a revision is + # selected/unselected + def _update_expectations(_) -> None: + # Hacky way to dynamically set an instruction message to the checkbox when + # a revision hash is selected/unselected. + checkbox._instruction = _get_expectations_str( + hf_cache_info, + selected_hashes=[choice["value"] for choice in checkbox.content_control.choices if choice["enabled"]], + ) + + checkbox.kb_func_lookup["toggle"].append({"func": _update_expectations}) + + # Finally display the form to the user. + try: + return checkbox.execute() + except KeyboardInterrupt: + return [] # Quit without deletion + + +@require_inquirer_py +def _ask_for_confirmation_tui(message: str, default: bool = True) -> bool: + """Ask for confirmation using Inquirer.""" + return inquirer.confirm(message, default=default).execute() + + +def _get_tui_choices_from_scan( + repos: Iterable[CachedRepoInfo], + preselected: List[str], + sort_by: Optional[SortingOption_T] = None, +) -> List: + """Build a list of choices from the scanned repos. + + Args: + repos (*Iterable[`CachedRepoInfo`]*): + List of scanned repos on which we want to delete revisions. + preselected (*List[`str`]*): + List of revision hashes that will be preselected. + sort_by (*Optional[SortingOption_T]*): + Sorting direction. Choices: "alphabetical", "lastUpdated", "lastUsed", "size". + + Return: + The list of choices to pass to `inquirer.checkbox`. + """ + choices: List[Union[Choice, Separator]] = [] + + # First choice is to cancel the deletion + choices.append( + Choice( + _CANCEL_DELETION_STR, + name="None of the following (if selected, nothing will be deleted).", + enabled=False, + ) + ) + + # Sort repos based on specified criteria + sorted_repos = sorted(repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by)) + + for repo in sorted_repos: + # Repo as separator + choices.append( + Separator( + f"\n{repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str}," + f" used {repo.last_accessed_str})" + ) + ) + for revision in sorted(repo.revisions, key=_revision_sorting_order): + # Revision as choice + choices.append( + Choice( + revision.commit_hash, + name=( + f"{revision.commit_hash[:8]}:" + f" {', '.join(sorted(revision.refs)) or '(detached)'} #" + f" modified {revision.last_modified_str}" + ), + enabled=revision.commit_hash in preselected, + ) + ) + + # Return choices + return choices + + +def _manual_review_no_tui( + hf_cache_info: HFCacheInfo, + preselected: List[str], + sort_by: Optional[SortingOption_T] = None, +) -> List[str]: + """Ask the user for a manual review of the revisions to delete. + + Used when TUI is disabled. Manual review happens in a separate tmp file that the + user can manually edit. + """ + # 1. Generate temporary file with delete commands. + fd, tmp_path = mkstemp(suffix=".txt") # suffix to make it easier to find by editors + os.close(fd) + + lines = [] + + sorted_repos = sorted(hf_cache_info.repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by)) + + for repo in sorted_repos: + lines.append( + f"\n# {repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str}," + f" used {repo.last_accessed_str})" + ) + for revision in sorted(repo.revisions, key=_revision_sorting_order): + lines.append( + # Deselect by prepending a '#' + f"{'' if revision.commit_hash in preselected else '#'} " + f" {revision.commit_hash} # Refs:" + # Print `refs` as comment on same line + f" {', '.join(sorted(revision.refs)) or '(detached)'} # modified" + # Print `last_modified` as comment on same line + f" {revision.last_modified_str}" + ) + + with open(tmp_path, "w") as f: + f.write(_MANUAL_REVIEW_NO_TUI_INSTRUCTIONS) + f.write("\n".join(lines)) + + # 2. Prompt instructions to user. + instructions = f""" + TUI is disabled. In order to select which revisions you want to delete, please edit + the following file using the text editor of your choice. Instructions for manual + editing are located at the beginning of the file. Edit the file, save it and confirm + to continue. + File to edit: {ANSI.bold(tmp_path)} + """ + print("\n".join(line.strip() for line in instructions.strip().split("\n"))) + + # 3. Wait for user confirmation. + while True: + selected_hashes = _read_manual_review_tmp_file(tmp_path) + if _ask_for_confirmation_no_tui( + _get_expectations_str(hf_cache_info, selected_hashes) + " Continue ?", + default=False, + ): + break + + # 4. Return selected_hashes sorted to maintain stable order + os.remove(tmp_path) + return sorted(selected_hashes) # Sort to maintain stable order + + +def _ask_for_confirmation_no_tui(message: str, default: bool = True) -> bool: + """Ask for confirmation using pure-python.""" + YES = ("y", "yes", "1") + NO = ("n", "no", "0") + DEFAULT = "" + ALL = YES + NO + (DEFAULT,) + full_message = message + (" (Y/n) " if default else " (y/N) ") + while True: + answer = input(full_message).lower() + if answer == DEFAULT: + return default + if answer in YES: + return True + if answer in NO: + return False + print(f"Invalid input. Must be one of {ALL}") + + +def _get_expectations_str(hf_cache_info: HFCacheInfo, selected_hashes: List[str]) -> str: + """Format a string to display to the user how much space would be saved. + + Example: + ``` + >>> _get_expectations_str(hf_cache_info, selected_hashes) + '7 revisions selected counting for 4.3G.' + ``` + """ + if _CANCEL_DELETION_STR in selected_hashes: + return "Nothing will be deleted." + strategy = hf_cache_info.delete_revisions(*selected_hashes) + return f"{len(selected_hashes)} revisions selected counting for {strategy.expected_freed_size_str}." + + +def _read_manual_review_tmp_file(tmp_path: str) -> List[str]: + """Read the manually reviewed instruction file and return a list of revision hash. + + Example: + ```txt + # This is the tmp file content + ### + + # Commented out line + 123456789 # revision hash + + # Something else + # a_newer_hash # 2 days ago + an_older_hash # 3 days ago + ``` + + ```py + >>> _read_manual_review_tmp_file(tmp_path) + ['123456789', 'an_older_hash'] + ``` + """ + with open(tmp_path) as f: + content = f.read() + + # Split lines + lines = [line.strip() for line in content.split("\n")] + + # Filter commented lines + selected_lines = [line for line in lines if not line.startswith("#")] + + # Select only before comment + selected_hashes = [line.split("#")[0].strip() for line in selected_lines] + + # Return revision hashes + return [hash for hash in selected_hashes if len(hash) > 0] + + +_MANUAL_REVIEW_NO_TUI_INSTRUCTIONS = f""" +# INSTRUCTIONS +# ------------ +# This is a temporary file created by running `huggingface-cli delete-cache` with the +# `--disable-tui` option. It contains a set of revisions that can be deleted from your +# local cache directory. +# +# Please manually review the revisions you want to delete: +# - Revision hashes can be commented out with '#'. +# - Only non-commented revisions in this file will be deleted. +# - Revision hashes that are removed from this file are ignored as well. +# - If `{_CANCEL_DELETION_STR}` line is uncommented, the all cache deletion is cancelled and +# no changes will be applied. +# +# Once you've manually reviewed this file, please confirm deletion in the terminal. This +# file will be automatically removed once done. +# ------------ + +# KILL SWITCH +# ------------ +# Un-comment following line to completely cancel the deletion process +# {_CANCEL_DELETION_STR} +# ------------ + +# REVISIONS +# ------------ +""".strip() + + +def _revision_sorting_order(revision: CachedRevisionInfo) -> Any: + # Sort by last modified (oldest first) + return revision.last_modified diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/download.py b/lib/python3.12/site-packages/huggingface_hub/commands/download.py new file mode 100644 index 0000000000000000000000000000000000000000..10e22c3d1eb83dbb52c4a633fb66f19b3f35d8e7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/download.py @@ -0,0 +1,200 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to download files from the Hub with the CLI. + +Usage: + huggingface-cli download --help + + # Download file + huggingface-cli download gpt2 config.json + + # Download entire repo + huggingface-cli download fffiloni/zeroscope --repo-type=space --revision=refs/pr/78 + + # Download repo with filters + huggingface-cli download gpt2 --include="*.safetensors" + + # Download with token + huggingface-cli download Wauplin/private-model --token=hf_*** + + # Download quietly (no progress bar, no warnings, only the returned path) + huggingface-cli download gpt2 config.json --quiet + + # Download to local dir + huggingface-cli download gpt2 --local-dir=./models/gpt2 +""" + +import warnings +from argparse import Namespace, _SubParsersAction +from typing import List, Optional + +from huggingface_hub import logging +from huggingface_hub._snapshot_download import snapshot_download +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.file_download import hf_hub_download +from huggingface_hub.utils import disable_progress_bars, enable_progress_bars + + +logger = logging.get_logger(__name__) + + +class DownloadCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + download_parser = parser.add_parser("download", help="Download files from the Hub") + download_parser.add_argument( + "repo_id", type=str, help="ID of the repo to download from (e.g. `username/repo-name`)." + ) + download_parser.add_argument( + "filenames", type=str, nargs="*", help="Files to download (e.g. `config.json`, `data/metadata.jsonl`)." + ) + download_parser.add_argument( + "--repo-type", + choices=["model", "dataset", "space"], + default="model", + help="Type of repo to download from (defaults to 'model').", + ) + download_parser.add_argument( + "--revision", + type=str, + help="An optional Git revision id which can be a branch name, a tag, or a commit hash.", + ) + download_parser.add_argument( + "--include", nargs="*", type=str, help="Glob patterns to match files to download." + ) + download_parser.add_argument( + "--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to download." + ) + download_parser.add_argument( + "--cache-dir", type=str, help="Path to the directory where to save the downloaded files." + ) + download_parser.add_argument( + "--local-dir", + type=str, + help=( + "If set, the downloaded file will be placed under this directory. Check out" + " https://huggingface.co/docs/huggingface_hub/guides/download#download-files-to-local-folder for more" + " details." + ), + ) + download_parser.add_argument( + "--local-dir-use-symlinks", + choices=["auto", "True", "False"], + help=("Deprecated and ignored. Downloading to a local directory does not use symlinks anymore."), + ) + download_parser.add_argument( + "--force-download", + action="store_true", + help="If True, the files will be downloaded even if they are already cached.", + ) + download_parser.add_argument( + "--resume-download", + action="store_true", + help="Deprecated and ignored. Downloading a file to local dir always attempts to resume previously interrupted downloads (unless hf-transfer is enabled).", + ) + download_parser.add_argument( + "--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens" + ) + download_parser.add_argument( + "--quiet", + action="store_true", + help="If True, progress bars are disabled and only the path to the download files is printed.", + ) + download_parser.add_argument( + "--max-workers", + type=int, + default=8, + help="Maximum number of workers to use for downloading files. Default is 8.", + ) + download_parser.set_defaults(func=DownloadCommand) + + def __init__(self, args: Namespace) -> None: + self.token = args.token + self.repo_id: str = args.repo_id + self.filenames: List[str] = args.filenames + self.repo_type: str = args.repo_type + self.revision: Optional[str] = args.revision + self.include: Optional[List[str]] = args.include + self.exclude: Optional[List[str]] = args.exclude + self.cache_dir: Optional[str] = args.cache_dir + self.local_dir: Optional[str] = args.local_dir + self.force_download: bool = args.force_download + self.resume_download: Optional[bool] = args.resume_download or None + self.quiet: bool = args.quiet + self.max_workers: int = args.max_workers + + if args.local_dir_use_symlinks is not None: + warnings.warn( + "Ignoring --local-dir-use-symlinks. Downloading to a local directory does not use symlinks anymore.", + FutureWarning, + ) + + def run(self) -> None: + if self.quiet: + disable_progress_bars() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + print(self._download()) # Print path to downloaded files + enable_progress_bars() + else: + logging.set_verbosity_info() + print(self._download()) # Print path to downloaded files + logging.set_verbosity_warning() + + def _download(self) -> str: + # Warn user if patterns are ignored + if len(self.filenames) > 0: + if self.include is not None and len(self.include) > 0: + warnings.warn("Ignoring `--include` since filenames have being explicitly set.") + if self.exclude is not None and len(self.exclude) > 0: + warnings.warn("Ignoring `--exclude` since filenames have being explicitly set.") + + # Single file to download: use `hf_hub_download` + if len(self.filenames) == 1: + return hf_hub_download( + repo_id=self.repo_id, + repo_type=self.repo_type, + revision=self.revision, + filename=self.filenames[0], + cache_dir=self.cache_dir, + resume_download=self.resume_download, + force_download=self.force_download, + token=self.token, + local_dir=self.local_dir, + library_name="huggingface-cli", + ) + + # Otherwise: use `snapshot_download` to ensure all files comes from same revision + elif len(self.filenames) == 0: + allow_patterns = self.include + ignore_patterns = self.exclude + else: + allow_patterns = self.filenames + ignore_patterns = None + + return snapshot_download( + repo_id=self.repo_id, + repo_type=self.repo_type, + revision=self.revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + resume_download=self.resume_download, + force_download=self.force_download, + cache_dir=self.cache_dir, + token=self.token, + local_dir=self.local_dir, + library_name="huggingface-cli", + max_workers=self.max_workers, + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/env.py b/lib/python3.12/site-packages/huggingface_hub/commands/env.py new file mode 100644 index 0000000000000000000000000000000000000000..23f2828bbfebda0a633b4b3c6883432e4a534c79 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/env.py @@ -0,0 +1,36 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to print information about the environment. + +Usage: + huggingface-cli env +""" + +from argparse import _SubParsersAction + +from ..utils import dump_environment_info +from . import BaseHuggingfaceCLICommand + + +class EnvironmentCommand(BaseHuggingfaceCLICommand): + def __init__(self, args): + self.args = args + + @staticmethod + def register_subcommand(parser: _SubParsersAction): + env_parser = parser.add_parser("env", help="Print information about the environment.") + env_parser.set_defaults(func=EnvironmentCommand) + + def run(self) -> None: + dump_environment_info() diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/huggingface_cli.py b/lib/python3.12/site-packages/huggingface_hub/commands/huggingface_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..1e790b5eb1b40710072ef5fc2597b9e6c3325355 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/huggingface_cli.py @@ -0,0 +1,61 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from argparse import ArgumentParser + +from huggingface_hub.commands.delete_cache import DeleteCacheCommand +from huggingface_hub.commands.download import DownloadCommand +from huggingface_hub.commands.env import EnvironmentCommand +from huggingface_hub.commands.lfs import LfsCommands +from huggingface_hub.commands.repo_files import RepoFilesCommand +from huggingface_hub.commands.scan_cache import ScanCacheCommand +from huggingface_hub.commands.tag import TagCommands +from huggingface_hub.commands.upload import UploadCommand +from huggingface_hub.commands.upload_large_folder import UploadLargeFolderCommand +from huggingface_hub.commands.user import UserCommands +from huggingface_hub.commands.version import VersionCommand + + +def main(): + parser = ArgumentParser("huggingface-cli", usage="huggingface-cli []") + commands_parser = parser.add_subparsers(help="huggingface-cli command helpers") + + # Register commands + DownloadCommand.register_subcommand(commands_parser) + UploadCommand.register_subcommand(commands_parser) + RepoFilesCommand.register_subcommand(commands_parser) + EnvironmentCommand.register_subcommand(commands_parser) + UserCommands.register_subcommand(commands_parser) + LfsCommands.register_subcommand(commands_parser) + ScanCacheCommand.register_subcommand(commands_parser) + DeleteCacheCommand.register_subcommand(commands_parser) + TagCommands.register_subcommand(commands_parser) + VersionCommand.register_subcommand(commands_parser) + + # Experimental + UploadLargeFolderCommand.register_subcommand(commands_parser) + + # Let's go + args = parser.parse_args() + if not hasattr(args, "func"): + parser.print_help() + exit(1) + + # Run + service = args.func(args) + service.run() + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/lfs.py b/lib/python3.12/site-packages/huggingface_hub/commands/lfs.py new file mode 100644 index 0000000000000000000000000000000000000000..e510e345e6a4bf6da03f71b35cbfa2a4f0eb7325 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/lfs.py @@ -0,0 +1,200 @@ +""" +Implementation of a custom transfer agent for the transfer type "multipart" for +git-lfs. + +Inspired by: +github.com/cbartz/git-lfs-swift-transfer-agent/blob/master/git_lfs_swift_transfer.py + +Spec is: github.com/git-lfs/git-lfs/blob/master/docs/custom-transfers.md + + +To launch debugger while developing: + +``` [lfs "customtransfer.multipart"] +path = /path/to/huggingface_hub/.env/bin/python args = -m debugpy --listen 5678 +--wait-for-client +/path/to/huggingface_hub/src/huggingface_hub/commands/huggingface_cli.py +lfs-multipart-upload ```""" + +import json +import os +import subprocess +import sys +from argparse import _SubParsersAction +from typing import Dict, List, Optional + +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.lfs import LFS_MULTIPART_UPLOAD_COMMAND + +from ..utils import get_session, hf_raise_for_status, logging +from ..utils._lfs import SliceFileObj + + +logger = logging.get_logger(__name__) + + +class LfsCommands(BaseHuggingfaceCLICommand): + """ + Implementation of a custom transfer agent for the transfer type "multipart" + for git-lfs. This lets users upload large files >5GB 🔥. Spec for LFS custom + transfer agent is: + https://github.com/git-lfs/git-lfs/blob/master/docs/custom-transfers.md + + This introduces two commands to the CLI: + + 1. $ huggingface-cli lfs-enable-largefiles + + This should be executed once for each model repo that contains a model file + >5GB. It's documented in the error message you get if you just try to git + push a 5GB file without having enabled it before. + + 2. $ huggingface-cli lfs-multipart-upload + + This command is called by lfs directly and is not meant to be called by the + user. + """ + + @staticmethod + def register_subcommand(parser: _SubParsersAction): + enable_parser = parser.add_parser( + "lfs-enable-largefiles", help="Configure your repository to enable upload of files > 5GB." + ) + enable_parser.add_argument("path", type=str, help="Local path to repository you want to configure.") + enable_parser.set_defaults(func=lambda args: LfsEnableCommand(args)) + + # Command will get called by git-lfs, do not call it directly. + upload_parser = parser.add_parser(LFS_MULTIPART_UPLOAD_COMMAND, add_help=False) + upload_parser.set_defaults(func=lambda args: LfsUploadCommand(args)) + + +class LfsEnableCommand: + def __init__(self, args): + self.args = args + + def run(self): + local_path = os.path.abspath(self.args.path) + if not os.path.isdir(local_path): + print("This does not look like a valid git repo.") + exit(1) + subprocess.run( + "git config lfs.customtransfer.multipart.path huggingface-cli".split(), + check=True, + cwd=local_path, + ) + subprocess.run( + f"git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}".split(), + check=True, + cwd=local_path, + ) + print("Local repo set up for largefiles") + + +def write_msg(msg: Dict): + """Write out the message in Line delimited JSON.""" + msg_str = json.dumps(msg) + "\n" + sys.stdout.write(msg_str) + sys.stdout.flush() + + +def read_msg() -> Optional[Dict]: + """Read Line delimited JSON from stdin.""" + msg = json.loads(sys.stdin.readline().strip()) + + if "terminate" in (msg.get("type"), msg.get("event")): + # terminate message received + return None + + if msg.get("event") not in ("download", "upload"): + logger.critical("Received unexpected message") + sys.exit(1) + + return msg + + +class LfsUploadCommand: + def __init__(self, args) -> None: + self.args = args + + def run(self) -> None: + # Immediately after invoking a custom transfer process, git-lfs + # sends initiation data to the process over stdin. + # This tells the process useful information about the configuration. + init_msg = json.loads(sys.stdin.readline().strip()) + if not (init_msg.get("event") == "init" and init_msg.get("operation") == "upload"): + write_msg({"error": {"code": 32, "message": "Wrong lfs init operation"}}) + sys.exit(1) + + # The transfer process should use the information it needs from the + # initiation structure, and also perform any one-off setup tasks it + # needs to do. It should then respond on stdout with a simple empty + # confirmation structure, as follows: + write_msg({}) + + # After the initiation exchange, git-lfs will send any number of + # transfer requests to the stdin of the transfer process, in a serial sequence. + while True: + msg = read_msg() + if msg is None: + # When all transfers have been processed, git-lfs will send + # a terminate event to the stdin of the transfer process. + # On receiving this message the transfer process should + # clean up and terminate. No response is expected. + sys.exit(0) + + oid = msg["oid"] + filepath = msg["path"] + completion_url = msg["action"]["href"] + header = msg["action"]["header"] + chunk_size = int(header.pop("chunk_size")) + presigned_urls: List[str] = list(header.values()) + + # Send a "started" progress event to allow other workers to start. + # Otherwise they're delayed until first "progress" event is reported, + # i.e. after the first 5GB by default (!) + write_msg( + { + "event": "progress", + "oid": oid, + "bytesSoFar": 1, + "bytesSinceLast": 0, + } + ) + + parts = [] + with open(filepath, "rb") as file: + for i, presigned_url in enumerate(presigned_urls): + with SliceFileObj( + file, + seek_from=i * chunk_size, + read_limit=chunk_size, + ) as data: + r = get_session().put(presigned_url, data=data) + hf_raise_for_status(r) + parts.append( + { + "etag": r.headers.get("etag"), + "partNumber": i + 1, + } + ) + # In order to support progress reporting while data is uploading / downloading, + # the transfer process should post messages to stdout + write_msg( + { + "event": "progress", + "oid": oid, + "bytesSoFar": (i + 1) * chunk_size, + "bytesSinceLast": chunk_size, + } + ) + # Not precise but that's ok. + + r = get_session().post( + completion_url, + json={ + "oid": oid, + "parts": parts, + }, + ) + hf_raise_for_status(r) + + write_msg({"event": "complete", "oid": oid}) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/repo_files.py b/lib/python3.12/site-packages/huggingface_hub/commands/repo_files.py new file mode 100644 index 0000000000000000000000000000000000000000..f15bbed04f3634d7783d8230324cdaee44df4f59 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/repo_files.py @@ -0,0 +1,128 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to update or delete files in a repository using the CLI. + +Usage: + # delete all + huggingface-cli repo-files delete "*" + + # delete single file + huggingface-cli repo-files delete file.txt + + # delete single folder + huggingface-cli repo-files delete folder/ + + # delete multiple + huggingface-cli repo-files delete file.txt folder/ file2.txt + + # delete multiple patterns + huggingface-cli repo-files delete file.txt "*.json" "folder/*.parquet" + + # delete from different revision / repo-type + huggingface-cli repo-files delete file.txt --revision=refs/pr/1 --repo-type=dataset +""" + +from argparse import _SubParsersAction +from typing import List, Optional + +from huggingface_hub import logging +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.hf_api import HfApi + + +logger = logging.get_logger(__name__) + + +class DeleteFilesSubCommand: + def __init__(self, args) -> None: + self.args = args + self.repo_id: str = args.repo_id + self.repo_type: Optional[str] = args.repo_type + self.revision: Optional[str] = args.revision + self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli") + self.patterns: List[str] = args.patterns + self.commit_message: Optional[str] = args.commit_message + self.commit_description: Optional[str] = args.commit_description + self.create_pr: bool = args.create_pr + self.token: Optional[str] = args.token + + def run(self) -> None: + logging.set_verbosity_info() + url = self.api.delete_files( + delete_patterns=self.patterns, + repo_id=self.repo_id, + repo_type=self.repo_type, + revision=self.revision, + commit_message=self.commit_message, + commit_description=self.commit_description, + create_pr=self.create_pr, + ) + print(f"Files correctly deleted from repo. Commit: {url}.") + logging.set_verbosity_warning() + + +class RepoFilesCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + repo_files_parser = parser.add_parser("repo-files", help="Manage files in a repo on the Hub") + repo_files_parser.add_argument( + "repo_id", type=str, help="The ID of the repo to manage (e.g. `username/repo-name`)." + ) + repo_files_subparsers = repo_files_parser.add_subparsers( + help="Action to execute against the files.", + required=True, + ) + delete_subparser = repo_files_subparsers.add_parser( + "delete", + help="Delete files from a repo on the Hub", + ) + delete_subparser.set_defaults(func=lambda args: DeleteFilesSubCommand(args)) + delete_subparser.add_argument( + "patterns", + nargs="+", + type=str, + help="Glob patterns to match files to delete.", + ) + delete_subparser.add_argument( + "--repo-type", + choices=["model", "dataset", "space"], + default="model", + help="Type of the repo to upload to (e.g. `dataset`).", + ) + delete_subparser.add_argument( + "--revision", + type=str, + help=( + "An optional Git revision to push to. It can be a branch name " + "or a PR reference. If revision does not" + " exist and `--create-pr` is not set, a branch will be automatically created." + ), + ) + delete_subparser.add_argument( + "--commit-message", type=str, help="The summary / title / first line of the generated commit." + ) + delete_subparser.add_argument( + "--commit-description", type=str, help="The description of the generated commit." + ) + delete_subparser.add_argument( + "--create-pr", action="store_true", help="Whether to create a new Pull Request for these changes." + ) + repo_files_parser.add_argument( + "--token", + type=str, + help="A User Access Token generated from https://huggingface.co/settings/tokens", + ) + + repo_files_parser.set_defaults(func=RepoFilesCommand) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/scan_cache.py b/lib/python3.12/site-packages/huggingface_hub/commands/scan_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..799b9ba5523134a668aa0171e9f3668694299341 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/scan_cache.py @@ -0,0 +1,181 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to scan the HF cache directory. + +Usage: + huggingface-cli scan-cache + huggingface-cli scan-cache -v + huggingface-cli scan-cache -vvv + huggingface-cli scan-cache --dir ~/.cache/huggingface/hub +""" + +import time +from argparse import Namespace, _SubParsersAction +from typing import Optional + +from ..utils import CacheNotFound, HFCacheInfo, scan_cache_dir +from . import BaseHuggingfaceCLICommand +from ._cli_utils import ANSI, tabulate + + +class ScanCacheCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + scan_cache_parser = parser.add_parser("scan-cache", help="Scan cache directory.") + + scan_cache_parser.add_argument( + "--dir", + type=str, + default=None, + help="cache directory to scan (optional). Default to the default HuggingFace cache.", + ) + scan_cache_parser.add_argument( + "-v", + "--verbose", + action="count", + default=0, + help="show a more verbose output", + ) + scan_cache_parser.set_defaults(func=ScanCacheCommand) + + def __init__(self, args: Namespace) -> None: + self.verbosity: int = args.verbose + self.cache_dir: Optional[str] = args.dir + + def run(self): + try: + t0 = time.time() + hf_cache_info = scan_cache_dir(self.cache_dir) + t1 = time.time() + except CacheNotFound as exc: + cache_dir = exc.cache_dir + print(f"Cache directory not found: {cache_dir}") + return + + self._print_hf_cache_info_as_table(hf_cache_info) + + print( + f"\nDone in {round(t1 - t0, 1)}s. Scanned {len(hf_cache_info.repos)} repo(s)" + f" for a total of {ANSI.red(hf_cache_info.size_on_disk_str)}." + ) + if len(hf_cache_info.warnings) > 0: + message = f"Got {len(hf_cache_info.warnings)} warning(s) while scanning." + if self.verbosity >= 3: + print(ANSI.gray(message)) + for warning in hf_cache_info.warnings: + print(ANSI.gray(warning)) + else: + print(ANSI.gray(message + " Use -vvv to print details.")) + + def _print_hf_cache_info_as_table(self, hf_cache_info: HFCacheInfo) -> None: + print(get_table(hf_cache_info, verbosity=self.verbosity)) + + +def get_table(hf_cache_info: HFCacheInfo, *, verbosity: int = 0) -> str: + """Generate a table from the [`HFCacheInfo`] object. + + Pass `verbosity=0` to get a table with a single row per repo, with columns + "repo_id", "repo_type", "size_on_disk", "nb_files", "last_accessed", "last_modified", "refs", "local_path". + + Pass `verbosity=1` to get a table with a row per repo and revision (thus multiple rows can appear for a single repo), with columns + "repo_id", "repo_type", "revision", "size_on_disk", "nb_files", "last_modified", "refs", "local_path". + + Example: + ```py + >>> from huggingface_hub.utils import scan_cache_dir + >>> from huggingface_hub.commands.scan_cache import get_table + + >>> hf_cache_info = scan_cache_dir() + HFCacheInfo(...) + + >>> print(get_table(hf_cache_info, verbosity=0)) + REPO ID REPO TYPE SIZE ON DISK NB FILES LAST_ACCESSED LAST_MODIFIED REFS LOCAL PATH + --------------------------------------------------- --------- ------------ -------- ------------- ------------- ---- -------------------------------------------------------------------------------------------------- + roberta-base model 2.7M 5 1 day ago 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--roberta-base + suno/bark model 8.8K 1 1 week ago 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--suno--bark + t5-base model 893.8M 4 4 days ago 7 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-base + t5-large model 3.0G 4 5 weeks ago 5 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-large + + >>> print(get_table(hf_cache_info, verbosity=1)) + REPO ID REPO TYPE REVISION SIZE ON DISK NB FILES LAST_MODIFIED REFS LOCAL PATH + --------------------------------------------------- --------- ---------------------------------------- ------------ -------- ------------- ---- ----------------------------------------------------------------------------------------------------------------------------------------------------- + roberta-base model e2da8e2f811d1448a5b465c236feacd80ffbac7b 2.7M 5 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--roberta-base\\snapshots\\e2da8e2f811d1448a5b465c236feacd80ffbac7b + suno/bark model 70a8a7d34168586dc5d028fa9666aceade177992 8.8K 1 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--suno--bark\\snapshots\\70a8a7d34168586dc5d028fa9666aceade177992 + t5-base model a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1 893.8M 4 7 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-base\\snapshots\\a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1 + t5-large model 150ebc2c4b72291e770f58e6057481c8d2ed331a 3.0G 4 5 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-large\\snapshots\\150ebc2c4b72291e770f58e6057481c8d2ed331a ``` + ``` + + Args: + hf_cache_info ([`HFCacheInfo`]): + The HFCacheInfo object to print. + verbosity (`int`, *optional*): + The verbosity level. Defaults to 0. + + Returns: + `str`: The table as a string. + """ + if verbosity == 0: + return tabulate( + rows=[ + [ + repo.repo_id, + repo.repo_type, + "{:>12}".format(repo.size_on_disk_str), + repo.nb_files, + repo.last_accessed_str, + repo.last_modified_str, + ", ".join(sorted(repo.refs)), + str(repo.repo_path), + ] + for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path) + ], + headers=[ + "REPO ID", + "REPO TYPE", + "SIZE ON DISK", + "NB FILES", + "LAST_ACCESSED", + "LAST_MODIFIED", + "REFS", + "LOCAL PATH", + ], + ) + else: + return tabulate( + rows=[ + [ + repo.repo_id, + repo.repo_type, + revision.commit_hash, + "{:>12}".format(revision.size_on_disk_str), + revision.nb_files, + revision.last_modified_str, + ", ".join(sorted(revision.refs)), + str(revision.snapshot_path), + ] + for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path) + for revision in sorted(repo.revisions, key=lambda revision: revision.commit_hash) + ], + headers=[ + "REPO ID", + "REPO TYPE", + "REVISION", + "SIZE ON DISK", + "NB FILES", + "LAST_MODIFIED", + "REFS", + "LOCAL PATH", + ], + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/tag.py b/lib/python3.12/site-packages/huggingface_hub/commands/tag.py new file mode 100644 index 0000000000000000000000000000000000000000..c3beab90a0a2858906c848fd1e3f54edfb9d4864 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/tag.py @@ -0,0 +1,159 @@ +# coding=utf-8 +# Copyright 2024-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contains commands to perform tag management with the CLI. + +Usage Examples: + - Create a tag: + $ huggingface-cli tag user/my-model 1.0 --message "First release" + $ huggingface-cli tag user/my-model 1.0 -m "First release" --revision develop + $ huggingface-cli tag user/my-dataset 1.0 -m "First release" --repo-type dataset + $ huggingface-cli tag user/my-space 1.0 + - List all tags: + $ huggingface-cli tag -l user/my-model + $ huggingface-cli tag --list user/my-dataset --repo-type dataset + - Delete a tag: + $ huggingface-cli tag -d user/my-model 1.0 + $ huggingface-cli tag --delete user/my-dataset 1.0 --repo-type dataset + $ huggingface-cli tag -d user/my-space 1.0 -y +""" + +from argparse import Namespace, _SubParsersAction + +from requests.exceptions import HTTPError + +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.constants import ( + REPO_TYPES, +) +from huggingface_hub.hf_api import HfApi + +from ..errors import HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError +from ._cli_utils import ANSI + + +class TagCommands(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + tag_parser = parser.add_parser("tag", help="(create, list, delete) tags for a repo in the hub") + + tag_parser.add_argument("repo_id", type=str, help="The ID of the repo to tag (e.g. `username/repo-name`).") + tag_parser.add_argument("tag", nargs="?", type=str, help="The name of the tag for creation or deletion.") + tag_parser.add_argument("-m", "--message", type=str, help="The description of the tag to create.") + tag_parser.add_argument("--revision", type=str, help="The git revision to tag.") + tag_parser.add_argument( + "--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens." + ) + tag_parser.add_argument( + "--repo-type", + choices=["model", "dataset", "space"], + default="model", + help="Set the type of repository (model, dataset, or space).", + ) + tag_parser.add_argument("-y", "--yes", action="store_true", help="Answer Yes to prompts automatically.") + + tag_parser.add_argument("-l", "--list", action="store_true", help="List tags for a repository.") + tag_parser.add_argument("-d", "--delete", action="store_true", help="Delete a tag for a repository.") + + tag_parser.set_defaults(func=lambda args: handle_commands(args)) + + +def handle_commands(args: Namespace): + if args.list: + return TagListCommand(args) + elif args.delete: + return TagDeleteCommand(args) + else: + return TagCreateCommand(args) + + +class TagCommand: + def __init__(self, args: Namespace): + self.args = args + self.api = HfApi(token=self.args.token) + self.repo_id = self.args.repo_id + self.repo_type = self.args.repo_type + if self.repo_type not in REPO_TYPES: + print("Invalid repo --repo-type") + exit(1) + + +class TagCreateCommand(TagCommand): + def run(self): + print(f"You are about to create tag {ANSI.bold(self.args.tag)} on {self.repo_type} {ANSI.bold(self.repo_id)}") + + try: + self.api.create_tag( + repo_id=self.repo_id, + tag=self.args.tag, + tag_message=self.args.message, + revision=self.args.revision, + repo_type=self.repo_type, + ) + except RepositoryNotFoundError: + print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.") + exit(1) + except RevisionNotFoundError: + print(f"Revision {ANSI.bold(self.args.revision)} not found.") + exit(1) + except HfHubHTTPError as e: + if e.response.status_code == 409: + print(f"Tag {ANSI.bold(self.args.tag)} already exists on {ANSI.bold(self.repo_id)}") + exit(1) + raise e + + print(f"Tag {ANSI.bold(self.args.tag)} created on {ANSI.bold(self.repo_id)}") + + +class TagListCommand(TagCommand): + def run(self): + try: + refs = self.api.list_repo_refs( + repo_id=self.repo_id, + repo_type=self.repo_type, + ) + except RepositoryNotFoundError: + print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.") + exit(1) + except HTTPError as e: + print(e) + print(ANSI.red(e.response.text)) + exit(1) + if len(refs.tags) == 0: + print("No tags found") + exit(0) + print(f"Tags for {self.repo_type} {ANSI.bold(self.repo_id)}:") + for tag in refs.tags: + print(tag.name) + + +class TagDeleteCommand(TagCommand): + def run(self): + print(f"You are about to delete tag {ANSI.bold(self.args.tag)} on {self.repo_type} {ANSI.bold(self.repo_id)}") + + if not self.args.yes: + choice = input("Proceed? [Y/n] ").lower() + if choice not in ("", "y", "yes"): + print("Abort") + exit() + try: + self.api.delete_tag(repo_id=self.repo_id, tag=self.args.tag, repo_type=self.repo_type) + except RepositoryNotFoundError: + print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.") + exit(1) + except RevisionNotFoundError: + print(f"Tag {ANSI.bold(self.args.tag)} not found on {ANSI.bold(self.repo_id)}") + exit(1) + print(f"Tag {ANSI.bold(self.args.tag)} deleted on {ANSI.bold(self.repo_id)}") diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/upload.py b/lib/python3.12/site-packages/huggingface_hub/commands/upload.py new file mode 100644 index 0000000000000000000000000000000000000000..3d4caebd5fec1872986db3730d1ce87407511d21 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/upload.py @@ -0,0 +1,314 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to upload a repo or file with the CLI. + +Usage: + # Upload file (implicit) + huggingface-cli upload my-cool-model ./my-cool-model.safetensors + + # Upload file (explicit) + huggingface-cli upload my-cool-model ./my-cool-model.safetensors model.safetensors + + # Upload directory (implicit). If `my-cool-model/` is a directory it will be uploaded, otherwise an exception is raised. + huggingface-cli upload my-cool-model + + # Upload directory (explicit) + huggingface-cli upload my-cool-model ./models/my-cool-model . + + # Upload filtered directory (example: tensorboard logs except for the last run) + huggingface-cli upload my-cool-model ./model/training /logs --include "*.tfevents.*" --exclude "*20230905*" + + # Upload with wildcard + huggingface-cli upload my-cool-model "./model/training/*.safetensors" + + # Upload private dataset + huggingface-cli upload Wauplin/my-cool-dataset ./data . --repo-type=dataset --private + + # Upload with token + huggingface-cli upload Wauplin/my-cool-model --token=hf_**** + + # Sync local Space with Hub (upload new files, delete removed files) + huggingface-cli upload Wauplin/space-example --repo-type=space --exclude="/logs/*" --delete="*" --commit-message="Sync local Space with Hub" + + # Schedule commits every 30 minutes + huggingface-cli upload Wauplin/my-cool-model --every=30 +""" + +import os +import time +import warnings +from argparse import Namespace, _SubParsersAction +from typing import List, Optional + +from huggingface_hub import logging +from huggingface_hub._commit_scheduler import CommitScheduler +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.constants import HF_HUB_ENABLE_HF_TRANSFER +from huggingface_hub.errors import RevisionNotFoundError +from huggingface_hub.hf_api import HfApi +from huggingface_hub.utils import disable_progress_bars, enable_progress_bars +from huggingface_hub.utils._runtime import is_xet_available + + +logger = logging.get_logger(__name__) + + +class UploadCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + upload_parser = parser.add_parser("upload", help="Upload a file or a folder to a repo on the Hub") + upload_parser.add_argument( + "repo_id", type=str, help="The ID of the repo to upload to (e.g. `username/repo-name`)." + ) + upload_parser.add_argument( + "local_path", + nargs="?", + help="Local path to the file or folder to upload. Wildcard patterns are supported. Defaults to current directory.", + ) + upload_parser.add_argument( + "path_in_repo", + nargs="?", + help="Path of the file or folder in the repo. Defaults to the relative path of the file or folder.", + ) + upload_parser.add_argument( + "--repo-type", + choices=["model", "dataset", "space"], + default="model", + help="Type of the repo to upload to (e.g. `dataset`).", + ) + upload_parser.add_argument( + "--revision", + type=str, + help=( + "An optional Git revision to push to. It can be a branch name or a PR reference. If revision does not" + " exist and `--create-pr` is not set, a branch will be automatically created." + ), + ) + upload_parser.add_argument( + "--private", + action="store_true", + help=( + "Whether to create a private repo if repo doesn't exist on the Hub. Ignored if the repo already" + " exists." + ), + ) + upload_parser.add_argument("--include", nargs="*", type=str, help="Glob patterns to match files to upload.") + upload_parser.add_argument( + "--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to upload." + ) + upload_parser.add_argument( + "--delete", + nargs="*", + type=str, + help="Glob patterns for file to be deleted from the repo while committing.", + ) + upload_parser.add_argument( + "--commit-message", type=str, help="The summary / title / first line of the generated commit." + ) + upload_parser.add_argument("--commit-description", type=str, help="The description of the generated commit.") + upload_parser.add_argument( + "--create-pr", action="store_true", help="Whether to upload content as a new Pull Request." + ) + upload_parser.add_argument( + "--every", + type=float, + help="If set, a background job is scheduled to create commits every `every` minutes.", + ) + upload_parser.add_argument( + "--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens" + ) + upload_parser.add_argument( + "--quiet", + action="store_true", + help="If True, progress bars are disabled and only the path to the uploaded files is printed.", + ) + upload_parser.set_defaults(func=UploadCommand) + + def __init__(self, args: Namespace) -> None: + self.repo_id: str = args.repo_id + self.repo_type: Optional[str] = args.repo_type + self.revision: Optional[str] = args.revision + self.private: bool = args.private + + self.include: Optional[List[str]] = args.include + self.exclude: Optional[List[str]] = args.exclude + self.delete: Optional[List[str]] = args.delete + + self.commit_message: Optional[str] = args.commit_message + self.commit_description: Optional[str] = args.commit_description + self.create_pr: bool = args.create_pr + self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli") + self.quiet: bool = args.quiet # disable warnings and progress bars + + # Check `--every` is valid + if args.every is not None and args.every <= 0: + raise ValueError(f"`every` must be a positive value (got '{args.every}')") + self.every: Optional[float] = args.every + + # Resolve `local_path` and `path_in_repo` + repo_name: str = args.repo_id.split("/")[-1] # e.g. "Wauplin/my-cool-model" => "my-cool-model" + self.local_path: str + self.path_in_repo: str + + if args.local_path is not None and any(c in args.local_path for c in ["*", "?", "["]): + if args.include is not None: + raise ValueError("Cannot set `--include` when passing a `local_path` containing a wildcard.") + if args.path_in_repo is not None and args.path_in_repo != ".": + raise ValueError("Cannot set `path_in_repo` when passing a `local_path` containing a wildcard.") + self.local_path = "." + self.include = args.local_path + self.path_in_repo = "." + elif args.local_path is None and os.path.isfile(repo_name): + # Implicit case 1: user provided only a repo_id which happen to be a local file as well => upload it with same name + self.local_path = repo_name + self.path_in_repo = repo_name + elif args.local_path is None and os.path.isdir(repo_name): + # Implicit case 2: user provided only a repo_id which happen to be a local folder as well => upload it at root + self.local_path = repo_name + self.path_in_repo = "." + elif args.local_path is None: + # Implicit case 3: user provided only a repo_id that does not match a local file or folder + # => the user must explicitly provide a local_path => raise exception + raise ValueError(f"'{repo_name}' is not a local file or folder. Please set `local_path` explicitly.") + elif args.path_in_repo is None and os.path.isfile(args.local_path): + # Explicit local path to file, no path in repo => upload it at root with same name + self.local_path = args.local_path + self.path_in_repo = os.path.basename(args.local_path) + elif args.path_in_repo is None: + # Explicit local path to folder, no path in repo => upload at root + self.local_path = args.local_path + self.path_in_repo = "." + else: + # Finally, if both paths are explicit + self.local_path = args.local_path + self.path_in_repo = args.path_in_repo + + def run(self) -> None: + if self.quiet: + disable_progress_bars() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + print(self._upload()) + enable_progress_bars() + else: + logging.set_verbosity_info() + print(self._upload()) + logging.set_verbosity_warning() + + def _upload(self) -> str: + if os.path.isfile(self.local_path): + if self.include is not None and len(self.include) > 0: + warnings.warn("Ignoring `--include` since a single file is uploaded.") + if self.exclude is not None and len(self.exclude) > 0: + warnings.warn("Ignoring `--exclude` since a single file is uploaded.") + if self.delete is not None and len(self.delete) > 0: + warnings.warn("Ignoring `--delete` since a single file is uploaded.") + + if not is_xet_available() and not HF_HUB_ENABLE_HF_TRANSFER: + logger.info( + "Consider using `hf_transfer` for faster uploads. This solution comes with some limitations. See" + " https://huggingface.co/docs/huggingface_hub/hf_transfer for more details." + ) + + # Schedule commits if `every` is set + if self.every is not None: + if os.path.isfile(self.local_path): + # If file => watch entire folder + use allow_patterns + folder_path = os.path.dirname(self.local_path) + path_in_repo = ( + self.path_in_repo[: -len(self.local_path)] # remove filename from path_in_repo + if self.path_in_repo.endswith(self.local_path) + else self.path_in_repo + ) + allow_patterns = [self.local_path] + ignore_patterns = [] + else: + folder_path = self.local_path + path_in_repo = self.path_in_repo + allow_patterns = self.include or [] + ignore_patterns = self.exclude or [] + if self.delete is not None and len(self.delete) > 0: + warnings.warn("Ignoring `--delete` when uploading with scheduled commits.") + + scheduler = CommitScheduler( + folder_path=folder_path, + repo_id=self.repo_id, + repo_type=self.repo_type, + revision=self.revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + path_in_repo=path_in_repo, + private=self.private, + every=self.every, + hf_api=self.api, + ) + print(f"Scheduling commits every {self.every} minutes to {scheduler.repo_id}.") + try: # Block main thread until KeyboardInterrupt + while True: + time.sleep(100) + except KeyboardInterrupt: + scheduler.stop() + return "Stopped scheduled commits." + + # Otherwise, create repo and proceed with the upload + if not os.path.isfile(self.local_path) and not os.path.isdir(self.local_path): + raise FileNotFoundError(f"No such file or directory: '{self.local_path}'.") + repo_id = self.api.create_repo( + repo_id=self.repo_id, + repo_type=self.repo_type, + exist_ok=True, + private=self.private, + space_sdk="gradio" if self.repo_type == "space" else None, + # ^ We don't want it to fail when uploading to a Space => let's set Gradio by default. + # ^ I'd rather not add CLI args to set it explicitly as we already have `huggingface-cli repo create` for that. + ).repo_id + + # Check if branch already exists and if not, create it + if self.revision is not None and not self.create_pr: + try: + self.api.repo_info(repo_id=repo_id, repo_type=self.repo_type, revision=self.revision) + except RevisionNotFoundError: + logger.info(f"Branch '{self.revision}' not found. Creating it...") + self.api.create_branch(repo_id=repo_id, repo_type=self.repo_type, branch=self.revision, exist_ok=True) + # ^ `exist_ok=True` to avoid race concurrency issues + + # File-based upload + if os.path.isfile(self.local_path): + return self.api.upload_file( + path_or_fileobj=self.local_path, + path_in_repo=self.path_in_repo, + repo_id=repo_id, + repo_type=self.repo_type, + revision=self.revision, + commit_message=self.commit_message, + commit_description=self.commit_description, + create_pr=self.create_pr, + ) + + # Folder-based upload + else: + return self.api.upload_folder( + folder_path=self.local_path, + path_in_repo=self.path_in_repo, + repo_id=repo_id, + repo_type=self.repo_type, + revision=self.revision, + commit_message=self.commit_message, + commit_description=self.commit_description, + create_pr=self.create_pr, + allow_patterns=self.include, + ignore_patterns=self.exclude, + delete_patterns=self.delete, + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/upload_large_folder.py b/lib/python3.12/site-packages/huggingface_hub/commands/upload_large_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..61c12a9f62f8e12591d8db4c9defc50dd91db705 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/upload_large_folder.py @@ -0,0 +1,129 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to upload a large folder with the CLI.""" + +import os +from argparse import Namespace, _SubParsersAction +from typing import List, Optional + +from huggingface_hub import logging +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.hf_api import HfApi +from huggingface_hub.utils import disable_progress_bars + +from ._cli_utils import ANSI + + +logger = logging.get_logger(__name__) + + +class UploadLargeFolderCommand(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + subparser = parser.add_parser("upload-large-folder", help="Upload a large folder to a repo on the Hub") + subparser.add_argument( + "repo_id", type=str, help="The ID of the repo to upload to (e.g. `username/repo-name`)." + ) + subparser.add_argument("local_path", type=str, help="Local path to the file or folder to upload.") + subparser.add_argument( + "--repo-type", + choices=["model", "dataset", "space"], + help="Type of the repo to upload to (e.g. `dataset`).", + ) + subparser.add_argument( + "--revision", + type=str, + help=("An optional Git revision to push to. It can be a branch name or a PR reference."), + ) + subparser.add_argument( + "--private", + action="store_true", + help=( + "Whether to create a private repo if repo doesn't exist on the Hub. Ignored if the repo already exists." + ), + ) + subparser.add_argument("--include", nargs="*", type=str, help="Glob patterns to match files to upload.") + subparser.add_argument("--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to upload.") + subparser.add_argument( + "--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens" + ) + subparser.add_argument( + "--num-workers", type=int, help="Number of workers to use to hash, upload and commit files." + ) + subparser.add_argument("--no-report", action="store_true", help="Whether to disable regular status report.") + subparser.add_argument("--no-bars", action="store_true", help="Whether to disable progress bars.") + subparser.set_defaults(func=UploadLargeFolderCommand) + + def __init__(self, args: Namespace) -> None: + self.repo_id: str = args.repo_id + self.local_path: str = args.local_path + self.repo_type: str = args.repo_type + self.revision: Optional[str] = args.revision + self.private: bool = args.private + + self.include: Optional[List[str]] = args.include + self.exclude: Optional[List[str]] = args.exclude + + self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli") + + self.num_workers: Optional[int] = args.num_workers + self.no_report: bool = args.no_report + self.no_bars: bool = args.no_bars + + if not os.path.isdir(self.local_path): + raise ValueError("Large upload is only supported for folders.") + + def run(self) -> None: + logging.set_verbosity_info() + + print( + ANSI.yellow( + "You are about to upload a large folder to the Hub using `huggingface-cli upload-large-folder`. " + "This is a new feature so feedback is very welcome!\n" + "\n" + "A few things to keep in mind:\n" + " - Repository limits still apply: https://huggingface.co/docs/hub/repositories-recommendations\n" + " - Do not start several processes in parallel.\n" + " - You can interrupt and resume the process at any time. " + "The script will pick up where it left off except for partially uploaded files that would have to be entirely reuploaded.\n" + " - Do not upload the same folder to several repositories. If you need to do so, you must delete the `./.cache/huggingface/` folder first.\n" + "\n" + f"Some temporary metadata will be stored under `{self.local_path}/.cache/huggingface`.\n" + " - You must not modify those files manually.\n" + " - You must not delete the `./.cache/huggingface/` folder while a process is running.\n" + " - You can delete the `./.cache/huggingface/` folder to reinitialize the upload state when process is not running. Files will have to be hashed and preuploaded again, except for already committed files.\n" + "\n" + "If the process output is too verbose, you can disable the progress bars with `--no-bars`. " + "You can also entirely disable the status report with `--no-report`.\n" + "\n" + "For more details, run `huggingface-cli upload-large-folder --help` or check the documentation at " + "https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-large-folder." + ) + ) + + if self.no_bars: + disable_progress_bars() + + self.api.upload_large_folder( + repo_id=self.repo_id, + folder_path=self.local_path, + repo_type=self.repo_type, + revision=self.revision, + private=self.private, + allow_patterns=self.include, + ignore_patterns=self.exclude, + num_workers=self.num_workers, + print_report=not self.no_report, + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/user.py b/lib/python3.12/site-packages/huggingface_hub/commands/user.py new file mode 100644 index 0000000000000000000000000000000000000000..9741a219f17c951c9d20582a5756750b9b92630f --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/user.py @@ -0,0 +1,304 @@ +# Copyright 2020 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains commands to authenticate to the Hugging Face Hub and interact with your repositories. + +Usage: + # login and save token locally. + huggingface-cli login --token=hf_*** --add-to-git-credential + + # switch between tokens + huggingface-cli auth switch + + # list all tokens + huggingface-cli auth list + + # logout from a specific token, if no token-name is provided, all tokens will be deleted from your machine. + huggingface-cli logout --token-name=your_token_name + + # find out which huggingface.co account you are logged in as + huggingface-cli whoami + + # create a new dataset repo on the Hub + huggingface-cli repo create mydataset --type=dataset + +""" + +import subprocess +from argparse import _SubParsersAction +from typing import List, Optional + +from requests.exceptions import HTTPError + +from huggingface_hub.commands import BaseHuggingfaceCLICommand +from huggingface_hub.constants import ENDPOINT, REPO_TYPES, REPO_TYPES_URL_PREFIXES, SPACES_SDK_TYPES +from huggingface_hub.hf_api import HfApi + +from .._login import ( # noqa: F401 # for backward compatibility # noqa: F401 # for backward compatibility + NOTEBOOK_LOGIN_PASSWORD_HTML, + NOTEBOOK_LOGIN_TOKEN_HTML_END, + NOTEBOOK_LOGIN_TOKEN_HTML_START, + auth_list, + auth_switch, + login, + logout, + notebook_login, +) +from ..utils import get_stored_tokens, get_token, logging +from ._cli_utils import ANSI + + +logger = logging.get_logger(__name__) + +try: + from InquirerPy import inquirer + from InquirerPy.base.control import Choice + + _inquirer_py_available = True +except ImportError: + _inquirer_py_available = False + + +class UserCommands(BaseHuggingfaceCLICommand): + @staticmethod + def register_subcommand(parser: _SubParsersAction): + login_parser = parser.add_parser("login", help="Log in using a token from huggingface.co/settings/tokens") + login_parser.add_argument( + "--token", + type=str, + help="Token generated from https://huggingface.co/settings/tokens", + ) + login_parser.add_argument( + "--add-to-git-credential", + action="store_true", + help="Optional: Save token to git credential helper.", + ) + login_parser.set_defaults(func=lambda args: LoginCommand(args)) + whoami_parser = parser.add_parser("whoami", help="Find out which huggingface.co account you are logged in as.") + whoami_parser.set_defaults(func=lambda args: WhoamiCommand(args)) + + logout_parser = parser.add_parser("logout", help="Log out") + logout_parser.add_argument( + "--token-name", + type=str, + help="Optional: Name of the access token to log out from.", + ) + logout_parser.set_defaults(func=lambda args: LogoutCommand(args)) + + auth_parser = parser.add_parser("auth", help="Other authentication related commands") + auth_subparsers = auth_parser.add_subparsers(help="Authentication subcommands") + auth_switch_parser = auth_subparsers.add_parser("switch", help="Switch between access tokens") + auth_switch_parser.add_argument( + "--token-name", + type=str, + help="Optional: Name of the access token to switch to.", + ) + auth_switch_parser.add_argument( + "--add-to-git-credential", + action="store_true", + help="Optional: Save token to git credential helper.", + ) + auth_switch_parser.set_defaults(func=lambda args: AuthSwitchCommand(args)) + auth_list_parser = auth_subparsers.add_parser("list", help="List all stored access tokens") + auth_list_parser.set_defaults(func=lambda args: AuthListCommand(args)) + # new system: git-based repo system + repo_parser = parser.add_parser("repo", help="{create} Commands to interact with your huggingface.co repos.") + repo_subparsers = repo_parser.add_subparsers(help="huggingface.co repos related commands") + repo_create_parser = repo_subparsers.add_parser("create", help="Create a new repo on huggingface.co") + repo_create_parser.add_argument( + "name", + type=str, + help="Name for your repo. Will be namespaced under your username to build the repo id.", + ) + repo_create_parser.add_argument( + "--type", + type=str, + help='Optional: repo_type: set to "dataset" or "space" if creating a dataset or space, default is model.', + ) + repo_create_parser.add_argument("--organization", type=str, help="Optional: organization namespace.") + repo_create_parser.add_argument( + "--space_sdk", + type=str, + help='Optional: Hugging Face Spaces SDK type. Required when --type is set to "space".', + choices=SPACES_SDK_TYPES, + ) + repo_create_parser.add_argument( + "-y", + "--yes", + action="store_true", + help="Optional: answer Yes to the prompt", + ) + repo_create_parser.set_defaults(func=lambda args: RepoCreateCommand(args)) + + +class BaseUserCommand: + def __init__(self, args): + self.args = args + self._api = HfApi() + + +class LoginCommand(BaseUserCommand): + def run(self): + logging.set_verbosity_info() + login( + token=self.args.token, + add_to_git_credential=self.args.add_to_git_credential, + ) + + +class LogoutCommand(BaseUserCommand): + def run(self): + logging.set_verbosity_info() + logout(token_name=self.args.token_name) + + +class AuthSwitchCommand(BaseUserCommand): + def run(self): + logging.set_verbosity_info() + token_name = self.args.token_name + if token_name is None: + token_name = self._select_token_name() + + if token_name is None: + print("No token name provided. Aborting.") + exit() + auth_switch(token_name, add_to_git_credential=self.args.add_to_git_credential) + + def _select_token_name(self) -> Optional[str]: + token_names = list(get_stored_tokens().keys()) + + if not token_names: + logger.error("No stored tokens found. Please login first.") + return None + + if _inquirer_py_available: + return self._select_token_name_tui(token_names) + # if inquirer is not available, use a simpler terminal UI + print("Available stored tokens:") + for i, token_name in enumerate(token_names, 1): + print(f"{i}. {token_name}") + while True: + try: + choice = input("Enter the number of the token to switch to (or 'q' to quit): ") + if choice.lower() == "q": + return None + index = int(choice) - 1 + if 0 <= index < len(token_names): + return token_names[index] + else: + print("Invalid selection. Please try again.") + except ValueError: + print("Invalid input. Please enter a number or 'q' to quit.") + + def _select_token_name_tui(self, token_names: List[str]) -> Optional[str]: + choices = [Choice(token_name, name=token_name) for token_name in token_names] + try: + return inquirer.select( + message="Select a token to switch to:", + choices=choices, + default=None, + ).execute() + except KeyboardInterrupt: + logger.info("Token selection cancelled.") + return None + + +class AuthListCommand(BaseUserCommand): + def run(self): + logging.set_verbosity_info() + auth_list() + + +class WhoamiCommand(BaseUserCommand): + def run(self): + token = get_token() + if token is None: + print("Not logged in") + exit() + try: + info = self._api.whoami(token) + print(info["name"]) + orgs = [org["name"] for org in info["orgs"]] + if orgs: + print(ANSI.bold("orgs: "), ",".join(orgs)) + + if ENDPOINT != "https://huggingface.co": + print(f"Authenticated through private endpoint: {ENDPOINT}") + except HTTPError as e: + print(e) + print(ANSI.red(e.response.text)) + exit(1) + + +class RepoCreateCommand(BaseUserCommand): + def run(self): + token = get_token() + if token is None: + print("Not logged in") + exit(1) + try: + stdout = subprocess.check_output(["git", "--version"]).decode("utf-8") + print(ANSI.gray(stdout.strip())) + except FileNotFoundError: + print("Looks like you do not have git installed, please install.") + + try: + stdout = subprocess.check_output(["git-lfs", "--version"]).decode("utf-8") + print(ANSI.gray(stdout.strip())) + except FileNotFoundError: + print( + ANSI.red( + "Looks like you do not have git-lfs installed, please install." + " You can install from https://git-lfs.github.com/." + " Then run `git lfs install` (you only have to do this once)." + ) + ) + print("") + + user = self._api.whoami(token)["name"] + namespace = self.args.organization if self.args.organization is not None else user + + repo_id = f"{namespace}/{self.args.name}" + + if self.args.type not in REPO_TYPES: + print("Invalid repo --type") + exit(1) + + if self.args.type in REPO_TYPES_URL_PREFIXES: + prefixed_repo_id = REPO_TYPES_URL_PREFIXES[self.args.type] + repo_id + else: + prefixed_repo_id = repo_id + + print(f"You are about to create {ANSI.bold(prefixed_repo_id)}") + + if not self.args.yes: + choice = input("Proceed? [Y/n] ").lower() + if not (choice == "" or choice == "y" or choice == "yes"): + print("Abort") + exit() + try: + url = self._api.create_repo( + repo_id=repo_id, + token=token, + repo_type=self.args.type, + space_sdk=self.args.space_sdk, + ) + except HTTPError as e: + print(e) + print(ANSI.red(e.response.text)) + exit(1) + print("\nYour repo now lives at:") + print(f" {ANSI.bold(url)}") + print("\nYou can clone it locally with the command below, and commit/push as usual.") + print(f"\n git clone {url}") + print("") diff --git a/lib/python3.12/site-packages/huggingface_hub/commands/version.py b/lib/python3.12/site-packages/huggingface_hub/commands/version.py new file mode 100644 index 0000000000000000000000000000000000000000..f7e866b76f1dcbfbb90a4ec494c47cf3d61c17dd --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/commands/version.py @@ -0,0 +1,37 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains command to print information about the version. + +Usage: + huggingface-cli version +""" + +from argparse import _SubParsersAction + +from huggingface_hub import __version__ + +from . import BaseHuggingfaceCLICommand + + +class VersionCommand(BaseHuggingfaceCLICommand): + def __init__(self, args): + self.args = args + + @staticmethod + def register_subcommand(parser: _SubParsersAction): + version_parser = parser.add_parser("version", help="Print information about the huggingface-cli version.") + version_parser.set_defaults(func=VersionCommand) + + def run(self) -> None: + print(f"huggingface_hub version: {__version__}") diff --git a/lib/python3.12/site-packages/huggingface_hub/community.py b/lib/python3.12/site-packages/huggingface_hub/community.py new file mode 100644 index 0000000000000000000000000000000000000000..16f2f02428dd5c2ce6437534af0397801bda45c5 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/community.py @@ -0,0 +1,355 @@ +""" +Data structures to interact with Discussions and Pull Requests on the Hub. + +See [the Discussions and Pull Requests guide](https://huggingface.co/docs/hub/repositories-pull-requests-discussions) +for more information on Pull Requests, Discussions, and the community tab. +""" + +from dataclasses import dataclass +from datetime import datetime +from typing import List, Literal, Optional, Union + +from . import constants +from .utils import parse_datetime + + +DiscussionStatus = Literal["open", "closed", "merged", "draft"] + + +@dataclass +class Discussion: + """ + A Discussion or Pull Request on the Hub. + + This dataclass is not intended to be instantiated directly. + + Attributes: + title (`str`): + The title of the Discussion / Pull Request + status (`str`): + The status of the Discussion / Pull Request. + It must be one of: + * `"open"` + * `"closed"` + * `"merged"` (only for Pull Requests ) + * `"draft"` (only for Pull Requests ) + num (`int`): + The number of the Discussion / Pull Request. + repo_id (`str`): + The id (`"{namespace}/{repo_name}"`) of the repo on which + the Discussion / Pull Request was open. + repo_type (`str`): + The type of the repo on which the Discussion / Pull Request was open. + Possible values are: `"model"`, `"dataset"`, `"space"`. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + is_pull_request (`bool`): + Whether or not this is a Pull Request. + created_at (`datetime`): + The `datetime` of creation of the Discussion / Pull Request. + endpoint (`str`): + Endpoint of the Hub. Default is https://huggingface.co. + git_reference (`str`, *optional*): + (property) Git reference to which changes can be pushed if this is a Pull Request, `None` otherwise. + url (`str`): + (property) URL of the discussion on the Hub. + """ + + title: str + status: DiscussionStatus + num: int + repo_id: str + repo_type: str + author: str + is_pull_request: bool + created_at: datetime + endpoint: str + + @property + def git_reference(self) -> Optional[str]: + """ + If this is a Pull Request , returns the git reference to which changes can be pushed. + Returns `None` otherwise. + """ + if self.is_pull_request: + return f"refs/pr/{self.num}" + return None + + @property + def url(self) -> str: + """Returns the URL of the discussion on the Hub.""" + if self.repo_type is None or self.repo_type == constants.REPO_TYPE_MODEL: + return f"{self.endpoint}/{self.repo_id}/discussions/{self.num}" + return f"{self.endpoint}/{self.repo_type}s/{self.repo_id}/discussions/{self.num}" + + +@dataclass +class DiscussionWithDetails(Discussion): + """ + Subclass of [`Discussion`]. + + Attributes: + title (`str`): + The title of the Discussion / Pull Request + status (`str`): + The status of the Discussion / Pull Request. + It can be one of: + * `"open"` + * `"closed"` + * `"merged"` (only for Pull Requests ) + * `"draft"` (only for Pull Requests ) + num (`int`): + The number of the Discussion / Pull Request. + repo_id (`str`): + The id (`"{namespace}/{repo_name}"`) of the repo on which + the Discussion / Pull Request was open. + repo_type (`str`): + The type of the repo on which the Discussion / Pull Request was open. + Possible values are: `"model"`, `"dataset"`, `"space"`. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + is_pull_request (`bool`): + Whether or not this is a Pull Request. + created_at (`datetime`): + The `datetime` of creation of the Discussion / Pull Request. + events (`list` of [`DiscussionEvent`]) + The list of [`DiscussionEvents`] in this Discussion or Pull Request. + conflicting_files (`Union[List[str], bool, None]`, *optional*): + A list of conflicting files if this is a Pull Request. + `None` if `self.is_pull_request` is `False`. + `True` if there are conflicting files but the list can't be retrieved. + target_branch (`str`, *optional*): + The branch into which changes are to be merged if this is a + Pull Request . `None` if `self.is_pull_request` is `False`. + merge_commit_oid (`str`, *optional*): + If this is a merged Pull Request , this is set to the OID / SHA of + the merge commit, `None` otherwise. + diff (`str`, *optional*): + The git diff if this is a Pull Request , `None` otherwise. + endpoint (`str`): + Endpoint of the Hub. Default is https://huggingface.co. + git_reference (`str`, *optional*): + (property) Git reference to which changes can be pushed if this is a Pull Request, `None` otherwise. + url (`str`): + (property) URL of the discussion on the Hub. + """ + + events: List["DiscussionEvent"] + conflicting_files: Union[List[str], bool, None] + target_branch: Optional[str] + merge_commit_oid: Optional[str] + diff: Optional[str] + + +@dataclass +class DiscussionEvent: + """ + An event in a Discussion or Pull Request. + + Use concrete classes: + * [`DiscussionComment`] + * [`DiscussionStatusChange`] + * [`DiscussionCommit`] + * [`DiscussionTitleChange`] + + Attributes: + id (`str`): + The ID of the event. An hexadecimal string. + type (`str`): + The type of the event. + created_at (`datetime`): + A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime) + object holding the creation timestamp for the event. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + """ + + id: str + type: str + created_at: datetime + author: str + + _event: dict + """Stores the original event data, in case we need to access it later.""" + + +@dataclass +class DiscussionComment(DiscussionEvent): + """A comment in a Discussion / Pull Request. + + Subclass of [`DiscussionEvent`]. + + + Attributes: + id (`str`): + The ID of the event. An hexadecimal string. + type (`str`): + The type of the event. + created_at (`datetime`): + A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime) + object holding the creation timestamp for the event. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + content (`str`): + The raw markdown content of the comment. Mentions, links and images are not rendered. + edited (`bool`): + Whether or not this comment has been edited. + hidden (`bool`): + Whether or not this comment has been hidden. + """ + + content: str + edited: bool + hidden: bool + + @property + def rendered(self) -> str: + """The rendered comment, as a HTML string""" + return self._event["data"]["latest"]["html"] + + @property + def last_edited_at(self) -> datetime: + """The last edit time, as a `datetime` object.""" + return parse_datetime(self._event["data"]["latest"]["updatedAt"]) + + @property + def last_edited_by(self) -> str: + """The last edit time, as a `datetime` object.""" + return self._event["data"]["latest"].get("author", {}).get("name", "deleted") + + @property + def edit_history(self) -> List[dict]: + """The edit history of the comment""" + return self._event["data"]["history"] + + @property + def number_of_edits(self) -> int: + return len(self.edit_history) + + +@dataclass +class DiscussionStatusChange(DiscussionEvent): + """A change of status in a Discussion / Pull Request. + + Subclass of [`DiscussionEvent`]. + + Attributes: + id (`str`): + The ID of the event. An hexadecimal string. + type (`str`): + The type of the event. + created_at (`datetime`): + A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime) + object holding the creation timestamp for the event. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + new_status (`str`): + The status of the Discussion / Pull Request after the change. + It can be one of: + * `"open"` + * `"closed"` + * `"merged"` (only for Pull Requests ) + """ + + new_status: str + + +@dataclass +class DiscussionCommit(DiscussionEvent): + """A commit in a Pull Request. + + Subclass of [`DiscussionEvent`]. + + Attributes: + id (`str`): + The ID of the event. An hexadecimal string. + type (`str`): + The type of the event. + created_at (`datetime`): + A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime) + object holding the creation timestamp for the event. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + summary (`str`): + The summary of the commit. + oid (`str`): + The OID / SHA of the commit, as a hexadecimal string. + """ + + summary: str + oid: str + + +@dataclass +class DiscussionTitleChange(DiscussionEvent): + """A rename event in a Discussion / Pull Request. + + Subclass of [`DiscussionEvent`]. + + Attributes: + id (`str`): + The ID of the event. An hexadecimal string. + type (`str`): + The type of the event. + created_at (`datetime`): + A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime) + object holding the creation timestamp for the event. + author (`str`): + The username of the Discussion / Pull Request author. + Can be `"deleted"` if the user has been deleted since. + old_title (`str`): + The previous title for the Discussion / Pull Request. + new_title (`str`): + The new title. + """ + + old_title: str + new_title: str + + +def deserialize_event(event: dict) -> DiscussionEvent: + """Instantiates a [`DiscussionEvent`] from a dict""" + event_id: str = event["id"] + event_type: str = event["type"] + created_at = parse_datetime(event["createdAt"]) + + common_args = dict( + id=event_id, + type=event_type, + created_at=created_at, + author=event.get("author", {}).get("name", "deleted"), + _event=event, + ) + + if event_type == "comment": + return DiscussionComment( + **common_args, + edited=event["data"]["edited"], + hidden=event["data"]["hidden"], + content=event["data"]["latest"]["raw"], + ) + if event_type == "status-change": + return DiscussionStatusChange( + **common_args, + new_status=event["data"]["status"], + ) + if event_type == "commit": + return DiscussionCommit( + **common_args, + summary=event["data"]["subject"], + oid=event["data"]["oid"], + ) + if event_type == "title-change": + return DiscussionTitleChange( + **common_args, + old_title=event["data"]["from"], + new_title=event["data"]["to"], + ) + + return DiscussionEvent(**common_args) diff --git a/lib/python3.12/site-packages/huggingface_hub/constants.py b/lib/python3.12/site-packages/huggingface_hub/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..33514f8aabcc6ff9eef6f2aa80b14d4182cc1ba7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/constants.py @@ -0,0 +1,285 @@ +import os +import re +import typing +from typing import Literal, Optional, Tuple + + +# Possible values for env variables + + +ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"}) + + +def _is_true(value: Optional[str]) -> bool: + if value is None: + return False + return value.upper() in ENV_VARS_TRUE_VALUES + + +def _as_int(value: Optional[str]) -> Optional[int]: + if value is None: + return None + return int(value) + + +# Constants for file downloads + +PYTORCH_WEIGHTS_NAME = "pytorch_model.bin" +TF2_WEIGHTS_NAME = "tf_model.h5" +TF_WEIGHTS_NAME = "model.ckpt" +FLAX_WEIGHTS_NAME = "flax_model.msgpack" +CONFIG_NAME = "config.json" +REPOCARD_NAME = "README.md" +DEFAULT_ETAG_TIMEOUT = 10 +DEFAULT_DOWNLOAD_TIMEOUT = 10 +DEFAULT_REQUEST_TIMEOUT = 10 +DOWNLOAD_CHUNK_SIZE = 10 * 1024 * 1024 +HF_TRANSFER_CONCURRENCY = 100 +MAX_HTTP_DOWNLOAD_SIZE = 50 * 1000 * 1000 * 1000 # 50 GB + +# Constants for serialization + +PYTORCH_WEIGHTS_FILE_PATTERN = "pytorch_model{suffix}.bin" # Unsafe pickle: use safetensors instead +SAFETENSORS_WEIGHTS_FILE_PATTERN = "model{suffix}.safetensors" +TF2_WEIGHTS_FILE_PATTERN = "tf_model{suffix}.h5" + +# Constants for safetensors repos + +SAFETENSORS_SINGLE_FILE = "model.safetensors" +SAFETENSORS_INDEX_FILE = "model.safetensors.index.json" +SAFETENSORS_MAX_HEADER_LENGTH = 25_000_000 + +# Timeout of aquiring file lock and logging the attempt +FILELOCK_LOG_EVERY_SECONDS = 10 + +# Git-related constants + +DEFAULT_REVISION = "main" +REGEX_COMMIT_OID = re.compile(r"[A-Fa-f0-9]{5,40}") + +HUGGINGFACE_CO_URL_HOME = "https://huggingface.co/" + +_staging_mode = _is_true(os.environ.get("HUGGINGFACE_CO_STAGING")) + +_HF_DEFAULT_ENDPOINT = "https://huggingface.co" +_HF_DEFAULT_STAGING_ENDPOINT = "https://hub-ci.huggingface.co" +ENDPOINT = os.getenv("HF_ENDPOINT", _HF_DEFAULT_ENDPOINT).rstrip("/") +HUGGINGFACE_CO_URL_TEMPLATE = ENDPOINT + "/{repo_id}/resolve/{revision}/{filename}" + +if _staging_mode: + ENDPOINT = _HF_DEFAULT_STAGING_ENDPOINT + HUGGINGFACE_CO_URL_TEMPLATE = _HF_DEFAULT_STAGING_ENDPOINT + "/{repo_id}/resolve/{revision}/{filename}" + +HUGGINGFACE_HEADER_X_REPO_COMMIT = "X-Repo-Commit" +HUGGINGFACE_HEADER_X_LINKED_ETAG = "X-Linked-Etag" +HUGGINGFACE_HEADER_X_LINKED_SIZE = "X-Linked-Size" +HUGGINGFACE_HEADER_X_BILL_TO = "X-HF-Bill-To" + +INFERENCE_ENDPOINT = os.environ.get("HF_INFERENCE_ENDPOINT", "https://api-inference.huggingface.co") + +# See https://huggingface.co/docs/inference-endpoints/index +INFERENCE_ENDPOINTS_ENDPOINT = "https://api.endpoints.huggingface.cloud/v2" +INFERENCE_CATALOG_ENDPOINT = "https://endpoints.huggingface.co/api/catalog" + +# See https://api.endpoints.huggingface.cloud/#post-/v2/endpoint/-namespace- +INFERENCE_ENDPOINT_IMAGE_KEYS = [ + "custom", + "huggingface", + "huggingfaceNeuron", + "llamacpp", + "tei", + "tgi", + "tgiNeuron", +] + +# Proxy for third-party providers +INFERENCE_PROXY_TEMPLATE = "https://router.huggingface.co/{provider}" + +REPO_ID_SEPARATOR = "--" +# ^ this substring is not allowed in repo_ids on hf.co +# and is the canonical one we use for serialization of repo ids elsewhere. + + +REPO_TYPE_DATASET = "dataset" +REPO_TYPE_SPACE = "space" +REPO_TYPE_MODEL = "model" +REPO_TYPES = [None, REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE] +SPACES_SDK_TYPES = ["gradio", "streamlit", "docker", "static"] + +REPO_TYPES_URL_PREFIXES = { + REPO_TYPE_DATASET: "datasets/", + REPO_TYPE_SPACE: "spaces/", +} +REPO_TYPES_MAPPING = { + "datasets": REPO_TYPE_DATASET, + "spaces": REPO_TYPE_SPACE, + "models": REPO_TYPE_MODEL, +} + +DiscussionTypeFilter = Literal["all", "discussion", "pull_request"] +DISCUSSION_TYPES: Tuple[DiscussionTypeFilter, ...] = typing.get_args(DiscussionTypeFilter) +DiscussionStatusFilter = Literal["all", "open", "closed"] +DISCUSSION_STATUS: Tuple[DiscussionTypeFilter, ...] = typing.get_args(DiscussionStatusFilter) + +# Webhook subscription types +WEBHOOK_DOMAIN_T = Literal["repo", "discussions"] + +# default cache +default_home = os.path.join(os.path.expanduser("~"), ".cache") +HF_HOME = os.path.expandvars( + os.path.expanduser( + os.getenv( + "HF_HOME", + os.path.join(os.getenv("XDG_CACHE_HOME", default_home), "huggingface"), + ) + ) +) +hf_cache_home = HF_HOME # for backward compatibility. TODO: remove this in 1.0.0 + +default_cache_path = os.path.join(HF_HOME, "hub") +default_assets_cache_path = os.path.join(HF_HOME, "assets") + +# Legacy env variables +HUGGINGFACE_HUB_CACHE = os.getenv("HUGGINGFACE_HUB_CACHE", default_cache_path) +HUGGINGFACE_ASSETS_CACHE = os.getenv("HUGGINGFACE_ASSETS_CACHE", default_assets_cache_path) + +# New env variables +HF_HUB_CACHE = os.path.expandvars( + os.path.expanduser( + os.getenv( + "HF_HUB_CACHE", + HUGGINGFACE_HUB_CACHE, + ) + ) +) +HF_ASSETS_CACHE = os.path.expandvars( + os.path.expanduser( + os.getenv( + "HF_ASSETS_CACHE", + HUGGINGFACE_ASSETS_CACHE, + ) + ) +) + +HF_HUB_OFFLINE = _is_true(os.environ.get("HF_HUB_OFFLINE") or os.environ.get("TRANSFORMERS_OFFLINE")) + +# If set, log level will be set to DEBUG and all requests made to the Hub will be logged +# as curl commands for reproducibility. +HF_DEBUG = _is_true(os.environ.get("HF_DEBUG")) + +# Opt-out from telemetry requests +HF_HUB_DISABLE_TELEMETRY = ( + _is_true(os.environ.get("HF_HUB_DISABLE_TELEMETRY")) # HF-specific env variable + or _is_true(os.environ.get("DISABLE_TELEMETRY")) + or _is_true(os.environ.get("DO_NOT_TRACK")) # https://consoledonottrack.com/ +) + +HF_TOKEN_PATH = os.path.expandvars( + os.path.expanduser( + os.getenv( + "HF_TOKEN_PATH", + os.path.join(HF_HOME, "token"), + ) + ) +) +HF_STORED_TOKENS_PATH = os.path.join(os.path.dirname(HF_TOKEN_PATH), "stored_tokens") + +if _staging_mode: + # In staging mode, we use a different cache to ensure we don't mix up production and staging data or tokens + # In practice in `huggingface_hub` tests, we monkeypatch these values with temporary directories. The following + # lines are only used in third-party libraries tests (e.g. `transformers`, `diffusers`, etc.). + _staging_home = os.path.join(os.path.expanduser("~"), ".cache", "huggingface_staging") + HUGGINGFACE_HUB_CACHE = os.path.join(_staging_home, "hub") + HF_TOKEN_PATH = os.path.join(_staging_home, "token") + +# Here, `True` will disable progress bars globally without possibility of enabling it +# programmatically. `False` will enable them without possibility of disabling them. +# If environment variable is not set (None), then the user is free to enable/disable +# them programmatically. +# TL;DR: env variable has priority over code +__HF_HUB_DISABLE_PROGRESS_BARS = os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS") +HF_HUB_DISABLE_PROGRESS_BARS: Optional[bool] = ( + _is_true(__HF_HUB_DISABLE_PROGRESS_BARS) if __HF_HUB_DISABLE_PROGRESS_BARS is not None else None +) + +# Disable warning on machines that do not support symlinks (e.g. Windows non-developer) +HF_HUB_DISABLE_SYMLINKS_WARNING: bool = _is_true(os.environ.get("HF_HUB_DISABLE_SYMLINKS_WARNING")) + +# Disable warning when using experimental features +HF_HUB_DISABLE_EXPERIMENTAL_WARNING: bool = _is_true(os.environ.get("HF_HUB_DISABLE_EXPERIMENTAL_WARNING")) + +# Disable sending the cached token by default is all HTTP requests to the Hub +HF_HUB_DISABLE_IMPLICIT_TOKEN: bool = _is_true(os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")) + +# Enable fast-download using external dependency "hf_transfer" +# See: +# - https://pypi.org/project/hf-transfer/ +# - https://github.com/huggingface/hf_transfer (private) +HF_HUB_ENABLE_HF_TRANSFER: bool = _is_true(os.environ.get("HF_HUB_ENABLE_HF_TRANSFER")) + + +# UNUSED +# We don't use symlinks in local dir anymore. +HF_HUB_LOCAL_DIR_AUTO_SYMLINK_THRESHOLD: int = ( + _as_int(os.environ.get("HF_HUB_LOCAL_DIR_AUTO_SYMLINK_THRESHOLD")) or 5 * 1024 * 1024 +) + +# Used to override the etag timeout on a system level +HF_HUB_ETAG_TIMEOUT: int = _as_int(os.environ.get("HF_HUB_ETAG_TIMEOUT")) or DEFAULT_ETAG_TIMEOUT + +# Used to override the get request timeout on a system level +HF_HUB_DOWNLOAD_TIMEOUT: int = _as_int(os.environ.get("HF_HUB_DOWNLOAD_TIMEOUT")) or DEFAULT_DOWNLOAD_TIMEOUT + +# Allows to add information about the requester in the user-agent (eg. partner name) +HF_HUB_USER_AGENT_ORIGIN: Optional[str] = os.environ.get("HF_HUB_USER_AGENT_ORIGIN") + +# List frameworks that are handled by the InferenceAPI service. Useful to scan endpoints and check which models are +# deployed and running. Since 95% of the models are using the top 4 frameworks listed below, we scan only those by +# default. We still keep the full list of supported frameworks in case we want to scan all of them. +MAIN_INFERENCE_API_FRAMEWORKS = [ + "diffusers", + "sentence-transformers", + "text-generation-inference", + "transformers", +] + +ALL_INFERENCE_API_FRAMEWORKS = MAIN_INFERENCE_API_FRAMEWORKS + [ + "adapter-transformers", + "allennlp", + "asteroid", + "bertopic", + "doctr", + "espnet", + "fairseq", + "fastai", + "fasttext", + "flair", + "k2", + "keras", + "mindspore", + "nemo", + "open_clip", + "paddlenlp", + "peft", + "pyannote-audio", + "sklearn", + "spacy", + "span-marker", + "speechbrain", + "stanza", + "timm", +] + +# Xet constants + + +HUGGINGFACE_HEADER_X_XET_ENDPOINT = "X-Xet-Cas-Url" +HUGGINGFACE_HEADER_X_XET_ACCESS_TOKEN = "X-Xet-Access-Token" +HUGGINGFACE_HEADER_X_XET_EXPIRATION = "X-Xet-Token-Expiration" +HUGGINGFACE_HEADER_X_XET_HASH = "X-Xet-Hash" +HUGGINGFACE_HEADER_X_XET_REFRESH_ROUTE = "X-Xet-Refresh-Route" +HUGGINGFACE_HEADER_LINK_XET_AUTH_KEY = "xet-auth" + +default_xet_cache_path = os.path.join(HF_HOME, "xet") +HF_XET_CACHE = os.getenv("HF_XET_CACHE", default_xet_cache_path) diff --git a/lib/python3.12/site-packages/huggingface_hub/dataclasses.py b/lib/python3.12/site-packages/huggingface_hub/dataclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..c5f8c7a3ea50aa774143c3ccdf7a1da6795013e0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/dataclasses.py @@ -0,0 +1,481 @@ +import inspect +from dataclasses import _MISSING_TYPE, MISSING, Field, field, fields +from functools import wraps +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + Type, + TypeVar, + Union, + get_args, + get_origin, + overload, +) + +from .errors import ( + StrictDataclassClassValidationError, + StrictDataclassDefinitionError, + StrictDataclassFieldValidationError, +) + + +Validator_T = Callable[[Any], None] +T = TypeVar("T") + + +# The overload decorator helps type checkers understand the different return types +@overload +def strict(cls: Type[T]) -> Type[T]: ... + + +@overload +def strict(*, accept_kwargs: bool = False) -> Callable[[Type[T]], Type[T]]: ... + + +def strict( + cls: Optional[Type[T]] = None, *, accept_kwargs: bool = False +) -> Union[Type[T], Callable[[Type[T]], Type[T]]]: + """ + Decorator to add strict validation to a dataclass. + + This decorator must be used on top of `@dataclass` to ensure IDEs and static typing tools + recognize the class as a dataclass. + + Can be used with or without arguments: + - `@strict` + - `@strict(accept_kwargs=True)` + + Args: + cls: + The class to convert to a strict dataclass. + accept_kwargs (`bool`, *optional*): + If True, allows arbitrary keyword arguments in `__init__`. Defaults to False. + + Returns: + The enhanced dataclass with strict validation on field assignment. + + Example: + ```py + >>> from dataclasses import dataclass + >>> from huggingface_hub.dataclasses import as_validated_field, strict, validated_field + + >>> @as_validated_field + >>> def positive_int(value: int): + ... if not value >= 0: + ... raise ValueError(f"Value must be positive, got {value}") + + >>> @strict(accept_kwargs=True) + ... @dataclass + ... class User: + ... name: str + ... age: int = positive_int(default=10) + + # Initialize + >>> User(name="John") + User(name='John', age=10) + + # Extra kwargs are accepted + >>> User(name="John", age=30, lastname="Doe") + User(name='John', age=30, *lastname='Doe') + + # Invalid type => raises + >>> User(name="John", age="30") + huggingface_hub.errors.StrictDataclassFieldValidationError: Validation error for field 'age': + TypeError: Field 'age' expected int, got str (value: '30') + + # Invalid value => raises + >>> User(name="John", age=-1) + huggingface_hub.errors.StrictDataclassFieldValidationError: Validation error for field 'age': + ValueError: Value must be positive, got -1 + ``` + """ + + def wrap(cls: Type[T]) -> Type[T]: + if not hasattr(cls, "__dataclass_fields__"): + raise StrictDataclassDefinitionError( + f"Class '{cls.__name__}' must be a dataclass before applying @strict." + ) + + # List and store validators + field_validators: Dict[str, List[Validator_T]] = {} + for f in fields(cls): # type: ignore [arg-type] + validators = [] + validators.append(_create_type_validator(f)) + custom_validator = f.metadata.get("validator") + if custom_validator is not None: + if not isinstance(custom_validator, list): + custom_validator = [custom_validator] + for validator in custom_validator: + if not _is_validator(validator): + raise StrictDataclassDefinitionError( + f"Invalid validator for field '{f.name}': {validator}. Must be a callable taking a single argument." + ) + validators.extend(custom_validator) + field_validators[f.name] = validators + cls.__validators__ = field_validators # type: ignore + + # Override __setattr__ to validate fields on assignment + original_setattr = cls.__setattr__ + + def __strict_setattr__(self: Any, name: str, value: Any) -> None: + """Custom __setattr__ method for strict dataclasses.""" + # Run all validators + for validator in self.__validators__.get(name, []): + try: + validator(value) + except (ValueError, TypeError) as e: + raise StrictDataclassFieldValidationError(field=name, cause=e) from e + + # If validation passed, set the attribute + original_setattr(self, name, value) + + cls.__setattr__ = __strict_setattr__ # type: ignore[method-assign] + + if accept_kwargs: + # (optional) Override __init__ to accept arbitrary keyword arguments + original_init = cls.__init__ + + @wraps(original_init) + def __init__(self, **kwargs: Any) -> None: + # Extract only the fields that are part of the dataclass + dataclass_fields = {f.name for f in fields(cls)} # type: ignore [arg-type] + standard_kwargs = {k: v for k, v in kwargs.items() if k in dataclass_fields} + + # Call the original __init__ with standard fields + original_init(self, **standard_kwargs) + + # Add any additional kwargs as attributes + for name, value in kwargs.items(): + if name not in dataclass_fields: + self.__setattr__(name, value) + + cls.__init__ = __init__ # type: ignore[method-assign] + + # (optional) Override __repr__ to include additional kwargs + original_repr = cls.__repr__ + + @wraps(original_repr) + def __repr__(self) -> str: + # Call the original __repr__ to get the standard fields + standard_repr = original_repr(self) + + # Get additional kwargs + additional_kwargs = [ + # add a '*' in front of additional kwargs to let the user know they are not part of the dataclass + f"*{k}={v!r}" + for k, v in self.__dict__.items() + if k not in cls.__dataclass_fields__ # type: ignore [attr-defined] + ] + additional_repr = ", ".join(additional_kwargs) + + # Combine both representations + return f"{standard_repr[:-1]}, {additional_repr})" if additional_kwargs else standard_repr + + cls.__repr__ = __repr__ # type: ignore [method-assign] + + # List all public methods starting with `validate_` => class validators. + class_validators = [] + + for name in dir(cls): + if not name.startswith("validate_"): + continue + method = getattr(cls, name) + if not callable(method): + continue + if len(inspect.signature(method).parameters) != 1: + raise StrictDataclassDefinitionError( + f"Class '{cls.__name__}' has a class validator '{name}' that takes more than one argument." + " Class validators must take only 'self' as an argument. Methods starting with 'validate_'" + " are considered to be class validators." + ) + class_validators.append(method) + + cls.__class_validators__ = class_validators # type: ignore [attr-defined] + + # Add `validate` method to the class, but first check if it already exists + def validate(self: T) -> None: + """Run class validators on the instance.""" + for validator in cls.__class_validators__: # type: ignore [attr-defined] + try: + validator(self) + except (ValueError, TypeError) as e: + raise StrictDataclassClassValidationError(validator=validator.__name__, cause=e) from e + + # Hack to be able to raise if `.validate()` already exists except if it was created by this decorator on a parent class + # (in which case we just override it) + validate.__is_defined_by_strict_decorator__ = True # type: ignore [attr-defined] + + if hasattr(cls, "validate"): + if not getattr(cls.validate, "__is_defined_by_strict_decorator__", False): # type: ignore [attr-defined] + raise StrictDataclassDefinitionError( + f"Class '{cls.__name__}' already implements a method called 'validate'." + " This method name is reserved when using the @strict decorator on a dataclass." + " If you want to keep your own method, please rename it." + ) + + cls.validate = validate # type: ignore + + # Run class validators after initialization + initial_init = cls.__init__ + + @wraps(initial_init) + def init_with_validate(self, *args, **kwargs) -> None: + """Run class validators after initialization.""" + initial_init(self, *args, **kwargs) # type: ignore [call-arg] + cls.validate(self) # type: ignore [attr-defined] + + setattr(cls, "__init__", init_with_validate) + + return cls + + # Return wrapped class or the decorator itself + return wrap(cls) if cls is not None else wrap + + +def validated_field( + validator: Union[List[Validator_T], Validator_T], + default: Union[Any, _MISSING_TYPE] = MISSING, + default_factory: Union[Callable[[], Any], _MISSING_TYPE] = MISSING, + init: bool = True, + repr: bool = True, + hash: Optional[bool] = None, + compare: bool = True, + metadata: Optional[Dict] = None, + **kwargs: Any, +) -> Any: + """ + Create a dataclass field with a custom validator. + + Useful to apply several checks to a field. If only applying one rule, check out the [`as_validated_field`] decorator. + + Args: + validator (`Callable` or `List[Callable]`): + A method that takes a value as input and raises ValueError/TypeError if the value is invalid. + Can be a list of validators to apply multiple checks. + **kwargs: + Additional arguments to pass to `dataclasses.field()`. + + Returns: + A field with the validator attached in metadata + """ + if not isinstance(validator, list): + validator = [validator] + if metadata is None: + metadata = {} + metadata["validator"] = validator + return field( # type: ignore + default=default, + default_factory=default_factory, + init=init, + repr=repr, + hash=hash, + compare=compare, + metadata=metadata, + **kwargs, + ) + + +def as_validated_field(validator: Validator_T): + """ + Decorates a validator function as a [`validated_field`] (i.e. a dataclass field with a custom validator). + + Args: + validator (`Callable`): + A method that takes a value as input and raises ValueError/TypeError if the value is invalid. + """ + + def _inner( + default: Union[Any, _MISSING_TYPE] = MISSING, + default_factory: Union[Callable[[], Any], _MISSING_TYPE] = MISSING, + init: bool = True, + repr: bool = True, + hash: Optional[bool] = None, + compare: bool = True, + metadata: Optional[Dict] = None, + **kwargs: Any, + ): + return validated_field( + validator, + default=default, + default_factory=default_factory, + init=init, + repr=repr, + hash=hash, + compare=compare, + metadata=metadata, + **kwargs, + ) + + return _inner + + +def type_validator(name: str, value: Any, expected_type: Any) -> None: + """Validate that 'value' matches 'expected_type'.""" + origin = get_origin(expected_type) + args = get_args(expected_type) + + if expected_type is Any: + return + elif validator := _BASIC_TYPE_VALIDATORS.get(origin): + validator(name, value, args) + elif isinstance(expected_type, type): # simple types + _validate_simple_type(name, value, expected_type) + else: + raise TypeError(f"Unsupported type for field '{name}': {expected_type}") + + +def _validate_union(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate that value matches one of the types in a Union.""" + errors = [] + for t in args: + try: + type_validator(name, value, t) + return # Valid if any type matches + except TypeError as e: + errors.append(str(e)) + + raise TypeError( + f"Field '{name}' with value {repr(value)} doesn't match any type in {args}. Errors: {'; '.join(errors)}" + ) + + +def _validate_literal(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate Literal type.""" + if value not in args: + raise TypeError(f"Field '{name}' expected one of {args}, got {value}") + + +def _validate_list(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate List[T] type.""" + if not isinstance(value, list): + raise TypeError(f"Field '{name}' expected a list, got {type(value).__name__}") + + # Validate each item in the list + item_type = args[0] + for i, item in enumerate(value): + try: + type_validator(f"{name}[{i}]", item, item_type) + except TypeError as e: + raise TypeError(f"Invalid item at index {i} in list '{name}'") from e + + +def _validate_dict(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate Dict[K, V] type.""" + if not isinstance(value, dict): + raise TypeError(f"Field '{name}' expected a dict, got {type(value).__name__}") + + # Validate keys and values + key_type, value_type = args + for k, v in value.items(): + try: + type_validator(f"{name}.key", k, key_type) + type_validator(f"{name}[{k!r}]", v, value_type) + except TypeError as e: + raise TypeError(f"Invalid key or value in dict '{name}'") from e + + +def _validate_tuple(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate Tuple type.""" + if not isinstance(value, tuple): + raise TypeError(f"Field '{name}' expected a tuple, got {type(value).__name__}") + + # Handle variable-length tuples: Tuple[T, ...] + if len(args) == 2 and args[1] is Ellipsis: + for i, item in enumerate(value): + try: + type_validator(f"{name}[{i}]", item, args[0]) + except TypeError as e: + raise TypeError(f"Invalid item at index {i} in tuple '{name}'") from e + # Handle fixed-length tuples: Tuple[T1, T2, ...] + elif len(args) != len(value): + raise TypeError(f"Field '{name}' expected a tuple of length {len(args)}, got {len(value)}") + else: + for i, (item, expected) in enumerate(zip(value, args)): + try: + type_validator(f"{name}[{i}]", item, expected) + except TypeError as e: + raise TypeError(f"Invalid item at index {i} in tuple '{name}'") from e + + +def _validate_set(name: str, value: Any, args: Tuple[Any, ...]) -> None: + """Validate Set[T] type.""" + if not isinstance(value, set): + raise TypeError(f"Field '{name}' expected a set, got {type(value).__name__}") + + # Validate each item in the set + item_type = args[0] + for i, item in enumerate(value): + try: + type_validator(f"{name} item", item, item_type) + except TypeError as e: + raise TypeError(f"Invalid item in set '{name}'") from e + + +def _validate_simple_type(name: str, value: Any, expected_type: type) -> None: + """Validate simple type (int, str, etc.).""" + if not isinstance(value, expected_type): + raise TypeError( + f"Field '{name}' expected {expected_type.__name__}, got {type(value).__name__} (value: {repr(value)})" + ) + + +def _create_type_validator(field: Field) -> Validator_T: + """Create a type validator function for a field.""" + # Hacky: we cannot use a lambda here because of reference issues + + def validator(value: Any) -> None: + type_validator(field.name, value, field.type) + + return validator + + +def _is_validator(validator: Any) -> bool: + """Check if a function is a validator. + + A validator is a Callable that can be called with a single positional argument. + The validator can have more arguments with default values. + + Basically, returns True if `validator(value)` is possible. + """ + if not callable(validator): + return False + + signature = inspect.signature(validator) + parameters = list(signature.parameters.values()) + if len(parameters) == 0: + return False + if parameters[0].kind not in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + ): + return False + for parameter in parameters[1:]: + if parameter.default == inspect.Parameter.empty: + return False + return True + + +_BASIC_TYPE_VALIDATORS = { + Union: _validate_union, + Literal: _validate_literal, + list: _validate_list, + dict: _validate_dict, + tuple: _validate_tuple, + set: _validate_set, +} + + +__all__ = [ + "strict", + "validated_field", + "Validator_T", + "StrictDataclassClassValidationError", + "StrictDataclassDefinitionError", + "StrictDataclassFieldValidationError", +] diff --git a/lib/python3.12/site-packages/huggingface_hub/errors.py b/lib/python3.12/site-packages/huggingface_hub/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..a0f7ed80e35a7cbe1dcc0f21dfa0354e467676f3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/errors.py @@ -0,0 +1,377 @@ +"""Contains all custom errors.""" + +from pathlib import Path +from typing import Optional, Union + +from requests import HTTPError, Response + + +# CACHE ERRORS + + +class CacheNotFound(Exception): + """Exception thrown when the Huggingface cache is not found.""" + + cache_dir: Union[str, Path] + + def __init__(self, msg: str, cache_dir: Union[str, Path], *args, **kwargs): + super().__init__(msg, *args, **kwargs) + self.cache_dir = cache_dir + + +class CorruptedCacheException(Exception): + """Exception for any unexpected structure in the Huggingface cache-system.""" + + +# HEADERS ERRORS + + +class LocalTokenNotFoundError(EnvironmentError): + """Raised if local token is required but not found.""" + + +# HTTP ERRORS + + +class OfflineModeIsEnabled(ConnectionError): + """Raised when a request is made but `HF_HUB_OFFLINE=1` is set as environment variable.""" + + +class HfHubHTTPError(HTTPError): + """ + HTTPError to inherit from for any custom HTTP Error raised in HF Hub. + + Any HTTPError is converted at least into a `HfHubHTTPError`. If some information is + sent back by the server, it will be added to the error message. + + Added details: + - Request id from "X-Request-Id" header if exists. If not, fallback to "X-Amzn-Trace-Id" header if exists. + - Server error message from the header "X-Error-Message". + - Server error message if we can found one in the response body. + + Example: + ```py + import requests + from huggingface_hub.utils import get_session, hf_raise_for_status, HfHubHTTPError + + response = get_session().post(...) + try: + hf_raise_for_status(response) + except HfHubHTTPError as e: + print(str(e)) # formatted message + e.request_id, e.server_message # details returned by server + + # Complete the error message with additional information once it's raised + e.append_to_message("\n`create_commit` expects the repository to exist.") + raise + ``` + """ + + def __init__(self, message: str, response: Optional[Response] = None, *, server_message: Optional[str] = None): + self.request_id = ( + response.headers.get("x-request-id") or response.headers.get("X-Amzn-Trace-Id") + if response is not None + else None + ) + self.server_message = server_message + + super().__init__( + message, + response=response, # type: ignore [arg-type] + request=response.request if response is not None else None, # type: ignore [arg-type] + ) + + def append_to_message(self, additional_message: str) -> None: + """Append additional information to the `HfHubHTTPError` initial message.""" + self.args = (self.args[0] + additional_message,) + self.args[1:] + + +# INFERENCE CLIENT ERRORS + + +class InferenceTimeoutError(HTTPError, TimeoutError): + """Error raised when a model is unavailable or the request times out.""" + + +# INFERENCE ENDPOINT ERRORS + + +class InferenceEndpointError(Exception): + """Generic exception when dealing with Inference Endpoints.""" + + +class InferenceEndpointTimeoutError(InferenceEndpointError, TimeoutError): + """Exception for timeouts while waiting for Inference Endpoint.""" + + +# SAFETENSORS ERRORS + + +class SafetensorsParsingError(Exception): + """Raised when failing to parse a safetensors file metadata. + + This can be the case if the file is not a safetensors file or does not respect the specification. + """ + + +class NotASafetensorsRepoError(Exception): + """Raised when a repo is not a Safetensors repo i.e. doesn't have either a `model.safetensors` or a + `model.safetensors.index.json` file. + """ + + +# TEXT GENERATION ERRORS + + +class TextGenerationError(HTTPError): + """Generic error raised if text-generation went wrong.""" + + +# Text Generation Inference Errors +class ValidationError(TextGenerationError): + """Server-side validation error.""" + + +class GenerationError(TextGenerationError): + pass + + +class OverloadedError(TextGenerationError): + pass + + +class IncompleteGenerationError(TextGenerationError): + pass + + +class UnknownError(TextGenerationError): + pass + + +# VALIDATION ERRORS + + +class HFValidationError(ValueError): + """Generic exception thrown by `huggingface_hub` validators. + + Inherits from [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError). + """ + + +# FILE METADATA ERRORS + + +class FileMetadataError(OSError): + """Error triggered when the metadata of a file on the Hub cannot be retrieved (missing ETag or commit_hash). + + Inherits from `OSError` for backward compatibility. + """ + + +# REPOSITORY ERRORS + + +class RepositoryNotFoundError(HfHubHTTPError): + """ + Raised when trying to access a hf.co URL with an invalid repository name, or + with a private repo name the user does not have access to. + + Example: + + ```py + >>> from huggingface_hub import model_info + >>> model_info("") + (...) + huggingface_hub.utils._errors.RepositoryNotFoundError: 401 Client Error. (Request ID: PvMw_VjBMjVdMz53WKIzP) + + Repository Not Found for url: https://huggingface.co/api/models/%3Cnon_existent_repository%3E. + Please make sure you specified the correct `repo_id` and `repo_type`. + If the repo is private, make sure you are authenticated. + Invalid username or password. + ``` + """ + + +class GatedRepoError(RepositoryNotFoundError): + """ + Raised when trying to access a gated repository for which the user is not on the + authorized list. + + Note: derives from `RepositoryNotFoundError` to ensure backward compatibility. + + Example: + + ```py + >>> from huggingface_hub import model_info + >>> model_info("") + (...) + huggingface_hub.utils._errors.GatedRepoError: 403 Client Error. (Request ID: ViT1Bf7O_026LGSQuVqfa) + + Cannot access gated repo for url https://huggingface.co/api/models/ardent-figment/gated-model. + Access to model ardent-figment/gated-model is restricted and you are not in the authorized list. + Visit https://huggingface.co/ardent-figment/gated-model to ask for access. + ``` + """ + + +class DisabledRepoError(HfHubHTTPError): + """ + Raised when trying to access a repository that has been disabled by its author. + + Example: + + ```py + >>> from huggingface_hub import dataset_info + >>> dataset_info("laion/laion-art") + (...) + huggingface_hub.utils._errors.DisabledRepoError: 403 Client Error. (Request ID: Root=1-659fc3fa-3031673e0f92c71a2260dbe2;bc6f4dfb-b30a-4862-af0a-5cfe827610d8) + + Cannot access repository for url https://huggingface.co/api/datasets/laion/laion-art. + Access to this resource is disabled. + ``` + """ + + +# REVISION ERROR + + +class RevisionNotFoundError(HfHubHTTPError): + """ + Raised when trying to access a hf.co URL with a valid repository but an invalid + revision. + + Example: + + ```py + >>> from huggingface_hub import hf_hub_download + >>> hf_hub_download('bert-base-cased', 'config.json', revision='') + (...) + huggingface_hub.utils._errors.RevisionNotFoundError: 404 Client Error. (Request ID: Mwhe_c3Kt650GcdKEFomX) + + Revision Not Found for url: https://huggingface.co/bert-base-cased/resolve/%3Cnon-existent-revision%3E/config.json. + ``` + """ + + +# ENTRY ERRORS +class EntryNotFoundError(HfHubHTTPError): + """ + Raised when trying to access a hf.co URL with a valid repository and revision + but an invalid filename. + + Example: + + ```py + >>> from huggingface_hub import hf_hub_download + >>> hf_hub_download('bert-base-cased', '') + (...) + huggingface_hub.utils._errors.EntryNotFoundError: 404 Client Error. (Request ID: 53pNl6M0MxsnG5Sw8JA6x) + + Entry Not Found for url: https://huggingface.co/bert-base-cased/resolve/main/%3Cnon-existent-file%3E. + ``` + """ + + +class LocalEntryNotFoundError(EntryNotFoundError, FileNotFoundError, ValueError): + """ + Raised when trying to access a file or snapshot that is not on the disk when network is + disabled or unavailable (connection issue). The entry may exist on the Hub. + + Note: `ValueError` type is to ensure backward compatibility. + Note: `LocalEntryNotFoundError` derives from `HTTPError` because of `EntryNotFoundError` + even when it is not a network issue. + + Example: + + ```py + >>> from huggingface_hub import hf_hub_download + >>> hf_hub_download('bert-base-cased', '', local_files_only=True) + (...) + huggingface_hub.utils._errors.LocalEntryNotFoundError: Cannot find the requested files in the disk cache and outgoing traffic has been disabled. To enable hf.co look-ups and downloads online, set 'local_files_only' to False. + ``` + """ + + def __init__(self, message: str): + super().__init__(message, response=None) + + +# REQUEST ERROR +class BadRequestError(HfHubHTTPError, ValueError): + """ + Raised by `hf_raise_for_status` when the server returns a HTTP 400 error. + + Example: + + ```py + >>> resp = requests.post("hf.co/api/check", ...) + >>> hf_raise_for_status(resp, endpoint_name="check") + huggingface_hub.utils._errors.BadRequestError: Bad request for check endpoint: {details} (Request ID: XXX) + ``` + """ + + +# DDUF file format ERROR + + +class DDUFError(Exception): + """Base exception for errors related to the DDUF format.""" + + +class DDUFCorruptedFileError(DDUFError): + """Exception thrown when the DDUF file is corrupted.""" + + +class DDUFExportError(DDUFError): + """Base exception for errors during DDUF export.""" + + +class DDUFInvalidEntryNameError(DDUFExportError): + """Exception thrown when the entry name is invalid.""" + + +# STRICT DATACLASSES ERRORS + + +class StrictDataclassError(Exception): + """Base exception for strict dataclasses.""" + + +class StrictDataclassDefinitionError(StrictDataclassError): + """Exception thrown when a strict dataclass is defined incorrectly.""" + + +class StrictDataclassFieldValidationError(StrictDataclassError): + """Exception thrown when a strict dataclass fails validation for a given field.""" + + def __init__(self, field: str, cause: Exception): + error_message = f"Validation error for field '{field}':" + error_message += f"\n {cause.__class__.__name__}: {cause}" + super().__init__(error_message) + + +class StrictDataclassClassValidationError(StrictDataclassError): + """Exception thrown when a strict dataclass fails validation on a class validator.""" + + def __init__(self, validator: str, cause: Exception): + error_message = f"Class validation error for validator '{validator}':" + error_message += f"\n {cause.__class__.__name__}: {cause}" + super().__init__(error_message) + + +# XET ERRORS + + +class XetError(Exception): + """Base exception for errors related to Xet Storage.""" + + +class XetAuthorizationError(XetError): + """Exception thrown when the user does not have the right authorization to use Xet Storage.""" + + +class XetRefreshTokenError(XetError): + """Exception thrown when the refresh token is invalid.""" + + +class XetDownloadError(Exception): + """Exception thrown when the download from Xet Storage fails.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/fastai_utils.py b/lib/python3.12/site-packages/huggingface_hub/fastai_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e75eba2a8baee7bdeb8d36a1c06bd950cf857c44 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/fastai_utils.py @@ -0,0 +1,425 @@ +import json +import os +from pathlib import Path +from pickle import DEFAULT_PROTOCOL, PicklingError +from typing import Any, Dict, List, Optional, Union + +from packaging import version + +from huggingface_hub import constants, snapshot_download +from huggingface_hub.hf_api import HfApi +from huggingface_hub.utils import ( + SoftTemporaryDirectory, + get_fastai_version, + get_fastcore_version, + get_python_version, +) + +from .utils import logging, validate_hf_hub_args +from .utils._runtime import _PY_VERSION # noqa: F401 # for backward compatibility... + + +logger = logging.get_logger(__name__) + + +def _check_fastai_fastcore_versions( + fastai_min_version: str = "2.4", + fastcore_min_version: str = "1.3.27", +): + """ + Checks that the installed fastai and fastcore versions are compatible for pickle serialization. + + Args: + fastai_min_version (`str`, *optional*): + The minimum fastai version supported. + fastcore_min_version (`str`, *optional*): + The minimum fastcore version supported. + + + Raises the following error: + + - [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + if the fastai or fastcore libraries are not available or are of an invalid version. + + + """ + + if (get_fastcore_version() or get_fastai_version()) == "N/A": + raise ImportError( + f"fastai>={fastai_min_version} and fastcore>={fastcore_min_version} are" + f" required. Currently using fastai=={get_fastai_version()} and" + f" fastcore=={get_fastcore_version()}." + ) + + current_fastai_version = version.Version(get_fastai_version()) + current_fastcore_version = version.Version(get_fastcore_version()) + + if current_fastai_version < version.Version(fastai_min_version): + raise ImportError( + "`push_to_hub_fastai` and `from_pretrained_fastai` require a" + f" fastai>={fastai_min_version} version, but you are using fastai version" + f" {get_fastai_version()} which is incompatible. Upgrade with `pip install" + " fastai==2.5.6`." + ) + + if current_fastcore_version < version.Version(fastcore_min_version): + raise ImportError( + "`push_to_hub_fastai` and `from_pretrained_fastai` require a" + f" fastcore>={fastcore_min_version} version, but you are using fastcore" + f" version {get_fastcore_version()} which is incompatible. Upgrade with" + " `pip install fastcore==1.3.27`." + ) + + +def _check_fastai_fastcore_pyproject_versions( + storage_folder: str, + fastai_min_version: str = "2.4", + fastcore_min_version: str = "1.3.27", +): + """ + Checks that the `pyproject.toml` file in the directory `storage_folder` has fastai and fastcore versions + that are compatible with `from_pretrained_fastai` and `push_to_hub_fastai`. If `pyproject.toml` does not exist + or does not contain versions for fastai and fastcore, then it logs a warning. + + Args: + storage_folder (`str`): + Folder to look for the `pyproject.toml` file. + fastai_min_version (`str`, *optional*): + The minimum fastai version supported. + fastcore_min_version (`str`, *optional*): + The minimum fastcore version supported. + + + Raises the following errors: + + - [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + if the `toml` module is not installed. + - [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + if the `pyproject.toml` indicates a lower than minimum supported version of fastai or fastcore. + + + """ + + try: + import toml + except ModuleNotFoundError: + raise ImportError( + "`push_to_hub_fastai` and `from_pretrained_fastai` require the toml module." + " Install it with `pip install toml`." + ) + + # Checks that a `pyproject.toml`, with `build-system` and `requires` sections, exists in the repository. If so, get a list of required packages. + if not os.path.isfile(f"{storage_folder}/pyproject.toml"): + logger.warning( + "There is no `pyproject.toml` in the repository that contains the fastai" + " `Learner`. The `pyproject.toml` would allow us to verify that your fastai" + " and fastcore versions are compatible with those of the model you want to" + " load." + ) + return + pyproject_toml = toml.load(f"{storage_folder}/pyproject.toml") + + if "build-system" not in pyproject_toml.keys(): + logger.warning( + "There is no `build-system` section in the pyproject.toml of the repository" + " that contains the fastai `Learner`. The `build-system` would allow us to" + " verify that your fastai and fastcore versions are compatible with those" + " of the model you want to load." + ) + return + build_system_toml = pyproject_toml["build-system"] + + if "requires" not in build_system_toml.keys(): + logger.warning( + "There is no `requires` section in the pyproject.toml of the repository" + " that contains the fastai `Learner`. The `requires` would allow us to" + " verify that your fastai and fastcore versions are compatible with those" + " of the model you want to load." + ) + return + package_versions = build_system_toml["requires"] + + # Extracts contains fastai and fastcore versions from `pyproject.toml` if available. + # If the package is specified but not the version (e.g. "fastai" instead of "fastai=2.4"), the default versions are the highest. + fastai_packages = [pck for pck in package_versions if pck.startswith("fastai")] + if len(fastai_packages) == 0: + logger.warning("The repository does not have a fastai version specified in the `pyproject.toml`.") + # fastai_version is an empty string if not specified + else: + fastai_version = str(fastai_packages[0]).partition("=")[2] + if fastai_version != "" and version.Version(fastai_version) < version.Version(fastai_min_version): + raise ImportError( + "`from_pretrained_fastai` requires" + f" fastai>={fastai_min_version} version but the model to load uses" + f" {fastai_version} which is incompatible." + ) + + fastcore_packages = [pck for pck in package_versions if pck.startswith("fastcore")] + if len(fastcore_packages) == 0: + logger.warning("The repository does not have a fastcore version specified in the `pyproject.toml`.") + # fastcore_version is an empty string if not specified + else: + fastcore_version = str(fastcore_packages[0]).partition("=")[2] + if fastcore_version != "" and version.Version(fastcore_version) < version.Version(fastcore_min_version): + raise ImportError( + "`from_pretrained_fastai` requires" + f" fastcore>={fastcore_min_version} version, but you are using fastcore" + f" version {fastcore_version} which is incompatible." + ) + + +README_TEMPLATE = """--- +tags: +- fastai +--- + +# Amazing! + +🥳 Congratulations on hosting your fastai model on the Hugging Face Hub! + +# Some next steps +1. Fill out this model card with more information (see the template below and the [documentation here](https://huggingface.co/docs/hub/model-repos))! + +2. Create a demo in Gradio or Streamlit using 🤗 Spaces ([documentation here](https://huggingface.co/docs/hub/spaces)). + +3. Join the fastai community on the [Fastai Discord](https://discord.com/invite/YKrxeNn)! + +Greetings fellow fastlearner 🤝! Don't forget to delete this content from your model card. + + +--- + + +# Model card + +## Model description +More information needed + +## Intended uses & limitations +More information needed + +## Training and evaluation data +More information needed +""" + +PYPROJECT_TEMPLATE = f"""[build-system] +requires = ["setuptools>=40.8.0", "wheel", "python={get_python_version()}", "fastai={get_fastai_version()}", "fastcore={get_fastcore_version()}"] +build-backend = "setuptools.build_meta:__legacy__" +""" + + +def _create_model_card(repo_dir: Path): + """ + Creates a model card for the repository. + + Args: + repo_dir (`Path`): + Directory where model card is created. + """ + readme_path = repo_dir / "README.md" + + if not readme_path.exists(): + with readme_path.open("w", encoding="utf-8") as f: + f.write(README_TEMPLATE) + + +def _create_model_pyproject(repo_dir: Path): + """ + Creates a `pyproject.toml` for the repository. + + Args: + repo_dir (`Path`): + Directory where `pyproject.toml` is created. + """ + pyproject_path = repo_dir / "pyproject.toml" + + if not pyproject_path.exists(): + with pyproject_path.open("w", encoding="utf-8") as f: + f.write(PYPROJECT_TEMPLATE) + + +def _save_pretrained_fastai( + learner, + save_directory: Union[str, Path], + config: Optional[Dict[str, Any]] = None, +): + """ + Saves a fastai learner to `save_directory` in pickle format using the default pickle protocol for the version of python used. + + Args: + learner (`Learner`): + The `fastai.Learner` you'd like to save. + save_directory (`str` or `Path`): + Specific directory in which you want to save the fastai learner. + config (`dict`, *optional*): + Configuration object. Will be uploaded as a .json file. Example: 'https://huggingface.co/espejelomar/fastai-pet-breeds-classification/blob/main/config.json'. + + + + Raises the following error: + + - [`RuntimeError`](https://docs.python.org/3/library/exceptions.html#RuntimeError) + if the config file provided is not a dictionary. + + + """ + _check_fastai_fastcore_versions() + + os.makedirs(save_directory, exist_ok=True) + + # if the user provides config then we update it with the fastai and fastcore versions in CONFIG_TEMPLATE. + if config is not None: + if not isinstance(config, dict): + raise RuntimeError(f"Provided config should be a dict. Got: '{type(config)}'") + path = os.path.join(save_directory, constants.CONFIG_NAME) + with open(path, "w") as f: + json.dump(config, f) + + _create_model_card(Path(save_directory)) + _create_model_pyproject(Path(save_directory)) + + # learner.export saves the model in `self.path`. + learner.path = Path(save_directory) + os.makedirs(save_directory, exist_ok=True) + try: + learner.export( + fname="model.pkl", + pickle_protocol=DEFAULT_PROTOCOL, + ) + except PicklingError: + raise PicklingError( + "You are using a lambda function, i.e., an anonymous function. `pickle`" + " cannot pickle function objects and requires that all functions have" + " names. One possible solution is to name the function." + ) + + +@validate_hf_hub_args +def from_pretrained_fastai( + repo_id: str, + revision: Optional[str] = None, +): + """ + Load pretrained fastai model from the Hub or from a local directory. + + Args: + repo_id (`str`): + The location where the pickled fastai.Learner is. It can be either of the two: + - Hosted on the Hugging Face Hub. E.g.: 'espejelomar/fatai-pet-breeds-classification' or 'distilgpt2'. + You can add a `revision` by appending `@` at the end of `repo_id`. E.g.: `dbmdz/bert-base-german-cased@main`. + Revision is the specific model version to use. Since we use a git-based system for storing models and other + artifacts on the Hugging Face Hub, it can be a branch name, a tag name, or a commit id. + - Hosted locally. `repo_id` would be a directory containing the pickle and a pyproject.toml + indicating the fastai and fastcore versions used to build the `fastai.Learner`. E.g.: `./my_model_directory/`. + revision (`str`, *optional*): + Revision at which the repo's files are downloaded. See documentation of `snapshot_download`. + + Returns: + The `fastai.Learner` model in the `repo_id` repo. + """ + _check_fastai_fastcore_versions() + + # Load the `repo_id` repo. + # `snapshot_download` returns the folder where the model was stored. + # `cache_dir` will be the default '/root/.cache/huggingface/hub' + if not os.path.isdir(repo_id): + storage_folder = snapshot_download( + repo_id=repo_id, + revision=revision, + library_name="fastai", + library_version=get_fastai_version(), + ) + else: + storage_folder = repo_id + + _check_fastai_fastcore_pyproject_versions(storage_folder) + + from fastai.learner import load_learner # type: ignore + + return load_learner(os.path.join(storage_folder, "model.pkl")) + + +@validate_hf_hub_args +def push_to_hub_fastai( + learner, + *, + repo_id: str, + commit_message: str = "Push FastAI model using huggingface_hub.", + private: Optional[bool] = None, + token: Optional[str] = None, + config: Optional[dict] = None, + branch: Optional[str] = None, + create_pr: Optional[bool] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + api_endpoint: Optional[str] = None, +): + """ + Upload learner checkpoint files to the Hub. + + Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use + `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more + details. + + Args: + learner (`Learner`): + The `fastai.Learner' you'd like to push to the Hub. + repo_id (`str`): + The repository id for your model in Hub in the format of "namespace/repo_name". The namespace can be your individual account or an organization to which you have write access (for example, 'stanfordnlp/stanza-de'). + commit_message (`str`, *optional*): + Message to commit while pushing. Will default to :obj:`"add model"`. + private (`bool`, *optional*): + Whether or not the repository created should be private. + If `None` (default), will default to been public except if the organization's default is private. + token (`str`, *optional*): + The Hugging Face account token to use as HTTP bearer authorization for remote files. If :obj:`None`, the token will be asked by a prompt. + config (`dict`, *optional*): + Configuration object to be saved alongside the model weights. + branch (`str`, *optional*): + The git branch on which to push the model. This defaults to + the default branch as specified in your repository, which + defaults to `"main"`. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request from `branch` with that commit. + Defaults to `False`. + api_endpoint (`str`, *optional*): + The API endpoint to use when pushing the model to the hub. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are pushed. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not pushed. + delete_patterns (`List[str]` or `str`, *optional*): + If provided, remote files matching any of the patterns will be deleted from the repo. + + Returns: + The url of the commit of your model in the given repository. + + + + Raises the following error: + + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if the user is not log on to the Hugging Face Hub. + + + """ + _check_fastai_fastcore_versions() + api = HfApi(endpoint=api_endpoint) + repo_id = api.create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True).repo_id + + # Push the files to the repo in a single commit + with SoftTemporaryDirectory() as tmp: + saved_path = Path(tmp) / repo_id + _save_pretrained_fastai(learner, saved_path, config=config) + return api.upload_folder( + repo_id=repo_id, + token=token, + folder_path=saved_path, + commit_message=commit_message, + revision=branch, + create_pr=create_pr, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + delete_patterns=delete_patterns, + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/file_download.py b/lib/python3.12/site-packages/huggingface_hub/file_download.py new file mode 100644 index 0000000000000000000000000000000000000000..2b3a4a1c4bac58bd62bab0109d64a200ff5d8722 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/file_download.py @@ -0,0 +1,1801 @@ +import copy +import errno +import inspect +import os +import re +import shutil +import stat +import time +import uuid +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Any, BinaryIO, Dict, Literal, NoReturn, Optional, Tuple, Union +from urllib.parse import quote, urlparse + +import requests + +from . import ( + __version__, # noqa: F401 # for backward compatibility + constants, +) +from ._local_folder import get_local_download_paths, read_download_metadata, write_download_metadata +from .constants import ( + HUGGINGFACE_CO_URL_TEMPLATE, # noqa: F401 # for backward compatibility + HUGGINGFACE_HUB_CACHE, # noqa: F401 # for backward compatibility +) +from .errors import ( + EntryNotFoundError, + FileMetadataError, + GatedRepoError, + HfHubHTTPError, + LocalEntryNotFoundError, + RepositoryNotFoundError, + RevisionNotFoundError, +) +from .utils import ( + OfflineModeIsEnabled, + SoftTemporaryDirectory, + WeakFileLock, + XetFileData, + build_hf_headers, + get_fastai_version, # noqa: F401 # for backward compatibility + get_fastcore_version, # noqa: F401 # for backward compatibility + get_graphviz_version, # noqa: F401 # for backward compatibility + get_jinja_version, # noqa: F401 # for backward compatibility + get_pydot_version, # noqa: F401 # for backward compatibility + get_tf_version, # noqa: F401 # for backward compatibility + get_torch_version, # noqa: F401 # for backward compatibility + hf_raise_for_status, + is_fastai_available, # noqa: F401 # for backward compatibility + is_fastcore_available, # noqa: F401 # for backward compatibility + is_graphviz_available, # noqa: F401 # for backward compatibility + is_jinja_available, # noqa: F401 # for backward compatibility + is_pydot_available, # noqa: F401 # for backward compatibility + is_tf_available, # noqa: F401 # for backward compatibility + is_torch_available, # noqa: F401 # for backward compatibility + logging, + parse_xet_file_data_from_response, + refresh_xet_connection_info, + reset_sessions, + tqdm, + validate_hf_hub_args, +) +from .utils._http import _adjust_range_header, http_backoff +from .utils._runtime import _PY_VERSION, is_xet_available # noqa: F401 # for backward compatibility +from .utils._typing import HTTP_METHOD_T +from .utils.sha import sha_fileobj +from .utils.tqdm import _get_progress_bar_context + + +logger = logging.get_logger(__name__) + +# Return value when trying to load a file from cache but the file does not exist in the distant repo. +_CACHED_NO_EXIST = object() +_CACHED_NO_EXIST_T = Any + +# Regex to get filename from a "Content-Disposition" header for CDN-served files +HEADER_FILENAME_PATTERN = re.compile(r'filename="(?P.*?)";') + +# Regex to check if the revision IS directly a commit_hash +REGEX_COMMIT_HASH = re.compile(r"^[0-9a-f]{40}$") + +# Regex to check if the file etag IS a valid sha256 +REGEX_SHA256 = re.compile(r"^[0-9a-f]{64}$") + +_are_symlinks_supported_in_dir: Dict[str, bool] = {} + + +def are_symlinks_supported(cache_dir: Union[str, Path, None] = None) -> bool: + """Return whether the symlinks are supported on the machine. + + Since symlinks support can change depending on the mounted disk, we need to check + on the precise cache folder. By default, the default HF cache directory is checked. + + Args: + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + + Returns: [bool] Whether symlinks are supported in the directory. + """ + # Defaults to HF cache + if cache_dir is None: + cache_dir = constants.HF_HUB_CACHE + cache_dir = str(Path(cache_dir).expanduser().resolve()) # make it unique + + # Check symlink compatibility only once (per cache directory) at first time use + if cache_dir not in _are_symlinks_supported_in_dir: + _are_symlinks_supported_in_dir[cache_dir] = True + + os.makedirs(cache_dir, exist_ok=True) + with SoftTemporaryDirectory(dir=cache_dir) as tmpdir: + src_path = Path(tmpdir) / "dummy_file_src" + src_path.touch() + dst_path = Path(tmpdir) / "dummy_file_dst" + + # Relative source path as in `_create_symlink`` + relative_src = os.path.relpath(src_path, start=os.path.dirname(dst_path)) + try: + os.symlink(relative_src, dst_path) + except OSError: + # Likely running on Windows + _are_symlinks_supported_in_dir[cache_dir] = False + + if not constants.HF_HUB_DISABLE_SYMLINKS_WARNING: + message = ( + "`huggingface_hub` cache-system uses symlinks by default to" + " efficiently store duplicated files but your machine does not" + f" support them in {cache_dir}. Caching files will still work" + " but in a degraded version that might require more space on" + " your disk. This warning can be disabled by setting the" + " `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For" + " more details, see" + " https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations." + ) + if os.name == "nt": + message += ( + "\nTo support symlinks on Windows, you either need to" + " activate Developer Mode or to run Python as an" + " administrator. In order to activate developer mode," + " see this article:" + " https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development" + ) + warnings.warn(message) + + return _are_symlinks_supported_in_dir[cache_dir] + + +@dataclass(frozen=True) +class HfFileMetadata: + """Data structure containing information about a file versioned on the Hub. + + Returned by [`get_hf_file_metadata`] based on a URL. + + Args: + commit_hash (`str`, *optional*): + The commit_hash related to the file. + etag (`str`, *optional*): + Etag of the file on the server. + location (`str`): + Location where to download the file. Can be a Hub url or not (CDN). + size (`size`): + Size of the file. In case of an LFS file, contains the size of the actual + LFS file, not the pointer. + xet_file_data (`XetFileData`, *optional*): + Xet information for the file. This is only set if the file is stored using Xet storage. + """ + + commit_hash: Optional[str] + etag: Optional[str] + location: str + size: Optional[int] + xet_file_data: Optional[XetFileData] + + +@validate_hf_hub_args +def hf_hub_url( + repo_id: str, + filename: str, + *, + subfolder: Optional[str] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + endpoint: Optional[str] = None, +) -> str: + """Construct the URL of a file from the given information. + + The resolved address can either be a huggingface.co-hosted url, or a link to + Cloudfront (a Content Delivery Network, or CDN) for large files which are + more than a few MBs. + + Args: + repo_id (`str`): + A namespace (user or an organization) name and a repo name separated + by a `/`. + filename (`str`): + The name of the file in the repo. + subfolder (`str`, *optional*): + An optional value corresponding to a folder inside the repo. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if downloading from a dataset or space, + `None` or `"model"` if downloading from a model. Default is `None`. + revision (`str`, *optional*): + An optional Git revision id which can be a branch name, a tag, or a + commit hash. + + Example: + + ```python + >>> from huggingface_hub import hf_hub_url + + >>> hf_hub_url( + ... repo_id="julien-c/EsperBERTo-small", filename="pytorch_model.bin" + ... ) + 'https://huggingface.co/julien-c/EsperBERTo-small/resolve/main/pytorch_model.bin' + ``` + + + + Notes: + + Cloudfront is replicated over the globe so downloads are way faster for + the end user (and it also lowers our bandwidth costs). + + Cloudfront aggressively caches files by default (default TTL is 24 + hours), however this is not an issue here because we implement a + git-based versioning system on huggingface.co, which means that we store + the files on S3/Cloudfront in a content-addressable way (i.e., the file + name is its hash). Using content-addressable filenames means cache can't + ever be stale. + + In terms of client-side caching from this library, we base our caching + on the objects' entity tag (`ETag`), which is an identifier of a + specific version of a resource [1]_. An object's ETag is: its git-sha1 + if stored in git, or its sha256 if stored in git-lfs. + + + + References: + + - [1] https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/ETag + """ + if subfolder == "": + subfolder = None + if subfolder is not None: + filename = f"{subfolder}/{filename}" + + if repo_type not in constants.REPO_TYPES: + raise ValueError("Invalid repo type") + + if repo_type in constants.REPO_TYPES_URL_PREFIXES: + repo_id = constants.REPO_TYPES_URL_PREFIXES[repo_type] + repo_id + + if revision is None: + revision = constants.DEFAULT_REVISION + url = HUGGINGFACE_CO_URL_TEMPLATE.format( + repo_id=repo_id, revision=quote(revision, safe=""), filename=quote(filename) + ) + # Update endpoint if provided + if endpoint is not None and url.startswith(constants.ENDPOINT): + url = endpoint + url[len(constants.ENDPOINT) :] + return url + + +def _request_wrapper( + method: HTTP_METHOD_T, url: str, *, follow_relative_redirects: bool = False, **params +) -> requests.Response: + """Wrapper around requests methods to follow relative redirects if `follow_relative_redirects=True` even when + `allow_redirection=False`. + + A backoff mechanism retries the HTTP call on 429, 503 and 504 errors. + + Args: + method (`str`): + HTTP method, such as 'GET' or 'HEAD'. + url (`str`): + The URL of the resource to fetch. + follow_relative_redirects (`bool`, *optional*, defaults to `False`) + If True, relative redirection (redirection to the same site) will be resolved even when `allow_redirection` + kwarg is set to False. Useful when we want to follow a redirection to a renamed repository without + following redirection to a CDN. + **params (`dict`, *optional*): + Params to pass to `requests.request`. + """ + # Recursively follow relative redirects + if follow_relative_redirects: + response = _request_wrapper( + method=method, + url=url, + follow_relative_redirects=False, + **params, + ) + + # If redirection, we redirect only relative paths. + # This is useful in case of a renamed repository. + if 300 <= response.status_code <= 399: + parsed_target = urlparse(response.headers["Location"]) + if parsed_target.netloc == "": + # This means it is a relative 'location' headers, as allowed by RFC 7231. + # (e.g. '/path/to/resource' instead of 'http://domain.tld/path/to/resource') + # We want to follow this relative redirect ! + # + # Highly inspired by `resolve_redirects` from requests library. + # See https://github.com/psf/requests/blob/main/requests/sessions.py#L159 + next_url = urlparse(url)._replace(path=parsed_target.path).geturl() + return _request_wrapper(method=method, url=next_url, follow_relative_redirects=True, **params) + return response + + # Perform request and return if status_code is not in the retry list. + response = http_backoff(method=method, url=url, **params, retry_on_exceptions=(), retry_on_status_codes=(429,)) + hf_raise_for_status(response) + return response + + +def _get_file_length_from_http_response(response: requests.Response) -> Optional[int]: + """ + Get the length of the file from the HTTP response headers. + + This function extracts the file size from the HTTP response headers, either from the + `Content-Range` or `Content-Length` header, if available (in that order). + The HTTP response object containing the headers. + `int` or `None`: The length of the file in bytes if the information is available, + otherwise `None`. + + Args: + response (`requests.Response`): + The HTTP response object. + + Returns: + `int` or `None`: The length of the file in bytes, or None if not available. + """ + + content_range = response.headers.get("Content-Range") + if content_range is not None: + return int(content_range.rsplit("/")[-1]) + + content_length = response.headers.get("Content-Length") + if content_length is not None: + return int(content_length) + + return None + + +def http_get( + url: str, + temp_file: BinaryIO, + *, + proxies: Optional[Dict] = None, + resume_size: int = 0, + headers: Optional[Dict[str, Any]] = None, + expected_size: Optional[int] = None, + displayed_filename: Optional[str] = None, + _nb_retries: int = 5, + _tqdm_bar: Optional[tqdm] = None, +) -> None: + """ + Download a remote file. Do not gobble up errors, and will return errors tailored to the Hugging Face Hub. + + If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely a + transient error (network outage?). We log a warning message and try to resume the download a few times before + giving up. The method gives up after 5 attempts if no new data has being received from the server. + + Args: + url (`str`): + The URL of the file to download. + temp_file (`BinaryIO`): + The file-like object where to save the file. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to `requests.request`. + resume_size (`int`, *optional*): + The number of bytes already downloaded. If set to 0 (default), the whole file is download. If set to a + positive number, the download will resume at the given position. + headers (`dict`, *optional*): + Dictionary of HTTP Headers to send with the request. + expected_size (`int`, *optional*): + The expected size of the file to download. If set, the download will raise an error if the size of the + received content is different from the expected one. + displayed_filename (`str`, *optional*): + The filename of the file that is being downloaded. Value is used only to display a nice progress bar. If + not set, the filename is guessed from the URL or the `Content-Disposition` header. + """ + if expected_size is not None and resume_size == expected_size: + # If the file is already fully downloaded, we don't need to download it again. + return + + has_custom_range_header = headers is not None and any(h.lower() == "range" for h in headers) + hf_transfer = None + if constants.HF_HUB_ENABLE_HF_TRANSFER: + if resume_size != 0: + warnings.warn("'hf_transfer' does not support `resume_size`: falling back to regular download method") + elif proxies is not None: + warnings.warn("'hf_transfer' does not support `proxies`: falling back to regular download method") + elif has_custom_range_header: + warnings.warn("'hf_transfer' ignores custom 'Range' headers; falling back to regular download method") + else: + try: + import hf_transfer # type: ignore[no-redef] + except ImportError: + raise ValueError( + "Fast download using 'hf_transfer' is enabled" + " (HF_HUB_ENABLE_HF_TRANSFER=1) but 'hf_transfer' package is not" + " available in your environment. Try `pip install hf_transfer`." + ) + + initial_headers = headers + headers = copy.deepcopy(headers) or {} + if resume_size > 0: + headers["Range"] = _adjust_range_header(headers.get("Range"), resume_size) + elif expected_size and expected_size > constants.MAX_HTTP_DOWNLOAD_SIZE: + # Any files over 50GB will not be available through basic http request. + # Setting the range header to 0-0 will force the server to return the file size in the Content-Range header. + # Since hf_transfer splits the download into chunks, the process will succeed afterwards. + if hf_transfer: + headers["Range"] = "bytes=0-0" + else: + raise ValueError( + "The file is too large to be downloaded using the regular download method. Use `hf_transfer` or `hf_xet` instead." + " Try `pip install hf_transfer` or `pip install hf_xet`." + ) + + r = _request_wrapper( + method="GET", url=url, stream=True, proxies=proxies, headers=headers, timeout=constants.HF_HUB_DOWNLOAD_TIMEOUT + ) + + hf_raise_for_status(r) + content_length = _get_file_length_from_http_response(r) + + # NOTE: 'total' is the total number of bytes to download, not the number of bytes in the file. + # If the file is compressed, the number of bytes in the saved file will be higher than 'total'. + total = resume_size + int(content_length) if content_length is not None else None + + if displayed_filename is None: + displayed_filename = url + content_disposition = r.headers.get("Content-Disposition") + if content_disposition is not None: + match = HEADER_FILENAME_PATTERN.search(content_disposition) + if match is not None: + # Means file is on CDN + displayed_filename = match.groupdict()["filename"] + + # Truncate filename if too long to display + if len(displayed_filename) > 40: + displayed_filename = f"(…){displayed_filename[-40:]}" + + consistency_error_message = ( + f"Consistency check failed: file should be of size {expected_size} but has size" + f" {{actual_size}} ({displayed_filename}).\nThis is usually due to network issues while downloading the file." + " Please retry with `force_download=True`." + ) + progress_cm = _get_progress_bar_context( + desc=displayed_filename, + log_level=logger.getEffectiveLevel(), + total=total, + initial=resume_size, + name="huggingface_hub.http_get", + _tqdm_bar=_tqdm_bar, + ) + + with progress_cm as progress: + if hf_transfer and total is not None and total > 5 * constants.DOWNLOAD_CHUNK_SIZE: + supports_callback = "callback" in inspect.signature(hf_transfer.download).parameters + if not supports_callback: + warnings.warn( + "You are using an outdated version of `hf_transfer`. " + "Consider upgrading to latest version to enable progress bars " + "using `pip install -U hf_transfer`." + ) + try: + hf_transfer.download( + url=url, + filename=temp_file.name, + max_files=constants.HF_TRANSFER_CONCURRENCY, + chunk_size=constants.DOWNLOAD_CHUNK_SIZE, + headers=initial_headers, + parallel_failures=3, + max_retries=5, + **({"callback": progress.update} if supports_callback else {}), + ) + except Exception as e: + raise RuntimeError( + "An error occurred while downloading using `hf_transfer`. Consider" + " disabling HF_HUB_ENABLE_HF_TRANSFER for better error handling." + ) from e + if not supports_callback: + progress.update(total) + if expected_size is not None and expected_size != os.path.getsize(temp_file.name): + raise EnvironmentError( + consistency_error_message.format( + actual_size=os.path.getsize(temp_file.name), + ) + ) + return + new_resume_size = resume_size + try: + for chunk in r.iter_content(chunk_size=constants.DOWNLOAD_CHUNK_SIZE): + if chunk: # filter out keep-alive new chunks + progress.update(len(chunk)) + temp_file.write(chunk) + new_resume_size += len(chunk) + # Some data has been downloaded from the server so we reset the number of retries. + _nb_retries = 5 + except (requests.ConnectionError, requests.ReadTimeout) as e: + # If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely + # a transient error (network outage?). We log a warning message and try to resume the download a few times + # before giving up. Tre retry mechanism is basic but should be enough in most cases. + if _nb_retries <= 0: + logger.warning("Error while downloading from %s: %s\nMax retries exceeded.", url, str(e)) + raise + logger.warning("Error while downloading from %s: %s\nTrying to resume download...", url, str(e)) + time.sleep(1) + reset_sessions() # In case of SSLError it's best to reset the shared requests.Session objects + return http_get( + url=url, + temp_file=temp_file, + proxies=proxies, + resume_size=new_resume_size, + headers=initial_headers, + expected_size=expected_size, + _nb_retries=_nb_retries - 1, + _tqdm_bar=_tqdm_bar, + ) + + if expected_size is not None and expected_size != temp_file.tell(): + raise EnvironmentError( + consistency_error_message.format( + actual_size=temp_file.tell(), + ) + ) + + +def xet_get( + *, + incomplete_path: Path, + xet_file_data: XetFileData, + headers: Dict[str, str], + expected_size: Optional[int] = None, + displayed_filename: Optional[str] = None, + _tqdm_bar: Optional[tqdm] = None, +) -> None: + """ + Download a file using Xet storage service. + + Args: + incomplete_path (`Path`): + The path to the file to download. + xet_file_data (`XetFileData`): + The file metadata needed to make the request to the xet storage service. + headers (`Dict[str, str]`): + The headers to send to the xet storage service. + expected_size (`int`, *optional*): + The expected size of the file to download. If set, the download will raise an error if the size of the + received content is different from the expected one. + displayed_filename (`str`, *optional*): + The filename of the file that is being downloaded. Value is used only to display a nice progress bar. If + not set, the filename is guessed from the URL or the `Content-Disposition` header. + + **How it works:** + The file download system uses Xet storage, which is a content-addressable storage system that breaks files into chunks + for efficient storage and transfer. + + `hf_xet.download_files` manages downloading files by: + - Taking a list of files to download (each with its unique content hash) + - Connecting to a storage server (CAS server) that knows how files are chunked + - Using authentication to ensure secure access + - Providing progress updates during download + + Authentication works by regularly refreshing access tokens through `refresh_xet_connection_info` to maintain a valid + connection to the storage server. + + The download process works like this: + 1. Create a local cache folder at `~/.cache/huggingface/xet/chunk-cache` to store reusable file chunks + 2. Download files in parallel: + 2.1. Prepare to write the file to disk + 2.2. Ask the server "how is this file split into chunks?" using the file's unique hash + The server responds with: + - Which chunks make up the complete file + - Where each chunk can be downloaded from + 2.3. For each needed chunk: + - Checks if we already have it in our local cache + - If not, download it from cloud storage (S3) + - Save it to cache for future use + - Assemble the chunks in order to recreate the original file + + """ + try: + from hf_xet import PyXetDownloadInfo, download_files # type: ignore[no-redef] + except ImportError: + raise ValueError( + "To use optimized download using Xet storage, you need to install the hf_xet package. " + 'Try `pip install "huggingface_hub[hf_xet]"` or `pip install hf_xet`.' + ) + + connection_info = refresh_xet_connection_info(file_data=xet_file_data, headers=headers) + + def token_refresher() -> Tuple[str, int]: + connection_info = refresh_xet_connection_info(file_data=xet_file_data, headers=headers) + if connection_info is None: + raise ValueError("Failed to refresh token using xet metadata.") + return connection_info.access_token, connection_info.expiration_unix_epoch + + xet_download_info = [ + PyXetDownloadInfo( + destination_path=str(incomplete_path.absolute()), hash=xet_file_data.file_hash, file_size=expected_size + ) + ] + + if not displayed_filename: + displayed_filename = incomplete_path.name + + # Truncate filename if too long to display + if len(displayed_filename) > 40: + displayed_filename = f"{displayed_filename[:40]}(…)" + + progress_cm = _get_progress_bar_context( + desc=displayed_filename, + log_level=logger.getEffectiveLevel(), + total=expected_size, + initial=0, + name="huggingface_hub.xet_get", + _tqdm_bar=_tqdm_bar, + ) + + with progress_cm as progress: + + def progress_updater(progress_bytes: float): + progress.update(progress_bytes) + + download_files( + xet_download_info, + endpoint=connection_info.endpoint, + token_info=(connection_info.access_token, connection_info.expiration_unix_epoch), + token_refresher=token_refresher, + progress_updater=[progress_updater], + ) + + +def _normalize_etag(etag: Optional[str]) -> Optional[str]: + """Normalize ETag HTTP header, so it can be used to create nice filepaths. + + The HTTP spec allows two forms of ETag: + ETag: W/"" + ETag: "" + + For now, we only expect the second form from the server, but we want to be future-proof so we support both. For + more context, see `TestNormalizeEtag` tests and https://github.com/huggingface/huggingface_hub/pull/1428. + + Args: + etag (`str`, *optional*): HTTP header + + Returns: + `str` or `None`: string that can be used as a nice directory name. + Returns `None` if input is None. + """ + if etag is None: + return None + return etag.lstrip("W/").strip('"') + + +def _create_relative_symlink(src: str, dst: str, new_blob: bool = False) -> None: + """Alias method used in `transformers` conversion script.""" + return _create_symlink(src=src, dst=dst, new_blob=new_blob) + + +def _create_symlink(src: str, dst: str, new_blob: bool = False) -> None: + """Create a symbolic link named dst pointing to src. + + By default, it will try to create a symlink using a relative path. Relative paths have 2 advantages: + - If the cache_folder is moved (example: back-up on a shared drive), relative paths within the cache folder will + not break. + - Relative paths seems to be better handled on Windows. Issue was reported 3 times in less than a week when + changing from relative to absolute paths. See https://github.com/huggingface/huggingface_hub/issues/1398, + https://github.com/huggingface/diffusers/issues/2729 and https://github.com/huggingface/transformers/pull/22228. + NOTE: The issue with absolute paths doesn't happen on admin mode. + When creating a symlink from the cache to a local folder, it is possible that a relative path cannot be created. + This happens when paths are not on the same volume. In that case, we use absolute paths. + + + The result layout looks something like + └── [ 128] snapshots + ├── [ 128] 2439f60ef33a0d46d85da5001d52aeda5b00ce9f + │ ├── [ 52] README.md -> ../../../blobs/d7edf6bd2a681fb0175f7735299831ee1b22b812 + │ └── [ 76] pytorch_model.bin -> ../../../blobs/403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + + If symlinks cannot be created on this platform (most likely to be Windows), the workaround is to avoid symlinks by + having the actual file in `dst`. If it is a new file (`new_blob=True`), we move it to `dst`. If it is not a new file + (`new_blob=False`), we don't know if the blob file is already referenced elsewhere. To avoid breaking existing + cache, the file is duplicated on the disk. + + In case symlinks are not supported, a warning message is displayed to the user once when loading `huggingface_hub`. + The warning message can be disabled with the `DISABLE_SYMLINKS_WARNING` environment variable. + """ + try: + os.remove(dst) + except OSError: + pass + + abs_src = os.path.abspath(os.path.expanduser(src)) + abs_dst = os.path.abspath(os.path.expanduser(dst)) + abs_dst_folder = os.path.dirname(abs_dst) + + # Use relative_dst in priority + try: + relative_src = os.path.relpath(abs_src, abs_dst_folder) + except ValueError: + # Raised on Windows if src and dst are not on the same volume. This is the case when creating a symlink to a + # local_dir instead of within the cache directory. + # See https://docs.python.org/3/library/os.path.html#os.path.relpath + relative_src = None + + try: + commonpath = os.path.commonpath([abs_src, abs_dst]) + _support_symlinks = are_symlinks_supported(commonpath) + except ValueError: + # Raised if src and dst are not on the same volume. Symlinks will still work on Linux/Macos. + # See https://docs.python.org/3/library/os.path.html#os.path.commonpath + _support_symlinks = os.name != "nt" + except PermissionError: + # Permission error means src and dst are not in the same volume (e.g. destination path has been provided + # by the user via `local_dir`. Let's test symlink support there) + _support_symlinks = are_symlinks_supported(abs_dst_folder) + except OSError as e: + # OS error (errno=30) means that the commonpath is readonly on Linux/MacOS. + if e.errno == errno.EROFS: + _support_symlinks = are_symlinks_supported(abs_dst_folder) + else: + raise + + # Symlinks are supported => let's create a symlink. + if _support_symlinks: + src_rel_or_abs = relative_src or abs_src + logger.debug(f"Creating pointer from {src_rel_or_abs} to {abs_dst}") + try: + os.symlink(src_rel_or_abs, abs_dst) + return + except FileExistsError: + if os.path.islink(abs_dst) and os.path.realpath(abs_dst) == os.path.realpath(abs_src): + # `abs_dst` already exists and is a symlink to the `abs_src` blob. It is most likely that the file has + # been cached twice concurrently (exactly between `os.remove` and `os.symlink`). Do nothing. + return + else: + # Very unlikely to happen. Means a file `dst` has been created exactly between `os.remove` and + # `os.symlink` and is not a symlink to the `abs_src` blob file. Raise exception. + raise + except PermissionError: + # Permission error means src and dst are not in the same volume (e.g. download to local dir) and symlink + # is supported on both volumes but not between them. Let's just make a hard copy in that case. + pass + + # Symlinks are not supported => let's move or copy the file. + if new_blob: + logger.info(f"Symlink not supported. Moving file from {abs_src} to {abs_dst}") + shutil.move(abs_src, abs_dst, copy_function=_copy_no_matter_what) + else: + logger.info(f"Symlink not supported. Copying file from {abs_src} to {abs_dst}") + shutil.copyfile(abs_src, abs_dst) + + +def _cache_commit_hash_for_specific_revision(storage_folder: str, revision: str, commit_hash: str) -> None: + """Cache reference between a revision (tag, branch or truncated commit hash) and the corresponding commit hash. + + Does nothing if `revision` is already a proper `commit_hash` or reference is already cached. + """ + if revision != commit_hash: + ref_path = Path(storage_folder) / "refs" / revision + ref_path.parent.mkdir(parents=True, exist_ok=True) + if not ref_path.exists() or commit_hash != ref_path.read_text(): + # Update ref only if has been updated. Could cause useless error in case + # repo is already cached and user doesn't have write access to cache folder. + # See https://github.com/huggingface/huggingface_hub/issues/1216. + ref_path.write_text(commit_hash) + + +@validate_hf_hub_args +def repo_folder_name(*, repo_id: str, repo_type: str) -> str: + """Return a serialized version of a hf.co repo name and type, safe for disk storage + as a single non-nested folder. + + Example: models--julien-c--EsperBERTo-small + """ + # remove all `/` occurrences to correctly convert repo to directory name + parts = [f"{repo_type}s", *repo_id.split("/")] + return constants.REPO_ID_SEPARATOR.join(parts) + + +def _check_disk_space(expected_size: int, target_dir: Union[str, Path]) -> None: + """Check disk usage and log a warning if there is not enough disk space to download the file. + + Args: + expected_size (`int`): + The expected size of the file in bytes. + target_dir (`str`): + The directory where the file will be stored after downloading. + """ + + target_dir = Path(target_dir) # format as `Path` + for path in [target_dir] + list(target_dir.parents): # first check target_dir, then each parents one by one + try: + target_dir_free = shutil.disk_usage(path).free + if target_dir_free < expected_size: + warnings.warn( + "Not enough free disk space to download the file. " + f"The expected file size is: {expected_size / 1e6:.2f} MB. " + f"The target location {target_dir} only has {target_dir_free / 1e6:.2f} MB free disk space." + ) + return + except OSError: # raise on anything: file does not exist or space disk cannot be checked + pass + + +@validate_hf_hub_args +def hf_hub_download( + repo_id: str, + filename: str, + *, + subfolder: Optional[str] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + cache_dir: Union[str, Path, None] = None, + local_dir: Union[str, Path, None] = None, + user_agent: Union[Dict, str, None] = None, + force_download: bool = False, + proxies: Optional[Dict] = None, + etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT, + token: Union[bool, str, None] = None, + local_files_only: bool = False, + headers: Optional[Dict[str, str]] = None, + endpoint: Optional[str] = None, + resume_download: Optional[bool] = None, + force_filename: Optional[str] = None, + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", +) -> str: + """Download a given file if it's not already present in the local cache. + + The new cache file layout looks like this: + - The cache directory contains one subfolder per repo_id (namespaced by repo type) + - inside each repo folder: + - refs is a list of the latest known revision => commit_hash pairs + - blobs contains the actual file blobs (identified by their git-sha or sha256, depending on + whether they're LFS files or not) + - snapshots contains one subfolder per commit, each "commit" contains the subset of the files + that have been resolved at that particular commit. Each filename is a symlink to the blob + at that particular commit. + + ``` + [ 96] . + └── [ 160] models--julien-c--EsperBERTo-small + ├── [ 160] blobs + │ ├── [321M] 403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + │ ├── [ 398] 7cb18dc9bafbfcf74629a4b760af1b160957a83e + │ └── [1.4K] d7edf6bd2a681fb0175f7735299831ee1b22b812 + ├── [ 96] refs + │ └── [ 40] main + └── [ 128] snapshots + ├── [ 128] 2439f60ef33a0d46d85da5001d52aeda5b00ce9f + │ ├── [ 52] README.md -> ../../blobs/d7edf6bd2a681fb0175f7735299831ee1b22b812 + │ └── [ 76] pytorch_model.bin -> ../../blobs/403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + └── [ 128] bbc77c8132af1cc5cf678da3f1ddf2de43606d48 + ├── [ 52] README.md -> ../../blobs/7cb18dc9bafbfcf74629a4b760af1b160957a83e + └── [ 76] pytorch_model.bin -> ../../blobs/403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + ``` + + If `local_dir` is provided, the file structure from the repo will be replicated in this location. When using this + option, the `cache_dir` will not be used and a `.cache/huggingface/` folder will be created at the root of `local_dir` + to store some metadata related to the downloaded files. While this mechanism is not as robust as the main + cache-system, it's optimized for regularly pulling the latest version of a repository. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + filename (`str`): + The name of the file in the repo. + subfolder (`str`, *optional*): + An optional value corresponding to a folder inside the model repo. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if downloading from a dataset or space, + `None` or `"model"` if downloading from a model. Default is `None`. + revision (`str`, *optional*): + An optional Git revision id which can be a branch name, a tag, or a + commit hash. + library_name (`str`, *optional*): + The name of the library to which the object corresponds. + library_version (`str`, *optional*): + The version of the library. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_dir (`str` or `Path`, *optional*): + If provided, the downloaded file will be placed under this directory. + user_agent (`dict`, `str`, *optional*): + The user-agent info in the form of a dictionary or a string. + force_download (`bool`, *optional*, defaults to `False`): + Whether the file should be downloaded even if it already exists in + the local cache. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to + `requests.request`. + etag_timeout (`float`, *optional*, defaults to `10`): + When fetching ETag, how many seconds to wait for the server to send + data before giving up which is passed to `requests.request`. + token (`str`, `bool`, *optional*): + A token to be used for the download. + - If `True`, the token is read from the HuggingFace config + folder. + - If a string, it's used as the authentication token. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the + local cached file if it exists. + headers (`dict`, *optional*): + Additional headers to be sent with the request. + + Returns: + `str`: Local path of file or if networking is off, last version of file cached on disk. + + Raises: + [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + [`~utils.EntryNotFoundError`] + If the file to download cannot be found. + [`~utils.LocalEntryNotFoundError`] + If network is disabled or unavailable and file is not found in cache. + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `token=True` but the token cannot be found. + [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) + If ETag cannot be determined. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If some parameter value is invalid. + + """ + if constants.HF_HUB_ETAG_TIMEOUT != constants.DEFAULT_ETAG_TIMEOUT: + # Respect environment variable above user value + etag_timeout = constants.HF_HUB_ETAG_TIMEOUT + + if force_filename is not None: + warnings.warn( + "The `force_filename` parameter is deprecated as a new caching system, " + "which keeps the filenames as they are on the Hub, is now in place.", + FutureWarning, + ) + if resume_download is not None: + warnings.warn( + "`resume_download` is deprecated and will be removed in version 1.0.0. " + "Downloads always resume when possible. " + "If you want to force a new download, use `force_download=True`.", + FutureWarning, + ) + + if cache_dir is None: + cache_dir = constants.HF_HUB_CACHE + if revision is None: + revision = constants.DEFAULT_REVISION + if isinstance(cache_dir, Path): + cache_dir = str(cache_dir) + if isinstance(local_dir, Path): + local_dir = str(local_dir) + + if subfolder == "": + subfolder = None + if subfolder is not None: + # This is used to create a URL, and not a local path, hence the forward slash. + filename = f"{subfolder}/{filename}" + + if repo_type is None: + repo_type = "model" + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type: {repo_type}. Accepted repo types are: {str(constants.REPO_TYPES)}") + + hf_headers = build_hf_headers( + token=token, + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + headers=headers, + ) + + if local_dir is not None: + if local_dir_use_symlinks != "auto": + warnings.warn( + "`local_dir_use_symlinks` parameter is deprecated and will be ignored. " + "The process to download files to a local folder has been updated and do " + "not rely on symlinks anymore. You only need to pass a destination folder " + "as`local_dir`.\n" + "For more details, check out https://huggingface.co/docs/huggingface_hub/main/en/guides/download#download-files-to-local-folder." + ) + + return _hf_hub_download_to_local_dir( + # Destination + local_dir=local_dir, + # File info + repo_id=repo_id, + repo_type=repo_type, + filename=filename, + revision=revision, + # HTTP info + endpoint=endpoint, + etag_timeout=etag_timeout, + headers=hf_headers, + proxies=proxies, + token=token, + # Additional options + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + ) + else: + return _hf_hub_download_to_cache_dir( + # Destination + cache_dir=cache_dir, + # File info + repo_id=repo_id, + filename=filename, + repo_type=repo_type, + revision=revision, + # HTTP info + endpoint=endpoint, + etag_timeout=etag_timeout, + headers=hf_headers, + proxies=proxies, + token=token, + # Additional options + local_files_only=local_files_only, + force_download=force_download, + ) + + +def _hf_hub_download_to_cache_dir( + *, + # Destination + cache_dir: str, + # File info + repo_id: str, + filename: str, + repo_type: str, + revision: str, + # HTTP info + endpoint: Optional[str], + etag_timeout: float, + headers: Dict[str, str], + proxies: Optional[Dict], + token: Optional[Union[bool, str]], + # Additional options + local_files_only: bool, + force_download: bool, +) -> str: + """Download a given file to a cache folder, if not already present. + + Method should not be called directly. Please use `hf_hub_download` instead. + """ + locks_dir = os.path.join(cache_dir, ".locks") + storage_folder = os.path.join(cache_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type)) + + # cross platform transcription of filename, to be used as a local file path. + relative_filename = os.path.join(*filename.split("/")) + if os.name == "nt": + if relative_filename.startswith("..\\") or "\\..\\" in relative_filename: + raise ValueError( + f"Invalid filename: cannot handle filename '{relative_filename}' on Windows. Please ask the repository" + " owner to rename this file." + ) + + # if user provides a commit_hash and they already have the file on disk, shortcut everything. + if REGEX_COMMIT_HASH.match(revision): + pointer_path = _get_pointer_path(storage_folder, revision, relative_filename) + if os.path.exists(pointer_path) and not force_download: + return pointer_path + + # Try to get metadata (etag, commit_hash, url, size) from the server. + # If we can't, a HEAD request error is returned. + (url_to_download, etag, commit_hash, expected_size, xet_file_data, head_call_error) = _get_metadata_or_catch_error( + repo_id=repo_id, + filename=filename, + repo_type=repo_type, + revision=revision, + endpoint=endpoint, + proxies=proxies, + etag_timeout=etag_timeout, + headers=headers, + token=token, + local_files_only=local_files_only, + storage_folder=storage_folder, + relative_filename=relative_filename, + ) + + # etag can be None for several reasons: + # 1. we passed local_files_only. + # 2. we don't have a connection + # 3. Hub is down (HTTP 500, 503, 504) + # 4. repo is not found -for example private or gated- and invalid/missing token sent + # 5. Hub is blocked by a firewall or proxy is not set correctly. + # => Try to get the last downloaded one from the specified revision. + # + # If the specified revision is a commit hash, look inside "snapshots". + # If the specified revision is a branch or tag, look inside "refs". + if head_call_error is not None: + # Couldn't make a HEAD call => let's try to find a local file + if not force_download: + commit_hash = None + if REGEX_COMMIT_HASH.match(revision): + commit_hash = revision + else: + ref_path = os.path.join(storage_folder, "refs", revision) + if os.path.isfile(ref_path): + with open(ref_path) as f: + commit_hash = f.read() + + # Return pointer file if exists + if commit_hash is not None: + pointer_path = _get_pointer_path(storage_folder, commit_hash, relative_filename) + if os.path.exists(pointer_path) and not force_download: + return pointer_path + + # Otherwise, raise appropriate error + _raise_on_head_call_error(head_call_error, force_download, local_files_only) + + # From now on, etag, commit_hash, url and size are not None. + assert etag is not None, "etag must have been retrieved from server" + assert commit_hash is not None, "commit_hash must have been retrieved from server" + assert url_to_download is not None, "file location must have been retrieved from server" + assert expected_size is not None, "expected_size must have been retrieved from server" + blob_path = os.path.join(storage_folder, "blobs", etag) + pointer_path = _get_pointer_path(storage_folder, commit_hash, relative_filename) + + os.makedirs(os.path.dirname(blob_path), exist_ok=True) + os.makedirs(os.path.dirname(pointer_path), exist_ok=True) + + # if passed revision is not identical to commit_hash + # then revision has to be a branch name or tag name. + # In that case store a ref. + _cache_commit_hash_for_specific_revision(storage_folder, revision, commit_hash) + + # If file already exists, return it (except if force_download=True) + if not force_download: + if os.path.exists(pointer_path): + return pointer_path + + if os.path.exists(blob_path): + # we have the blob already, but not the pointer + _create_symlink(blob_path, pointer_path, new_blob=False) + return pointer_path + + # Prevent parallel downloads of the same file with a lock. + # etag could be duplicated across repos, + lock_path = os.path.join(locks_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type), f"{etag}.lock") + + # Some Windows versions do not allow for paths longer than 255 characters. + # In this case, we must specify it as an extended path by using the "\\?\" prefix. + if os.name == "nt" and len(os.path.abspath(lock_path)) > 255: + lock_path = "\\\\?\\" + os.path.abspath(lock_path) + + if os.name == "nt" and len(os.path.abspath(blob_path)) > 255: + blob_path = "\\\\?\\" + os.path.abspath(blob_path) + + # Local file doesn't exist or etag isn't a match => retrieve file from remote (or cache) + + Path(lock_path).parent.mkdir(parents=True, exist_ok=True) + with WeakFileLock(lock_path): + _download_to_tmp_and_move( + incomplete_path=Path(blob_path + ".incomplete"), + destination_path=Path(blob_path), + url_to_download=url_to_download, + proxies=proxies, + headers=headers, + expected_size=expected_size, + filename=filename, + force_download=force_download, + etag=etag, + xet_file_data=xet_file_data, + ) + if not os.path.exists(pointer_path): + _create_symlink(blob_path, pointer_path, new_blob=True) + + return pointer_path + + +def _hf_hub_download_to_local_dir( + *, + # Destination + local_dir: Union[str, Path], + # File info + repo_id: str, + repo_type: str, + filename: str, + revision: str, + # HTTP info + endpoint: Optional[str], + etag_timeout: float, + headers: Dict[str, str], + proxies: Optional[Dict], + token: Union[bool, str, None], + # Additional options + cache_dir: str, + force_download: bool, + local_files_only: bool, +) -> str: + """Download a given file to a local folder, if not already present. + + Method should not be called directly. Please use `hf_hub_download` instead. + """ + # Some Windows versions do not allow for paths longer than 255 characters. + # In this case, we must specify it as an extended path by using the "\\?\" prefix. + if os.name == "nt" and len(os.path.abspath(local_dir)) > 255: + local_dir = "\\\\?\\" + os.path.abspath(local_dir) + local_dir = Path(local_dir) + paths = get_local_download_paths(local_dir=local_dir, filename=filename) + local_metadata = read_download_metadata(local_dir=local_dir, filename=filename) + + # Local file exists + metadata exists + commit_hash matches => return file + if ( + not force_download + and REGEX_COMMIT_HASH.match(revision) + and paths.file_path.is_file() + and local_metadata is not None + and local_metadata.commit_hash == revision + ): + return str(paths.file_path) + + # Local file doesn't exist or commit_hash doesn't match => we need the etag + (url_to_download, etag, commit_hash, expected_size, xet_file_data, head_call_error) = _get_metadata_or_catch_error( + repo_id=repo_id, + filename=filename, + repo_type=repo_type, + revision=revision, + endpoint=endpoint, + proxies=proxies, + etag_timeout=etag_timeout, + headers=headers, + token=token, + local_files_only=local_files_only, + ) + + if head_call_error is not None: + # No HEAD call but local file exists => default to local file + if not force_download and paths.file_path.is_file(): + logger.warning( + f"Couldn't access the Hub to check for update but local file already exists. Defaulting to existing file. (error: {head_call_error})" + ) + return str(paths.file_path) + # Otherwise => raise + _raise_on_head_call_error(head_call_error, force_download, local_files_only) + + # From now on, etag, commit_hash, url and size are not None. + assert etag is not None, "etag must have been retrieved from server" + assert commit_hash is not None, "commit_hash must have been retrieved from server" + assert url_to_download is not None, "file location must have been retrieved from server" + assert expected_size is not None, "expected_size must have been retrieved from server" + + # Local file exists => check if it's up-to-date + if not force_download and paths.file_path.is_file(): + # etag matches => update metadata and return file + if local_metadata is not None and local_metadata.etag == etag: + write_download_metadata(local_dir=local_dir, filename=filename, commit_hash=commit_hash, etag=etag) + return str(paths.file_path) + + # metadata is outdated + etag is a sha256 + # => means it's an LFS file (large) + # => let's compute local hash and compare + # => if match, update metadata and return file + if local_metadata is None and REGEX_SHA256.match(etag) is not None: + with open(paths.file_path, "rb") as f: + file_hash = sha_fileobj(f).hex() + if file_hash == etag: + write_download_metadata(local_dir=local_dir, filename=filename, commit_hash=commit_hash, etag=etag) + return str(paths.file_path) + + # Local file doesn't exist or etag isn't a match => retrieve file from remote (or cache) + + # If we are lucky enough, the file is already in the cache => copy it + if not force_download: + cached_path = try_to_load_from_cache( + repo_id=repo_id, + filename=filename, + cache_dir=cache_dir, + revision=commit_hash, + repo_type=repo_type, + ) + if isinstance(cached_path, str): + with WeakFileLock(paths.lock_path): + paths.file_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(cached_path, paths.file_path) + write_download_metadata(local_dir=local_dir, filename=filename, commit_hash=commit_hash, etag=etag) + return str(paths.file_path) + + # Otherwise, let's download the file! + with WeakFileLock(paths.lock_path): + paths.file_path.unlink(missing_ok=True) # delete outdated file first + _download_to_tmp_and_move( + incomplete_path=paths.incomplete_path(etag), + destination_path=paths.file_path, + url_to_download=url_to_download, + proxies=proxies, + headers=headers, + expected_size=expected_size, + filename=filename, + force_download=force_download, + etag=etag, + xet_file_data=xet_file_data, + ) + + write_download_metadata(local_dir=local_dir, filename=filename, commit_hash=commit_hash, etag=etag) + return str(paths.file_path) + + +@validate_hf_hub_args +def try_to_load_from_cache( + repo_id: str, + filename: str, + cache_dir: Union[str, Path, None] = None, + revision: Optional[str] = None, + repo_type: Optional[str] = None, +) -> Union[str, _CACHED_NO_EXIST_T, None]: + """ + Explores the cache to return the latest cached file for a given revision if found. + + This function will not raise any exception if the file in not cached. + + Args: + cache_dir (`str` or `os.PathLike`): + The folder where the cached files lie. + repo_id (`str`): + The ID of the repo on huggingface.co. + filename (`str`): + The filename to look for inside `repo_id`. + revision (`str`, *optional*): + The specific model version to use. Will default to `"main"` if it's not provided and no `commit_hash` is + provided either. + repo_type (`str`, *optional*): + The type of the repository. Will default to `"model"`. + + Returns: + `Optional[str]` or `_CACHED_NO_EXIST`: + Will return `None` if the file was not cached. Otherwise: + - The exact path to the cached file if it's found in the cache + - A special value `_CACHED_NO_EXIST` if the file does not exist at the given commit hash and this fact was + cached. + + Example: + + ```python + from huggingface_hub import try_to_load_from_cache, _CACHED_NO_EXIST + + filepath = try_to_load_from_cache() + if isinstance(filepath, str): + # file exists and is cached + ... + elif filepath is _CACHED_NO_EXIST: + # non-existence of file is cached + ... + else: + # file is not cached + ... + ``` + """ + if revision is None: + revision = "main" + if repo_type is None: + repo_type = "model" + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type: {repo_type}. Accepted repo types are: {str(constants.REPO_TYPES)}") + if cache_dir is None: + cache_dir = constants.HF_HUB_CACHE + + object_id = repo_id.replace("/", "--") + repo_cache = os.path.join(cache_dir, f"{repo_type}s--{object_id}") + if not os.path.isdir(repo_cache): + # No cache for this model + return None + + refs_dir = os.path.join(repo_cache, "refs") + snapshots_dir = os.path.join(repo_cache, "snapshots") + no_exist_dir = os.path.join(repo_cache, ".no_exist") + + # Resolve refs (for instance to convert main to the associated commit sha) + if os.path.isdir(refs_dir): + revision_file = os.path.join(refs_dir, revision) + if os.path.isfile(revision_file): + with open(revision_file) as f: + revision = f.read() + + # Check if file is cached as "no_exist" + if os.path.isfile(os.path.join(no_exist_dir, revision, filename)): + return _CACHED_NO_EXIST + + # Check if revision folder exists + if not os.path.exists(snapshots_dir): + return None + cached_shas = os.listdir(snapshots_dir) + if revision not in cached_shas: + # No cache for this revision and we won't try to return a random revision + return None + + # Check if file exists in cache + cached_file = os.path.join(snapshots_dir, revision, filename) + return cached_file if os.path.isfile(cached_file) else None + + +@validate_hf_hub_args +def get_hf_file_metadata( + url: str, + token: Union[bool, str, None] = None, + proxies: Optional[Dict] = None, + timeout: Optional[float] = constants.DEFAULT_REQUEST_TIMEOUT, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, + headers: Optional[Dict[str, str]] = None, +) -> HfFileMetadata: + """Fetch metadata of a file versioned on the Hub for a given url. + + Args: + url (`str`): + File url, for example returned by [`hf_hub_url`]. + token (`str` or `bool`, *optional*): + A token to be used for the download. + - If `True`, the token is read from the HuggingFace config + folder. + - If `False` or `None`, no token is provided. + - If a string, it's used as the authentication token. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to + `requests.request`. + timeout (`float`, *optional*, defaults to 10): + How many seconds to wait for the server to send metadata before giving up. + library_name (`str`, *optional*): + The name of the library to which the object corresponds. + library_version (`str`, *optional*): + The version of the library. + user_agent (`dict`, `str`, *optional*): + The user-agent info in the form of a dictionary or a string. + headers (`dict`, *optional*): + Additional headers to be sent with the request. + + Returns: + A [`HfFileMetadata`] object containing metadata such as location, etag, size and + commit_hash. + """ + hf_headers = build_hf_headers( + token=token, + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + headers=headers, + ) + hf_headers["Accept-Encoding"] = "identity" # prevent any compression => we want to know the real size of the file + + # Retrieve metadata + r = _request_wrapper( + method="HEAD", + url=url, + headers=hf_headers, + allow_redirects=False, + follow_relative_redirects=True, + proxies=proxies, + timeout=timeout, + ) + hf_raise_for_status(r) + + # Return + return HfFileMetadata( + commit_hash=r.headers.get(constants.HUGGINGFACE_HEADER_X_REPO_COMMIT), + # We favor a custom header indicating the etag of the linked resource, and + # we fallback to the regular etag header. + etag=_normalize_etag(r.headers.get(constants.HUGGINGFACE_HEADER_X_LINKED_ETAG) or r.headers.get("ETag")), + # Either from response headers (if redirected) or defaults to request url + # Do not use directly `url`, as `_request_wrapper` might have followed relative + # redirects. + location=r.headers.get("Location") or r.request.url, # type: ignore + size=_int_or_none( + r.headers.get(constants.HUGGINGFACE_HEADER_X_LINKED_SIZE) or r.headers.get("Content-Length") + ), + xet_file_data=parse_xet_file_data_from_response(r), # type: ignore + ) + + +def _get_metadata_or_catch_error( + *, + repo_id: str, + filename: str, + repo_type: str, + revision: str, + endpoint: Optional[str], + proxies: Optional[Dict], + etag_timeout: Optional[float], + headers: Dict[str, str], # mutated inplace! + token: Union[bool, str, None], + local_files_only: bool, + relative_filename: Optional[str] = None, # only used to store `.no_exists` in cache + storage_folder: Optional[str] = None, # only used to store `.no_exists` in cache +) -> Union[ + # Either an exception is caught and returned + Tuple[None, None, None, None, None, Exception], + # Or the metadata is returned as + # `(url_to_download, etag, commit_hash, expected_size, xet_file_data, None)` + Tuple[str, str, str, int, Optional[XetFileData], None], +]: + """Get metadata for a file on the Hub, safely handling network issues. + + Returns either the etag, commit_hash and expected size of the file, or the error + raised while fetching the metadata. + + NOTE: This function mutates `headers` inplace! It removes the `authorization` header + if the file is a LFS blob and the domain of the url is different from the + domain of the location (typically an S3 bucket). + """ + if local_files_only: + return ( + None, + None, + None, + None, + None, + OfflineModeIsEnabled( + f"Cannot access file since 'local_files_only=True' as been set. (repo_id: {repo_id}, repo_type: {repo_type}, revision: {revision}, filename: {filename})" + ), + ) + + url = hf_hub_url(repo_id, filename, repo_type=repo_type, revision=revision, endpoint=endpoint) + url_to_download: str = url + etag: Optional[str] = None + commit_hash: Optional[str] = None + expected_size: Optional[int] = None + head_error_call: Optional[Exception] = None + xet_file_data: Optional[XetFileData] = None + + # Try to get metadata from the server. + # Do not raise yet if the file is not found or not accessible. + if not local_files_only: + try: + try: + metadata = get_hf_file_metadata( + url=url, proxies=proxies, timeout=etag_timeout, headers=headers, token=token + ) + except EntryNotFoundError as http_error: + if storage_folder is not None and relative_filename is not None: + # Cache the non-existence of the file + commit_hash = http_error.response.headers.get(constants.HUGGINGFACE_HEADER_X_REPO_COMMIT) + if commit_hash is not None: + no_exist_file_path = Path(storage_folder) / ".no_exist" / commit_hash / relative_filename + try: + no_exist_file_path.parent.mkdir(parents=True, exist_ok=True) + no_exist_file_path.touch() + except OSError as e: + logger.error( + f"Could not cache non-existence of file. Will ignore error and continue. Error: {e}" + ) + _cache_commit_hash_for_specific_revision(storage_folder, revision, commit_hash) + raise + + # Commit hash must exist + commit_hash = metadata.commit_hash + if commit_hash is None: + raise FileMetadataError( + "Distant resource does not seem to be on huggingface.co. It is possible that a configuration issue" + " prevents you from downloading resources from https://huggingface.co. Please check your firewall" + " and proxy settings and make sure your SSL certificates are updated." + ) + + # Etag must exist + # If we don't have any of those, raise an error. + etag = metadata.etag + if etag is None: + raise FileMetadataError( + "Distant resource does not have an ETag, we won't be able to reliably ensure reproducibility." + ) + + # Size must exist + expected_size = metadata.size + if expected_size is None: + raise FileMetadataError("Distant resource does not have a Content-Length.") + + xet_file_data = metadata.xet_file_data + + # In case of a redirect, save an extra redirect on the request.get call, + # and ensure we download the exact atomic version even if it changed + # between the HEAD and the GET (unlikely, but hey). + # + # If url domain is different => we are downloading from a CDN => url is signed => don't send auth + # If url domain is the same => redirect due to repo rename AND downloading a regular file => keep auth + if xet_file_data is None and url != metadata.location: + url_to_download = metadata.location + if urlparse(url).netloc != urlparse(metadata.location).netloc: + # Remove authorization header when downloading a LFS blob + headers.pop("authorization", None) + except (requests.exceptions.SSLError, requests.exceptions.ProxyError): + # Actually raise for those subclasses of ConnectionError + raise + except ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + OfflineModeIsEnabled, + ) as error: + # Otherwise, our Internet connection is down. + # etag is None + head_error_call = error + except (RevisionNotFoundError, EntryNotFoundError): + # The repo was found but the revision or entry doesn't exist on the Hub (never existed or got deleted) + raise + except requests.HTTPError as error: + # Multiple reasons for an http error: + # - Repository is private and invalid/missing token sent + # - Repository is gated and invalid/missing token sent + # - Hub is down (error 500 or 504) + # => let's switch to 'local_files_only=True' to check if the files are already cached. + # (if it's not the case, the error will be re-raised) + head_error_call = error + except FileMetadataError as error: + # Multiple reasons for a FileMetadataError: + # - Wrong network configuration (proxy, firewall, SSL certificates) + # - Inconsistency on the Hub + # => let's switch to 'local_files_only=True' to check if the files are already cached. + # (if it's not the case, the error will be re-raised) + head_error_call = error + + if not (local_files_only or etag is not None or head_error_call is not None): + raise RuntimeError("etag is empty due to uncovered problems") + + return (url_to_download, etag, commit_hash, expected_size, xet_file_data, head_error_call) # type: ignore [return-value] + + +def _raise_on_head_call_error(head_call_error: Exception, force_download: bool, local_files_only: bool) -> NoReturn: + """Raise an appropriate error when the HEAD call failed and we cannot locate a local file.""" + # No head call => we cannot force download. + if force_download: + if local_files_only: + raise ValueError("Cannot pass 'force_download=True' and 'local_files_only=True' at the same time.") + elif isinstance(head_call_error, OfflineModeIsEnabled): + raise ValueError("Cannot pass 'force_download=True' when offline mode is enabled.") from head_call_error + else: + raise ValueError("Force download failed due to the above error.") from head_call_error + + # No head call + couldn't find an appropriate file on disk => raise an error. + if local_files_only: + raise LocalEntryNotFoundError( + "Cannot find the requested files in the disk cache and outgoing traffic has been disabled. To enable" + " hf.co look-ups and downloads online, set 'local_files_only' to False." + ) + elif isinstance(head_call_error, (RepositoryNotFoundError, GatedRepoError)) or ( + isinstance(head_call_error, HfHubHTTPError) and head_call_error.response.status_code == 401 + ): + # Repo not found or gated => let's raise the actual error + # Unauthorized => likely a token issue => let's raise the actual error + raise head_call_error + else: + # Otherwise: most likely a connection issue or Hub downtime => let's warn the user + raise LocalEntryNotFoundError( + "An error happened while trying to locate the file on the Hub and we cannot find the requested files" + " in the local cache. Please check your connection and try again or make sure your Internet connection" + " is on." + ) from head_call_error + + +def _download_to_tmp_and_move( + incomplete_path: Path, + destination_path: Path, + url_to_download: str, + proxies: Optional[Dict], + headers: Dict[str, str], + expected_size: Optional[int], + filename: str, + force_download: bool, + etag: Optional[str], + xet_file_data: Optional[XetFileData], +) -> None: + """Download content from a URL to a destination path. + + Internal logic: + - return early if file is already downloaded + - resume download if possible (from incomplete file) + - do not resume download if `force_download=True` or `HF_HUB_ENABLE_HF_TRANSFER=True` + - check disk space before downloading + - download content to a temporary file + - set correct permissions on temporary file + - move the temporary file to the destination path + + Both `incomplete_path` and `destination_path` must be on the same volume to avoid a local copy. + """ + if destination_path.exists() and not force_download: + # Do nothing if already exists (except if force_download=True) + return + + if incomplete_path.exists() and (force_download or (constants.HF_HUB_ENABLE_HF_TRANSFER and not proxies)): + # By default, we will try to resume the download if possible. + # However, if the user has set `force_download=True` or if `hf_transfer` is enabled, then we should + # not resume the download => delete the incomplete file. + message = f"Removing incomplete file '{incomplete_path}'" + if force_download: + message += " (force_download=True)" + elif constants.HF_HUB_ENABLE_HF_TRANSFER and not proxies: + message += " (hf_transfer=True)" + logger.info(message) + incomplete_path.unlink(missing_ok=True) + + with incomplete_path.open("ab") as f: + resume_size = f.tell() + message = f"Downloading '{filename}' to '{incomplete_path}'" + if resume_size > 0 and expected_size is not None: + message += f" (resume from {resume_size}/{expected_size})" + logger.info(message) + + if expected_size is not None: # might be None if HTTP header not set correctly + # Check disk space in both tmp and destination path + _check_disk_space(expected_size, incomplete_path.parent) + _check_disk_space(expected_size, destination_path.parent) + + if xet_file_data is not None and is_xet_available(): + logger.info("Xet Storage is enabled for this repo. Downloading file from Xet Storage..") + xet_get( + incomplete_path=incomplete_path, + xet_file_data=xet_file_data, + headers=headers, + expected_size=expected_size, + displayed_filename=filename, + ) + else: + if xet_file_data is not None: + logger.warning( + "Xet Storage is enabled for this repo, but the 'hf_xet' package is not installed. " + "Falling back to regular HTTP download. " + "For better performance, install the package with: `pip install huggingface_hub[hf_xet]` or `pip install hf_xet`" + ) + + http_get( + url_to_download, + f, + proxies=proxies, + resume_size=resume_size, + headers=headers, + expected_size=expected_size, + ) + + logger.info(f"Download complete. Moving file to {destination_path}") + _chmod_and_move(incomplete_path, destination_path) + + +def _int_or_none(value: Optional[str]) -> Optional[int]: + try: + return int(value) # type: ignore + except (TypeError, ValueError): + return None + + +def _chmod_and_move(src: Path, dst: Path) -> None: + """Set correct permission before moving a blob from tmp directory to cache dir. + + Do not take into account the `umask` from the process as there is no convenient way + to get it that is thread-safe. + + See: + - About umask: https://docs.python.org/3/library/os.html#os.umask + - Thread-safety: https://stackoverflow.com/a/70343066 + - About solution: https://github.com/huggingface/huggingface_hub/pull/1220#issuecomment-1326211591 + - Fix issue: https://github.com/huggingface/huggingface_hub/issues/1141 + - Fix issue: https://github.com/huggingface/huggingface_hub/issues/1215 + """ + # Get umask by creating a temporary file in the cached repo folder. + tmp_file = dst.parent.parent / f"tmp_{uuid.uuid4()}" + try: + tmp_file.touch() + cache_dir_mode = Path(tmp_file).stat().st_mode + os.chmod(str(src), stat.S_IMODE(cache_dir_mode)) + except OSError as e: + logger.warning( + f"Could not set the permissions on the file '{src}'. Error: {e}.\nContinuing without setting permissions." + ) + finally: + try: + tmp_file.unlink() + except OSError: + # fails if `tmp_file.touch()` failed => do nothing + # See https://github.com/huggingface/huggingface_hub/issues/2359 + pass + + shutil.move(str(src), str(dst), copy_function=_copy_no_matter_what) + + +def _copy_no_matter_what(src: str, dst: str) -> None: + """Copy file from src to dst. + + If `shutil.copy2` fails, fallback to `shutil.copyfile`. + """ + try: + # Copy file with metadata and permission + # Can fail e.g. if dst is an S3 mount + shutil.copy2(src, dst) + except OSError: + # Copy only file content + shutil.copyfile(src, dst) + + +def _get_pointer_path(storage_folder: str, revision: str, relative_filename: str) -> str: + # Using `os.path.abspath` instead of `Path.resolve()` to avoid resolving symlinks + snapshot_path = os.path.join(storage_folder, "snapshots") + pointer_path = os.path.join(snapshot_path, revision, relative_filename) + if Path(os.path.abspath(snapshot_path)) not in Path(os.path.abspath(pointer_path)).parents: + raise ValueError( + "Invalid pointer path: cannot create pointer path in snapshot folder if" + f" `storage_folder='{storage_folder}'`, `revision='{revision}'` and" + f" `relative_filename='{relative_filename}'`." + ) + return pointer_path diff --git a/lib/python3.12/site-packages/huggingface_hub/hf_api.py b/lib/python3.12/site-packages/huggingface_hub/hf_api.py new file mode 100644 index 0000000000000000000000000000000000000000..2c945198d9bb3b5e69af5baba9557d6c25692d4b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/hf_api.py @@ -0,0 +1,10040 @@ +# coding=utf-8 +# Copyright 2019-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import inspect +import io +import json +import re +import struct +import warnings +from collections import defaultdict +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import asdict, dataclass, field +from datetime import datetime +from functools import wraps +from itertools import islice +from pathlib import Path +from typing import ( + Any, + BinaryIO, + Callable, + Dict, + Iterable, + Iterator, + List, + Literal, + Optional, + Tuple, + TypeVar, + Union, + overload, +) +from urllib.parse import quote, unquote + +import requests +from requests.exceptions import HTTPError +from tqdm.auto import tqdm as base_tqdm +from tqdm.contrib.concurrent import thread_map + +from . import constants +from ._commit_api import ( + CommitOperation, + CommitOperationAdd, + CommitOperationCopy, + CommitOperationDelete, + _fetch_files_to_copy, + _fetch_upload_modes, + _prepare_commit_payload, + _upload_lfs_files, + _upload_xet_files, + _warn_on_overwriting_operations, +) +from ._inference_endpoints import InferenceEndpoint, InferenceEndpointType +from ._space_api import SpaceHardware, SpaceRuntime, SpaceStorage, SpaceVariable +from ._upload_large_folder import upload_large_folder_internal +from .community import ( + Discussion, + DiscussionComment, + DiscussionStatusChange, + DiscussionTitleChange, + DiscussionWithDetails, + deserialize_event, +) +from .constants import ( + DEFAULT_ETAG_TIMEOUT, # noqa: F401 # kept for backward compatibility + DEFAULT_REQUEST_TIMEOUT, # noqa: F401 # kept for backward compatibility + DEFAULT_REVISION, # noqa: F401 # kept for backward compatibility + DISCUSSION_STATUS, # noqa: F401 # kept for backward compatibility + DISCUSSION_TYPES, # noqa: F401 # kept for backward compatibility + ENDPOINT, # noqa: F401 # kept for backward compatibility + INFERENCE_ENDPOINTS_ENDPOINT, # noqa: F401 # kept for backward compatibility + REGEX_COMMIT_OID, # noqa: F401 # kept for backward compatibility + REPO_TYPE_MODEL, # noqa: F401 # kept for backward compatibility + REPO_TYPES, # noqa: F401 # kept for backward compatibility + REPO_TYPES_MAPPING, # noqa: F401 # kept for backward compatibility + REPO_TYPES_URL_PREFIXES, # noqa: F401 # kept for backward compatibility + SAFETENSORS_INDEX_FILE, # noqa: F401 # kept for backward compatibility + SAFETENSORS_MAX_HEADER_LENGTH, # noqa: F401 # kept for backward compatibility + SAFETENSORS_SINGLE_FILE, # noqa: F401 # kept for backward compatibility + SPACES_SDK_TYPES, # noqa: F401 # kept for backward compatibility + WEBHOOK_DOMAIN_T, # noqa: F401 # kept for backward compatibility + DiscussionStatusFilter, # noqa: F401 # kept for backward compatibility + DiscussionTypeFilter, # noqa: F401 # kept for backward compatibility +) +from .errors import ( + BadRequestError, + EntryNotFoundError, + GatedRepoError, + HfHubHTTPError, + RepositoryNotFoundError, + RevisionNotFoundError, +) +from .file_download import HfFileMetadata, get_hf_file_metadata, hf_hub_url +from .repocard_data import DatasetCardData, ModelCardData, SpaceCardData +from .utils import ( + DEFAULT_IGNORE_PATTERNS, + HfFolder, # noqa: F401 # kept for backward compatibility + LocalTokenNotFoundError, + NotASafetensorsRepoError, + SafetensorsFileMetadata, + SafetensorsParsingError, + SafetensorsRepoMetadata, + TensorInfo, + build_hf_headers, + chunk_iterable, + experimental, + filter_repo_objects, + fix_hf_endpoint_in_url, + get_session, + get_token, + hf_raise_for_status, + logging, + paginate, + parse_datetime, + validate_hf_hub_args, +) +from .utils import tqdm as hf_tqdm +from .utils._auth import _get_token_from_environment, _get_token_from_file, _get_token_from_google_colab +from .utils._deprecation import _deprecate_method +from .utils._runtime import is_xet_available +from .utils._typing import CallableT +from .utils.endpoint_helpers import _is_emission_within_threshold + + +R = TypeVar("R") # Return type +CollectionItemType_T = Literal["model", "dataset", "space", "paper"] + +ExpandModelProperty_T = Literal[ + "author", + "baseModels", + "cardData", + "childrenModelCount", + "config", + "createdAt", + "disabled", + "downloads", + "downloadsAllTime", + "gated", + "gguf", + "inference", + "inferenceProviderMapping", + "lastModified", + "library_name", + "likes", + "mask_token", + "model-index", + "pipeline_tag", + "private", + "resourceGroup", + "safetensors", + "sha", + "siblings", + "spaces", + "tags", + "transformersInfo", + "trendingScore", + "usedStorage", + "widgetData", + "xetEnabled", +] + +ExpandDatasetProperty_T = Literal[ + "author", + "cardData", + "citation", + "createdAt", + "description", + "disabled", + "downloads", + "downloadsAllTime", + "gated", + "lastModified", + "likes", + "paperswithcode_id", + "private", + "resourceGroup", + "sha", + "siblings", + "tags", + "trendingScore", + "usedStorage", + "xetEnabled", +] + +ExpandSpaceProperty_T = Literal[ + "author", + "cardData", + "createdAt", + "datasets", + "disabled", + "lastModified", + "likes", + "models", + "private", + "resourceGroup", + "runtime", + "sdk", + "sha", + "siblings", + "subdomain", + "tags", + "trendingScore", + "usedStorage", + "xetEnabled", +] + +USERNAME_PLACEHOLDER = "hf_user" +_REGEX_DISCUSSION_URL = re.compile(r".*/discussions/(\d+)$") + +_CREATE_COMMIT_NO_REPO_ERROR_MESSAGE = ( + "\nNote: Creating a commit assumes that the repo already exists on the" + " Huggingface Hub. Please use `create_repo` if it's not the case." +) +_AUTH_CHECK_NO_REPO_ERROR_MESSAGE = ( + "\nNote: The repository either does not exist or you do not have access rights." + " Please check the repository ID and your access permissions." + " If this is a private repository, ensure that your token is correct." +) +logger = logging.get_logger(__name__) + + +def repo_type_and_id_from_hf_id(hf_id: str, hub_url: Optional[str] = None) -> Tuple[Optional[str], Optional[str], str]: + """ + Returns the repo type and ID from a huggingface.co URL linking to a + repository + + Args: + hf_id (`str`): + An URL or ID of a repository on the HF hub. Accepted values are: + + - https://huggingface.co/// + - https://huggingface.co// + - hf://// + - hf:/// + - // + - / + - + hub_url (`str`, *optional*): + The URL of the HuggingFace Hub, defaults to https://huggingface.co + + Returns: + A tuple with three items: repo_type (`str` or `None`), namespace (`str` or + `None`) and repo_id (`str`). + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If URL cannot be parsed. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `repo_type` is unknown. + """ + input_hf_id = hf_id + + hub_url = re.sub(r"https?://", "", hub_url if hub_url is not None else constants.ENDPOINT) + is_hf_url = hub_url in hf_id and "@" not in hf_id + + HFFS_PREFIX = "hf://" + if hf_id.startswith(HFFS_PREFIX): # Remove "hf://" prefix if exists + hf_id = hf_id[len(HFFS_PREFIX) :] + + url_segments = hf_id.split("/") + is_hf_id = len(url_segments) <= 3 + + namespace: Optional[str] + if is_hf_url: + namespace, repo_id = url_segments[-2:] + if namespace == hub_url: + namespace = None + if len(url_segments) > 2 and hub_url not in url_segments[-3]: + repo_type = url_segments[-3] + elif namespace in constants.REPO_TYPES_MAPPING: + # Mean canonical dataset or model + repo_type = constants.REPO_TYPES_MAPPING[namespace] + namespace = None + else: + repo_type = None + elif is_hf_id: + if len(url_segments) == 3: + # Passed // or // + repo_type, namespace, repo_id = url_segments[-3:] + elif len(url_segments) == 2: + if url_segments[0] in constants.REPO_TYPES_MAPPING: + # Passed '' or 'datasets/' for a canonical model or dataset + repo_type = constants.REPO_TYPES_MAPPING[url_segments[0]] + namespace = None + repo_id = hf_id.split("/")[-1] + else: + # Passed / or / + namespace, repo_id = hf_id.split("/")[-2:] + repo_type = None + else: + # Passed + repo_id = url_segments[0] + namespace, repo_type = None, None + else: + raise ValueError(f"Unable to retrieve user and repo ID from the passed HF ID: {hf_id}") + + # Check if repo type is known (mapping "spaces" => "space" + empty value => `None`) + if repo_type in constants.REPO_TYPES_MAPPING: + repo_type = constants.REPO_TYPES_MAPPING[repo_type] + if repo_type == "": + repo_type = None + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Unknown `repo_type`: '{repo_type}' ('{input_hf_id}')") + + return repo_type, namespace, repo_id + + +@dataclass +class LastCommitInfo(dict): + oid: str + title: str + date: datetime + + def __post_init__(self): # hack to make LastCommitInfo backward compatible + self.update(asdict(self)) + + +@dataclass +class BlobLfsInfo(dict): + size: int + sha256: str + pointer_size: int + + def __post_init__(self): # hack to make BlobLfsInfo backward compatible + self.update(asdict(self)) + + +@dataclass +class BlobSecurityInfo(dict): + safe: bool # duplicate information with "status" field, keeping it for backward compatibility + status: str + av_scan: Optional[Dict] + pickle_import_scan: Optional[Dict] + + def __post_init__(self): # hack to make BlogSecurityInfo backward compatible + self.update(asdict(self)) + + +@dataclass +class TransformersInfo(dict): + auto_model: str + custom_class: Optional[str] = None + # possible `pipeline_tag` values: https://github.com/huggingface/huggingface.js/blob/3ee32554b8620644a6287e786b2a83bf5caf559c/packages/tasks/src/pipelines.ts#L72 + pipeline_tag: Optional[str] = None + processor: Optional[str] = None + + def __post_init__(self): # hack to make TransformersInfo backward compatible + self.update(asdict(self)) + + +@dataclass +class SafeTensorsInfo(dict): + parameters: Dict[str, int] + total: int + + def __post_init__(self): # hack to make SafeTensorsInfo backward compatible + self.update(asdict(self)) + + +@dataclass +class CommitInfo(str): + """Data structure containing information about a newly created commit. + + Returned by any method that creates a commit on the Hub: [`create_commit`], [`upload_file`], [`upload_folder`], + [`delete_file`], [`delete_folder`]. It inherits from `str` for backward compatibility but using methods specific + to `str` is deprecated. + + Attributes: + commit_url (`str`): + Url where to find the commit. + + commit_message (`str`): + The summary (first line) of the commit that has been created. + + commit_description (`str`): + Description of the commit that has been created. Can be empty. + + oid (`str`): + Commit hash id. Example: `"91c54ad1727ee830252e457677f467be0bfd8a57"`. + + pr_url (`str`, *optional*): + Url to the PR that has been created, if any. Populated when `create_pr=True` + is passed. + + pr_revision (`str`, *optional*): + Revision of the PR that has been created, if any. Populated when + `create_pr=True` is passed. Example: `"refs/pr/1"`. + + pr_num (`int`, *optional*): + Number of the PR discussion that has been created, if any. Populated when + `create_pr=True` is passed. Can be passed as `discussion_num` in + [`get_discussion_details`]. Example: `1`. + + repo_url (`RepoUrl`): + Repo URL of the commit containing info like repo_id, repo_type, etc. + + _url (`str`, *optional*): + Legacy url for `str` compatibility. Can be the url to the uploaded file on the Hub (if returned by + [`upload_file`]), to the uploaded folder on the Hub (if returned by [`upload_folder`]) or to the commit on + the Hub (if returned by [`create_commit`]). Defaults to `commit_url`. It is deprecated to use this + attribute. Please use `commit_url` instead. + """ + + commit_url: str + commit_message: str + commit_description: str + oid: str + pr_url: Optional[str] = None + + # Computed from `commit_url` in `__post_init__` + repo_url: RepoUrl = field(init=False) + + # Computed from `pr_url` in `__post_init__` + pr_revision: Optional[str] = field(init=False) + pr_num: Optional[str] = field(init=False) + + # legacy url for `str` compatibility (ex: url to uploaded file, url to uploaded folder, url to PR, etc.) + _url: str = field(repr=False, default=None) # type: ignore # defaults to `commit_url` + + def __new__(cls, *args, commit_url: str, _url: Optional[str] = None, **kwargs): + return str.__new__(cls, _url or commit_url) + + def __post_init__(self): + """Populate pr-related fields after initialization. + + See https://docs.python.org/3.10/library/dataclasses.html#post-init-processing. + """ + # Repo info + self.repo_url = RepoUrl(self.commit_url.split("/commit/")[0]) + + # PR info + if self.pr_url is not None: + self.pr_revision = _parse_revision_from_pr_url(self.pr_url) + self.pr_num = int(self.pr_revision.split("/")[-1]) + else: + self.pr_revision = None + self.pr_num = None + + +@dataclass +class AccessRequest: + """Data structure containing information about a user access request. + + Attributes: + username (`str`): + Username of the user who requested access. + fullname (`str`): + Fullname of the user who requested access. + email (`Optional[str]`): + Email of the user who requested access. + Can only be `None` in the /accepted list if the user was granted access manually. + timestamp (`datetime`): + Timestamp of the request. + status (`Literal["pending", "accepted", "rejected"]`): + Status of the request. Can be one of `["pending", "accepted", "rejected"]`. + fields (`Dict[str, Any]`, *optional*): + Additional fields filled by the user in the gate form. + """ + + username: str + fullname: str + email: Optional[str] + timestamp: datetime + status: Literal["pending", "accepted", "rejected"] + + # Additional fields filled by the user in the gate form + fields: Optional[Dict[str, Any]] = None + + +@dataclass +class WebhookWatchedItem: + """Data structure containing information about the items watched by a webhook. + + Attributes: + type (`Literal["dataset", "model", "org", "space", "user"]`): + Type of the item to be watched. Can be one of `["dataset", "model", "org", "space", "user"]`. + name (`str`): + Name of the item to be watched. Can be the username, organization name, model name, dataset name or space name. + """ + + type: Literal["dataset", "model", "org", "space", "user"] + name: str + + +@dataclass +class WebhookInfo: + """Data structure containing information about a webhook. + + Attributes: + id (`str`): + ID of the webhook. + url (`str`): + URL of the webhook. + watched (`List[WebhookWatchedItem]`): + List of items watched by the webhook, see [`WebhookWatchedItem`]. + domains (`List[WEBHOOK_DOMAIN_T]`): + List of domains the webhook is watching. Can be one of `["repo", "discussions"]`. + secret (`str`, *optional*): + Secret of the webhook. + disabled (`bool`): + Whether the webhook is disabled or not. + """ + + id: str + url: str + watched: List[WebhookWatchedItem] + domains: List[constants.WEBHOOK_DOMAIN_T] + secret: Optional[str] + disabled: bool + + +class RepoUrl(str): + """Subclass of `str` describing a repo URL on the Hub. + + `RepoUrl` is returned by `HfApi.create_repo`. It inherits from `str` for backward + compatibility. At initialization, the URL is parsed to populate properties: + - endpoint (`str`) + - namespace (`Optional[str]`) + - repo_name (`str`) + - repo_id (`str`) + - repo_type (`Literal["model", "dataset", "space"]`) + - url (`str`) + + Args: + url (`Any`): + String value of the repo url. + endpoint (`str`, *optional*): + Endpoint of the Hub. Defaults to . + + Example: + ```py + >>> RepoUrl('https://huggingface.co/gpt2') + RepoUrl('https://huggingface.co/gpt2', endpoint='https://huggingface.co', repo_type='model', repo_id='gpt2') + + >>> RepoUrl('https://hub-ci.huggingface.co/datasets/dummy_user/dummy_dataset', endpoint='https://hub-ci.huggingface.co') + RepoUrl('https://hub-ci.huggingface.co/datasets/dummy_user/dummy_dataset', endpoint='https://hub-ci.huggingface.co', repo_type='dataset', repo_id='dummy_user/dummy_dataset') + + >>> RepoUrl('hf://datasets/my-user/my-dataset') + RepoUrl('hf://datasets/my-user/my-dataset', endpoint='https://huggingface.co', repo_type='dataset', repo_id='user/dataset') + + >>> HfApi.create_repo("dummy_model") + RepoUrl('https://huggingface.co/Wauplin/dummy_model', endpoint='https://huggingface.co', repo_type='model', repo_id='Wauplin/dummy_model') + ``` + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If URL cannot be parsed. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `repo_type` is unknown. + """ + + def __new__(cls, url: Any, endpoint: Optional[str] = None): + url = fix_hf_endpoint_in_url(url, endpoint=endpoint) + return super(RepoUrl, cls).__new__(cls, url) + + def __init__(self, url: Any, endpoint: Optional[str] = None) -> None: + super().__init__() + # Parse URL + self.endpoint = endpoint or constants.ENDPOINT + repo_type, namespace, repo_name = repo_type_and_id_from_hf_id(self, hub_url=self.endpoint) + + # Populate fields + self.namespace = namespace + self.repo_name = repo_name + self.repo_id = repo_name if namespace is None else f"{namespace}/{repo_name}" + self.repo_type = repo_type or constants.REPO_TYPE_MODEL + self.url = str(self) # just in case it's needed + + def __repr__(self) -> str: + return f"RepoUrl('{self}', endpoint='{self.endpoint}', repo_type='{self.repo_type}', repo_id='{self.repo_id}')" + + +@dataclass +class RepoSibling: + """ + Contains basic information about a repo file inside a repo on the Hub. + + + + All attributes of this class are optional except `rfilename`. This is because only the file names are returned when + listing repositories on the Hub (with [`list_models`], [`list_datasets`] or [`list_spaces`]). If you need more + information like file size, blob id or lfs details, you must request them specifically from one repo at a time + (using [`model_info`], [`dataset_info`] or [`space_info`]) as it adds more constraints on the backend server to + retrieve these. + + + + Attributes: + rfilename (str): + file name, relative to the repo root. + size (`int`, *optional*): + The file's size, in bytes. This attribute is defined when `files_metadata` argument of [`repo_info`] is set + to `True`. It's `None` otherwise. + blob_id (`str`, *optional*): + The file's git OID. This attribute is defined when `files_metadata` argument of [`repo_info`] is set to + `True`. It's `None` otherwise. + lfs (`BlobLfsInfo`, *optional*): + The file's LFS metadata. This attribute is defined when`files_metadata` argument of [`repo_info`] is set to + `True` and the file is stored with Git LFS. It's `None` otherwise. + """ + + rfilename: str + size: Optional[int] = None + blob_id: Optional[str] = None + lfs: Optional[BlobLfsInfo] = None + + +@dataclass +class RepoFile: + """ + Contains information about a file on the Hub. + + Attributes: + path (str): + file path relative to the repo root. + size (`int`): + The file's size, in bytes. + blob_id (`str`): + The file's git OID. + lfs (`BlobLfsInfo`): + The file's LFS metadata. + last_commit (`LastCommitInfo`, *optional*): + The file's last commit metadata. Only defined if [`list_repo_tree`] and [`get_paths_info`] + are called with `expand=True`. + security (`BlobSecurityInfo`, *optional*): + The file's security scan metadata. Only defined if [`list_repo_tree`] and [`get_paths_info`] + are called with `expand=True`. + """ + + path: str + size: int + blob_id: str + lfs: Optional[BlobLfsInfo] = None + last_commit: Optional[LastCommitInfo] = None + security: Optional[BlobSecurityInfo] = None + + def __init__(self, **kwargs): + self.path = kwargs.pop("path") + self.size = kwargs.pop("size") + self.blob_id = kwargs.pop("oid") + lfs = kwargs.pop("lfs", None) + if lfs is not None: + lfs = BlobLfsInfo(size=lfs["size"], sha256=lfs["oid"], pointer_size=lfs["pointerSize"]) + self.lfs = lfs + last_commit = kwargs.pop("lastCommit", None) or kwargs.pop("last_commit", None) + if last_commit is not None: + last_commit = LastCommitInfo( + oid=last_commit["id"], title=last_commit["title"], date=parse_datetime(last_commit["date"]) + ) + self.last_commit = last_commit + security = kwargs.pop("securityFileStatus", None) + if security is not None: + safe = security["status"] == "safe" + security = BlobSecurityInfo( + safe=safe, + status=security["status"], + av_scan=security["avScan"], + pickle_import_scan=security["pickleImportScan"], + ) + self.security = security + + # backwards compatibility + self.rfilename = self.path + self.lastCommit = self.last_commit + + +@dataclass +class RepoFolder: + """ + Contains information about a folder on the Hub. + + Attributes: + path (str): + folder path relative to the repo root. + tree_id (`str`): + The folder's git OID. + last_commit (`LastCommitInfo`, *optional*): + The folder's last commit metadata. Only defined if [`list_repo_tree`] and [`get_paths_info`] + are called with `expand=True`. + """ + + path: str + tree_id: str + last_commit: Optional[LastCommitInfo] = None + + def __init__(self, **kwargs): + self.path = kwargs.pop("path") + self.tree_id = kwargs.pop("oid") + last_commit = kwargs.pop("lastCommit", None) or kwargs.pop("last_commit", None) + if last_commit is not None: + last_commit = LastCommitInfo( + oid=last_commit["id"], title=last_commit["title"], date=parse_datetime(last_commit["date"]) + ) + self.last_commit = last_commit + + +@dataclass +class InferenceProviderMapping: + hf_model_id: str + status: Literal["live", "staging"] + provider_id: str + task: str + + adapter: Optional[str] = None + adapter_weights_path: Optional[str] = None + + def __init__(self, **kwargs): + self.hf_model_id = kwargs.pop("hf_model_id") + self.status = kwargs.pop("status") + self.provider_id = kwargs.pop("providerId") + self.task = kwargs.pop("task") + self.adapter = kwargs.pop("adapter", None) + self.adapter_weights_path = kwargs.pop("adapterWeightsPath", None) + self.__dict__.update(**kwargs) + + +@dataclass +class ModelInfo: + """ + Contains information about a model on the Hub. + + + + Most attributes of this class are optional. This is because the data returned by the Hub depends on the query made. + In general, the more specific the query, the more information is returned. On the contrary, when listing models + using [`list_models`] only a subset of the attributes are returned. + + + + Attributes: + id (`str`): + ID of model. + author (`str`, *optional*): + Author of the model. + sha (`str`, *optional*): + Repo SHA at this particular revision. + created_at (`datetime`, *optional*): + Date of creation of the repo on the Hub. Note that the lowest value is `2022-03-02T23:29:04.000Z`, + corresponding to the date when we began to store creation dates. + last_modified (`datetime`, *optional*): + Date of last commit to the repo. + private (`bool`): + Is the repo private. + disabled (`bool`, *optional*): + Is the repo disabled. + downloads (`int`): + Number of downloads of the model over the last 30 days. + downloads_all_time (`int`): + Cumulated number of downloads of the model since its creation. + gated (`Literal["auto", "manual", False]`, *optional*): + Is the repo gated. + If so, whether there is manual or automatic approval. + gguf (`Dict`, *optional*): + GGUF information of the model. + inference (`Literal["cold", "frozen", "warm"]`, *optional*): + Status of the model on the inference API. + Warm models are available for immediate use. Cold models will be loaded on first inference call. + Frozen models are not available in Inference API. + inference_provider_mapping (`Dict`, *optional*): + Model's inference provider mapping. + likes (`int`): + Number of likes of the model. + library_name (`str`, *optional*): + Library associated with the model. + tags (`List[str]`): + List of tags of the model. Compared to `card_data.tags`, contains extra tags computed by the Hub + (e.g. supported libraries, model's arXiv). + pipeline_tag (`str`, *optional*): + Pipeline tag associated with the model. + mask_token (`str`, *optional*): + Mask token used by the model. + widget_data (`Any`, *optional*): + Widget data associated with the model. + model_index (`Dict`, *optional*): + Model index for evaluation. + config (`Dict`, *optional*): + Model configuration. + transformers_info (`TransformersInfo`, *optional*): + Transformers-specific info (auto class, processor, etc.) associated with the model. + trending_score (`int`, *optional*): + Trending score of the model. + card_data (`ModelCardData`, *optional*): + Model Card Metadata as a [`huggingface_hub.repocard_data.ModelCardData`] object. + siblings (`List[RepoSibling]`): + List of [`huggingface_hub.hf_api.RepoSibling`] objects that constitute the model. + spaces (`List[str]`, *optional*): + List of spaces using the model. + safetensors (`SafeTensorsInfo`, *optional*): + Model's safetensors information. + security_repo_status (`Dict`, *optional*): + Model's security scan status. + """ + + id: str + author: Optional[str] + sha: Optional[str] + created_at: Optional[datetime] + last_modified: Optional[datetime] + private: Optional[bool] + disabled: Optional[bool] + downloads: Optional[int] + downloads_all_time: Optional[int] + gated: Optional[Literal["auto", "manual", False]] + gguf: Optional[Dict] + inference: Optional[Literal["warm", "cold", "frozen"]] + inference_provider_mapping: Optional[Dict[str, InferenceProviderMapping]] + likes: Optional[int] + library_name: Optional[str] + tags: Optional[List[str]] + pipeline_tag: Optional[str] + mask_token: Optional[str] + card_data: Optional[ModelCardData] + widget_data: Optional[Any] + model_index: Optional[Dict] + config: Optional[Dict] + transformers_info: Optional[TransformersInfo] + trending_score: Optional[int] + siblings: Optional[List[RepoSibling]] + spaces: Optional[List[str]] + safetensors: Optional[SafeTensorsInfo] + security_repo_status: Optional[Dict] + xet_enabled: Optional[bool] + + def __init__(self, **kwargs): + self.id = kwargs.pop("id") + self.author = kwargs.pop("author", None) + self.sha = kwargs.pop("sha", None) + last_modified = kwargs.pop("lastModified", None) or kwargs.pop("last_modified", None) + self.last_modified = parse_datetime(last_modified) if last_modified else None + created_at = kwargs.pop("createdAt", None) or kwargs.pop("created_at", None) + self.created_at = parse_datetime(created_at) if created_at else None + self.private = kwargs.pop("private", None) + self.gated = kwargs.pop("gated", None) + self.disabled = kwargs.pop("disabled", None) + self.downloads = kwargs.pop("downloads", None) + self.downloads_all_time = kwargs.pop("downloadsAllTime", None) + self.likes = kwargs.pop("likes", None) + self.library_name = kwargs.pop("library_name", None) + self.gguf = kwargs.pop("gguf", None) + + self.inference = kwargs.pop("inference", None) + self.inference_provider_mapping = kwargs.pop("inferenceProviderMapping", None) + if self.inference_provider_mapping: + self.inference_provider_mapping = { + provider: InferenceProviderMapping( + **{**value, "hf_model_id": self.id} + ) # little hack to simplify Inference Providers logic + for provider, value in self.inference_provider_mapping.items() + } + + self.tags = kwargs.pop("tags", None) + self.pipeline_tag = kwargs.pop("pipeline_tag", None) + self.mask_token = kwargs.pop("mask_token", None) + self.trending_score = kwargs.pop("trendingScore", None) + + card_data = kwargs.pop("cardData", None) or kwargs.pop("card_data", None) + self.card_data = ( + ModelCardData(**card_data, ignore_metadata_errors=True) if isinstance(card_data, dict) else card_data + ) + + self.widget_data = kwargs.pop("widgetData", None) + self.model_index = kwargs.pop("model-index", None) or kwargs.pop("model_index", None) + self.config = kwargs.pop("config", None) + transformers_info = kwargs.pop("transformersInfo", None) or kwargs.pop("transformers_info", None) + self.transformers_info = TransformersInfo(**transformers_info) if transformers_info else None + siblings = kwargs.pop("siblings", None) + self.siblings = ( + [ + RepoSibling( + rfilename=sibling["rfilename"], + size=sibling.get("size"), + blob_id=sibling.get("blobId"), + lfs=( + BlobLfsInfo( + size=sibling["lfs"]["size"], + sha256=sibling["lfs"]["sha256"], + pointer_size=sibling["lfs"]["pointerSize"], + ) + if sibling.get("lfs") + else None + ), + ) + for sibling in siblings + ] + if siblings is not None + else None + ) + self.spaces = kwargs.pop("spaces", None) + safetensors = kwargs.pop("safetensors", None) + self.safetensors = ( + SafeTensorsInfo( + parameters=safetensors["parameters"], + total=safetensors["total"], + ) + if safetensors + else None + ) + self.security_repo_status = kwargs.pop("securityRepoStatus", None) + self.xet_enabled = kwargs.pop("xetEnabled", None) + # backwards compatibility + self.lastModified = self.last_modified + self.cardData = self.card_data + self.transformersInfo = self.transformers_info + self.__dict__.update(**kwargs) + + +@dataclass +class DatasetInfo: + """ + Contains information about a dataset on the Hub. + + + + Most attributes of this class are optional. This is because the data returned by the Hub depends on the query made. + In general, the more specific the query, the more information is returned. On the contrary, when listing datasets + using [`list_datasets`] only a subset of the attributes are returned. + + + + Attributes: + id (`str`): + ID of dataset. + author (`str`): + Author of the dataset. + sha (`str`): + Repo SHA at this particular revision. + created_at (`datetime`, *optional*): + Date of creation of the repo on the Hub. Note that the lowest value is `2022-03-02T23:29:04.000Z`, + corresponding to the date when we began to store creation dates. + last_modified (`datetime`, *optional*): + Date of last commit to the repo. + private (`bool`): + Is the repo private. + disabled (`bool`, *optional*): + Is the repo disabled. + gated (`Literal["auto", "manual", False]`, *optional*): + Is the repo gated. + If so, whether there is manual or automatic approval. + downloads (`int`): + Number of downloads of the dataset over the last 30 days. + downloads_all_time (`int`): + Cumulated number of downloads of the model since its creation. + likes (`int`): + Number of likes of the dataset. + tags (`List[str]`): + List of tags of the dataset. + card_data (`DatasetCardData`, *optional*): + Model Card Metadata as a [`huggingface_hub.repocard_data.DatasetCardData`] object. + siblings (`List[RepoSibling]`): + List of [`huggingface_hub.hf_api.RepoSibling`] objects that constitute the dataset. + paperswithcode_id (`str`, *optional*): + Papers with code ID of the dataset. + trending_score (`int`, *optional*): + Trending score of the dataset. + """ + + id: str + author: Optional[str] + sha: Optional[str] + created_at: Optional[datetime] + last_modified: Optional[datetime] + private: Optional[bool] + gated: Optional[Literal["auto", "manual", False]] + disabled: Optional[bool] + downloads: Optional[int] + downloads_all_time: Optional[int] + likes: Optional[int] + paperswithcode_id: Optional[str] + tags: Optional[List[str]] + trending_score: Optional[int] + card_data: Optional[DatasetCardData] + siblings: Optional[List[RepoSibling]] + xet_enabled: Optional[bool] + + def __init__(self, **kwargs): + self.id = kwargs.pop("id") + self.author = kwargs.pop("author", None) + self.sha = kwargs.pop("sha", None) + created_at = kwargs.pop("createdAt", None) or kwargs.pop("created_at", None) + self.created_at = parse_datetime(created_at) if created_at else None + last_modified = kwargs.pop("lastModified", None) or kwargs.pop("last_modified", None) + self.last_modified = parse_datetime(last_modified) if last_modified else None + self.private = kwargs.pop("private", None) + self.gated = kwargs.pop("gated", None) + self.disabled = kwargs.pop("disabled", None) + self.downloads = kwargs.pop("downloads", None) + self.downloads_all_time = kwargs.pop("downloadsAllTime", None) + self.likes = kwargs.pop("likes", None) + self.paperswithcode_id = kwargs.pop("paperswithcode_id", None) + self.tags = kwargs.pop("tags", None) + self.trending_score = kwargs.pop("trendingScore", None) + + card_data = kwargs.pop("cardData", None) or kwargs.pop("card_data", None) + self.card_data = ( + DatasetCardData(**card_data, ignore_metadata_errors=True) if isinstance(card_data, dict) else card_data + ) + siblings = kwargs.pop("siblings", None) + self.siblings = ( + [ + RepoSibling( + rfilename=sibling["rfilename"], + size=sibling.get("size"), + blob_id=sibling.get("blobId"), + lfs=( + BlobLfsInfo( + size=sibling["lfs"]["size"], + sha256=sibling["lfs"]["sha256"], + pointer_size=sibling["lfs"]["pointerSize"], + ) + if sibling.get("lfs") + else None + ), + ) + for sibling in siblings + ] + if siblings is not None + else None + ) + self.xet_enabled = kwargs.pop("xetEnabled", None) + # backwards compatibility + self.lastModified = self.last_modified + self.cardData = self.card_data + self.__dict__.update(**kwargs) + + +@dataclass +class SpaceInfo: + """ + Contains information about a Space on the Hub. + + + + Most attributes of this class are optional. This is because the data returned by the Hub depends on the query made. + In general, the more specific the query, the more information is returned. On the contrary, when listing spaces + using [`list_spaces`] only a subset of the attributes are returned. + + + + Attributes: + id (`str`): + ID of the Space. + author (`str`, *optional*): + Author of the Space. + sha (`str`, *optional*): + Repo SHA at this particular revision. + created_at (`datetime`, *optional*): + Date of creation of the repo on the Hub. Note that the lowest value is `2022-03-02T23:29:04.000Z`, + corresponding to the date when we began to store creation dates. + last_modified (`datetime`, *optional*): + Date of last commit to the repo. + private (`bool`): + Is the repo private. + gated (`Literal["auto", "manual", False]`, *optional*): + Is the repo gated. + If so, whether there is manual or automatic approval. + disabled (`bool`, *optional*): + Is the Space disabled. + host (`str`, *optional*): + Host URL of the Space. + subdomain (`str`, *optional*): + Subdomain of the Space. + likes (`int`): + Number of likes of the Space. + tags (`List[str]`): + List of tags of the Space. + siblings (`List[RepoSibling]`): + List of [`huggingface_hub.hf_api.RepoSibling`] objects that constitute the Space. + card_data (`SpaceCardData`, *optional*): + Space Card Metadata as a [`huggingface_hub.repocard_data.SpaceCardData`] object. + runtime (`SpaceRuntime`, *optional*): + Space runtime information as a [`huggingface_hub.hf_api.SpaceRuntime`] object. + sdk (`str`, *optional*): + SDK used by the Space. + models (`List[str]`, *optional*): + List of models used by the Space. + datasets (`List[str]`, *optional*): + List of datasets used by the Space. + trending_score (`int`, *optional*): + Trending score of the Space. + """ + + id: str + author: Optional[str] + sha: Optional[str] + created_at: Optional[datetime] + last_modified: Optional[datetime] + private: Optional[bool] + gated: Optional[Literal["auto", "manual", False]] + disabled: Optional[bool] + host: Optional[str] + subdomain: Optional[str] + likes: Optional[int] + sdk: Optional[str] + tags: Optional[List[str]] + siblings: Optional[List[RepoSibling]] + trending_score: Optional[int] + card_data: Optional[SpaceCardData] + runtime: Optional[SpaceRuntime] + models: Optional[List[str]] + datasets: Optional[List[str]] + xet_enabled: Optional[bool] + + def __init__(self, **kwargs): + self.id = kwargs.pop("id") + self.author = kwargs.pop("author", None) + self.sha = kwargs.pop("sha", None) + created_at = kwargs.pop("createdAt", None) or kwargs.pop("created_at", None) + self.created_at = parse_datetime(created_at) if created_at else None + last_modified = kwargs.pop("lastModified", None) or kwargs.pop("last_modified", None) + self.last_modified = parse_datetime(last_modified) if last_modified else None + self.private = kwargs.pop("private", None) + self.gated = kwargs.pop("gated", None) + self.disabled = kwargs.pop("disabled", None) + self.host = kwargs.pop("host", None) + self.subdomain = kwargs.pop("subdomain", None) + self.likes = kwargs.pop("likes", None) + self.sdk = kwargs.pop("sdk", None) + self.tags = kwargs.pop("tags", None) + self.trending_score = kwargs.pop("trendingScore", None) + card_data = kwargs.pop("cardData", None) or kwargs.pop("card_data", None) + self.card_data = ( + SpaceCardData(**card_data, ignore_metadata_errors=True) if isinstance(card_data, dict) else card_data + ) + siblings = kwargs.pop("siblings", None) + self.siblings = ( + [ + RepoSibling( + rfilename=sibling["rfilename"], + size=sibling.get("size"), + blob_id=sibling.get("blobId"), + lfs=( + BlobLfsInfo( + size=sibling["lfs"]["size"], + sha256=sibling["lfs"]["sha256"], + pointer_size=sibling["lfs"]["pointerSize"], + ) + if sibling.get("lfs") + else None + ), + ) + for sibling in siblings + ] + if siblings is not None + else None + ) + runtime = kwargs.pop("runtime", None) + self.runtime = SpaceRuntime(runtime) if runtime else None + self.models = kwargs.pop("models", None) + self.datasets = kwargs.pop("datasets", None) + self.xet_enabled = kwargs.pop("xetEnabled", None) + # backwards compatibility + self.lastModified = self.last_modified + self.cardData = self.card_data + self.__dict__.update(**kwargs) + + +@dataclass +class CollectionItem: + """ + Contains information about an item of a Collection (model, dataset, Space or paper). + + Attributes: + item_object_id (`str`): + Unique ID of the item in the collection. + item_id (`str`): + ID of the underlying object on the Hub. Can be either a repo_id or a paper id + e.g. `"jbilcke-hf/ai-comic-factory"`, `"2307.09288"`. + item_type (`str`): + Type of the underlying object. Can be one of `"model"`, `"dataset"`, `"space"` or `"paper"`. + position (`int`): + Position of the item in the collection. + note (`str`, *optional*): + Note associated with the item, as plain text. + """ + + item_object_id: str # id in database + item_id: str # repo_id or paper id + item_type: str + position: int + note: Optional[str] = None + + def __init__( + self, _id: str, id: str, type: CollectionItemType_T, position: int, note: Optional[Dict] = None, **kwargs + ) -> None: + self.item_object_id: str = _id # id in database + self.item_id: str = id # repo_id or paper id + self.item_type: CollectionItemType_T = type + self.position: int = position + self.note: str = note["text"] if note is not None else None + + +@dataclass +class Collection: + """ + Contains information about a Collection on the Hub. + + Attributes: + slug (`str`): + Slug of the collection. E.g. `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + title (`str`): + Title of the collection. E.g. `"Recent models"`. + owner (`str`): + Owner of the collection. E.g. `"TheBloke"`. + items (`List[CollectionItem]`): + List of items in the collection. + last_updated (`datetime`): + Date of the last update of the collection. + position (`int`): + Position of the collection in the list of collections of the owner. + private (`bool`): + Whether the collection is private or not. + theme (`str`): + Theme of the collection. E.g. `"green"`. + upvotes (`int`): + Number of upvotes of the collection. + description (`str`, *optional*): + Description of the collection, as plain text. + url (`str`): + (property) URL of the collection on the Hub. + """ + + slug: str + title: str + owner: str + items: List[CollectionItem] + last_updated: datetime + position: int + private: bool + theme: str + upvotes: int + description: Optional[str] = None + + def __init__(self, **kwargs) -> None: + self.slug = kwargs.pop("slug") + self.title = kwargs.pop("title") + self.owner = kwargs.pop("owner") + self.items = [CollectionItem(**item) for item in kwargs.pop("items")] + self.last_updated = parse_datetime(kwargs.pop("lastUpdated")) + self.position = kwargs.pop("position") + self.private = kwargs.pop("private") + self.theme = kwargs.pop("theme") + self.upvotes = kwargs.pop("upvotes") + self.description = kwargs.pop("description", None) + endpoint = kwargs.pop("endpoint", None) + if endpoint is None: + endpoint = constants.ENDPOINT + self._url = f"{endpoint}/collections/{self.slug}" + + @property + def url(self) -> str: + """Returns the URL of the collection on the Hub.""" + return self._url + + +@dataclass +class GitRefInfo: + """ + Contains information about a git reference for a repo on the Hub. + + Attributes: + name (`str`): + Name of the reference (e.g. tag name or branch name). + ref (`str`): + Full git ref on the Hub (e.g. `"refs/heads/main"` or `"refs/tags/v1.0"`). + target_commit (`str`): + OID of the target commit for the ref (e.g. `"e7da7f221d5bf496a48136c0cd264e630fe9fcc8"`) + """ + + name: str + ref: str + target_commit: str + + +@dataclass +class GitRefs: + """ + Contains information about all git references for a repo on the Hub. + + Object is returned by [`list_repo_refs`]. + + Attributes: + branches (`List[GitRefInfo]`): + A list of [`GitRefInfo`] containing information about branches on the repo. + converts (`List[GitRefInfo]`): + A list of [`GitRefInfo`] containing information about "convert" refs on the repo. + Converts are refs used (internally) to push preprocessed data in Dataset repos. + tags (`List[GitRefInfo]`): + A list of [`GitRefInfo`] containing information about tags on the repo. + pull_requests (`List[GitRefInfo]`, *optional*): + A list of [`GitRefInfo`] containing information about pull requests on the repo. + Only returned if `include_prs=True` is set. + """ + + branches: List[GitRefInfo] + converts: List[GitRefInfo] + tags: List[GitRefInfo] + pull_requests: Optional[List[GitRefInfo]] = None + + +@dataclass +class GitCommitInfo: + """ + Contains information about a git commit for a repo on the Hub. Check out [`list_repo_commits`] for more details. + + Attributes: + commit_id (`str`): + OID of the commit (e.g. `"e7da7f221d5bf496a48136c0cd264e630fe9fcc8"`) + authors (`List[str]`): + List of authors of the commit. + created_at (`datetime`): + Datetime when the commit was created. + title (`str`): + Title of the commit. This is a free-text value entered by the authors. + message (`str`): + Description of the commit. This is a free-text value entered by the authors. + formatted_title (`str`): + Title of the commit formatted as HTML. Only returned if `formatted=True` is set. + formatted_message (`str`): + Description of the commit formatted as HTML. Only returned if `formatted=True` is set. + """ + + commit_id: str + + authors: List[str] + created_at: datetime + title: str + message: str + + formatted_title: Optional[str] + formatted_message: Optional[str] + + +@dataclass +class UserLikes: + """ + Contains information about a user likes on the Hub. + + Attributes: + user (`str`): + Name of the user for which we fetched the likes. + total (`int`): + Total number of likes. + datasets (`List[str]`): + List of datasets liked by the user (as repo_ids). + models (`List[str]`): + List of models liked by the user (as repo_ids). + spaces (`List[str]`): + List of spaces liked by the user (as repo_ids). + """ + + # Metadata + user: str + total: int + + # User likes + datasets: List[str] + models: List[str] + spaces: List[str] + + +@dataclass +class Organization: + """ + Contains information about an organization on the Hub. + + Attributes: + avatar_url (`str`): + URL of the organization's avatar. + name (`str`): + Name of the organization on the Hub (unique). + fullname (`str`): + Organization's full name. + """ + + avatar_url: str + name: str + fullname: str + + def __init__(self, **kwargs) -> None: + self.avatar_url = kwargs.pop("avatarUrl", "") + self.name = kwargs.pop("name", "") + self.fullname = kwargs.pop("fullname", "") + + # forward compatibility + self.__dict__.update(**kwargs) + + +@dataclass +class User: + """ + Contains information about a user on the Hub. + + Attributes: + username (`str`): + Name of the user on the Hub (unique). + fullname (`str`): + User's full name. + avatar_url (`str`): + URL of the user's avatar. + details (`str`, *optional*): + User's details. + is_following (`bool`, *optional*): + Whether the authenticated user is following this user. + is_pro (`bool`, *optional*): + Whether the user is a pro user. + num_models (`int`, *optional*): + Number of models created by the user. + num_datasets (`int`, *optional*): + Number of datasets created by the user. + num_spaces (`int`, *optional*): + Number of spaces created by the user. + num_discussions (`int`, *optional*): + Number of discussions initiated by the user. + num_papers (`int`, *optional*): + Number of papers authored by the user. + num_upvotes (`int`, *optional*): + Number of upvotes received by the user. + num_likes (`int`, *optional*): + Number of likes given by the user. + num_following (`int`, *optional*): + Number of users this user is following. + num_followers (`int`, *optional*): + Number of users following this user. + orgs (list of [`Organization`]): + List of organizations the user is part of. + """ + + # Metadata + username: str + fullname: str + avatar_url: str + details: Optional[str] = None + is_following: Optional[bool] = None + is_pro: Optional[bool] = None + num_models: Optional[int] = None + num_datasets: Optional[int] = None + num_spaces: Optional[int] = None + num_discussions: Optional[int] = None + num_papers: Optional[int] = None + num_upvotes: Optional[int] = None + num_likes: Optional[int] = None + num_following: Optional[int] = None + num_followers: Optional[int] = None + orgs: List[Organization] = field(default_factory=list) + + def __init__(self, **kwargs) -> None: + self.username = kwargs.pop("user", "") + self.fullname = kwargs.pop("fullname", "") + self.avatar_url = kwargs.pop("avatarUrl", "") + self.is_following = kwargs.pop("isFollowing", None) + self.is_pro = kwargs.pop("isPro", None) + self.details = kwargs.pop("details", None) + self.num_models = kwargs.pop("numModels", None) + self.num_datasets = kwargs.pop("numDatasets", None) + self.num_spaces = kwargs.pop("numSpaces", None) + self.num_discussions = kwargs.pop("numDiscussions", None) + self.num_papers = kwargs.pop("numPapers", None) + self.num_upvotes = kwargs.pop("numUpvotes", None) + self.num_likes = kwargs.pop("numLikes", None) + self.num_following = kwargs.pop("numFollowing", None) + self.num_followers = kwargs.pop("numFollowers", None) + self.user_type = kwargs.pop("type", None) + self.orgs = [Organization(**org) for org in kwargs.pop("orgs", [])] + + # forward compatibility + self.__dict__.update(**kwargs) + + +@dataclass +class PaperInfo: + """ + Contains information about a paper on the Hub. + + Attributes: + id (`str`): + arXiv paper ID. + authors (`List[str]`, **optional**): + Names of paper authors + published_at (`datetime`, **optional**): + Date paper published. + title (`str`, **optional**): + Title of the paper. + summary (`str`, **optional**): + Summary of the paper. + upvotes (`int`, **optional**): + Number of upvotes for the paper on the Hub. + discussion_id (`str`, **optional**): + Discussion ID for the paper on the Hub. + source (`str`, **optional**): + Source of the paper. + comments (`int`, **optional**): + Number of comments for the paper on the Hub. + submitted_at (`datetime`, **optional**): + Date paper appeared in daily papers on the Hub. + submitted_by (`User`, **optional**): + Information about who submitted the daily paper. + """ + + id: str + authors: Optional[List[str]] + published_at: Optional[datetime] + title: Optional[str] + summary: Optional[str] + upvotes: Optional[int] + discussion_id: Optional[str] + source: Optional[str] + comments: Optional[int] + submitted_at: Optional[datetime] + submitted_by: Optional[User] + + def __init__(self, **kwargs) -> None: + paper = kwargs.pop("paper", {}) + self.id = kwargs.pop("id", None) or paper.pop("id", None) + authors = paper.pop("authors", None) or kwargs.pop("authors", None) + self.authors = [author.pop("name", None) for author in authors] if authors else None + published_at = paper.pop("publishedAt", None) or kwargs.pop("publishedAt", None) + self.published_at = parse_datetime(published_at) if published_at else None + self.title = kwargs.pop("title", None) + self.source = kwargs.pop("source", None) + self.summary = paper.pop("summary", None) or kwargs.pop("summary", None) + self.upvotes = paper.pop("upvotes", None) or kwargs.pop("upvotes", None) + self.discussion_id = paper.pop("discussionId", None) or kwargs.pop("discussionId", None) + self.comments = kwargs.pop("numComments", 0) + submitted_at = kwargs.pop("publishedAt", None) or kwargs.pop("submittedOnDailyAt", None) + self.submitted_at = parse_datetime(submitted_at) if submitted_at else None + submitted_by = kwargs.pop("submittedBy", None) or kwargs.pop("submittedOnDailyBy", None) + self.submitted_by = User(**submitted_by) if submitted_by else None + + # forward compatibility + self.__dict__.update(**kwargs) + + +@dataclass +class LFSFileInfo: + """ + Contains information about a file stored as LFS on a repo on the Hub. + + Used in the context of listing and permanently deleting LFS files from a repo to free-up space. + See [`list_lfs_files`] and [`permanently_delete_lfs_files`] for more details. + + Git LFS files are tracked using SHA-256 object IDs, rather than file paths, to optimize performance + This approach is necessary because a single object can be referenced by multiple paths across different commits, + making it impractical to search and resolve these connections. Check out [our documentation](https://huggingface.co/docs/hub/storage-limits#advanced-track-lfs-file-references) + to learn how to know which filename(s) is(are) associated with each SHA. + + Attributes: + file_oid (`str`): + SHA-256 object ID of the file. This is the identifier to pass when permanently deleting the file. + filename (`str`): + Possible filename for the LFS object. See the note above for more information. + oid (`str`): + OID of the LFS object. + pushed_at (`datetime`): + Date the LFS object was pushed to the repo. + ref (`str`, *optional*): + Ref where the LFS object has been pushed (if any). + size (`int`): + Size of the LFS object. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> lfs_files = api.list_lfs_files("username/my-cool-repo") + + # Filter files files to delete based on a combination of `filename`, `pushed_at`, `ref` or `size`. + # e.g. select only LFS files in the "checkpoints" folder + >>> lfs_files_to_delete = (lfs_file for lfs_file in lfs_files if lfs_file.filename.startswith("checkpoints/")) + + # Permanently delete LFS files + >>> api.permanently_delete_lfs_files("username/my-cool-repo", lfs_files_to_delete) + ``` + """ + + file_oid: str + filename: str + oid: str + pushed_at: datetime + ref: Optional[str] + size: int + + def __init__(self, **kwargs) -> None: + self.file_oid = kwargs.pop("fileOid") + self.filename = kwargs.pop("filename") + self.oid = kwargs.pop("oid") + self.pushed_at = parse_datetime(kwargs.pop("pushedAt")) + self.ref = kwargs.pop("ref", None) + self.size = kwargs.pop("size") + + # forward compatibility + self.__dict__.update(**kwargs) + + +def future_compatible(fn: CallableT) -> CallableT: + """Wrap a method of `HfApi` to handle `run_as_future=True`. + + A method flagged as "future_compatible" will be called in a thread if `run_as_future=True` and return a + `concurrent.futures.Future` instance. Otherwise, it will be called normally and return the result. + """ + sig = inspect.signature(fn) + args_params = list(sig.parameters)[1:] # remove "self" from list + + @wraps(fn) + def _inner(self, *args, **kwargs): + # Get `run_as_future` value if provided (default to False) + if "run_as_future" in kwargs: + run_as_future = kwargs["run_as_future"] + kwargs["run_as_future"] = False # avoid recursion error + else: + run_as_future = False + for param, value in zip(args_params, args): + if param == "run_as_future": + run_as_future = value + break + + # Call the function in a thread if `run_as_future=True` + if run_as_future: + return self.run_as_future(fn, self, *args, **kwargs) + + # Otherwise, call the function normally + return fn(self, *args, **kwargs) + + _inner.is_future_compatible = True # type: ignore + return _inner # type: ignore + + +class HfApi: + """ + Client to interact with the Hugging Face Hub via HTTP. + + The client is initialized with some high-level settings used in all requests + made to the Hub (HF endpoint, authentication, user agents...). Using the `HfApi` + client is preferred but not mandatory as all of its public methods are exposed + directly at the root of `huggingface_hub`. + + Args: + endpoint (`str`, *optional*): + Endpoint of the Hub. Defaults to . + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + library_name (`str`, *optional*): + The name of the library that is making the HTTP request. Will be added to + the user-agent header. Example: `"transformers"`. + library_version (`str`, *optional*): + The version of the library that is making the HTTP request. Will be added + to the user-agent header. Example: `"4.24.0"`. + user_agent (`str`, `dict`, *optional*): + The user agent info in the form of a dictionary or a single string. It will + be completed with information about the installed packages. + headers (`dict`, *optional*): + Additional headers to be sent with each request. Example: `{"X-My-Header": "value"}`. + Headers passed here are taking precedence over the default headers. + """ + + def __init__( + self, + endpoint: Optional[str] = None, + token: Union[str, bool, None] = None, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, + headers: Optional[Dict[str, str]] = None, + ) -> None: + self.endpoint = endpoint if endpoint is not None else constants.ENDPOINT + self.token = token + self.library_name = library_name + self.library_version = library_version + self.user_agent = user_agent + self.headers = headers + self._thread_pool: Optional[ThreadPoolExecutor] = None + + def run_as_future(self, fn: Callable[..., R], *args, **kwargs) -> Future[R]: + """ + Run a method in the background and return a Future instance. + + The main goal is to run methods without blocking the main thread (e.g. to push data during a training). + Background jobs are queued to preserve order but are not ran in parallel. If you need to speed-up your scripts + by parallelizing lots of call to the API, you must setup and use your own [ThreadPoolExecutor](https://docs.python.org/3/library/concurrent.futures.html#threadpoolexecutor). + + Note: Most-used methods like [`upload_file`], [`upload_folder`] and [`create_commit`] have a `run_as_future: bool` + argument to directly call them in the background. This is equivalent to calling `api.run_as_future(...)` on them + but less verbose. + + Args: + fn (`Callable`): + The method to run in the background. + *args, **kwargs: + Arguments with which the method will be called. + + Return: + `Future`: a [Future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) instance to + get the result of the task. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> future = api.run_as_future(api.whoami) # instant + >>> future.done() + False + >>> future.result() # wait until complete and return result + (...) + >>> future.done() + True + ``` + """ + if self._thread_pool is None: + self._thread_pool = ThreadPoolExecutor(max_workers=1) + self._thread_pool + return self._thread_pool.submit(fn, *args, **kwargs) + + @validate_hf_hub_args + def whoami(self, token: Union[bool, str, None] = None) -> Dict: + """ + Call HF API to know "whoami". + + Args: + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + # Get the effective token using the helper function get_token + effective_token = token or self.token or get_token() or True + r = get_session().get( + f"{self.endpoint}/api/whoami-v2", + headers=self._build_hf_headers(token=effective_token), + ) + try: + hf_raise_for_status(r) + except HTTPError as e: + error_message = "Invalid user token." + # Check which token is the effective one and generate the error message accordingly + if effective_token == _get_token_from_google_colab(): + error_message += " The token from Google Colab vault is invalid. Please update it from the UI." + elif effective_token == _get_token_from_environment(): + error_message += ( + " The token from HF_TOKEN environment variable is invalid. " + "Note that HF_TOKEN takes precedence over `huggingface-cli login`." + ) + elif effective_token == _get_token_from_file(): + error_message += " The token stored is invalid. Please run `huggingface-cli login` to update it." + raise HTTPError(error_message, request=e.request, response=e.response) from e + return r.json() + + @_deprecate_method( + version="1.0", + message=( + "Permissions are more complex than when `get_token_permission` was first introduced. " + "OAuth and fine-grain tokens allows for more detailed permissions. " + "If you need to know the permissions associated with a token, please use `whoami` and check the `'auth'` key." + ), + ) + def get_token_permission( + self, token: Union[bool, str, None] = None + ) -> Literal["read", "write", "fineGrained", None]: + """ + Check if a given `token` is valid and return its permissions. + + + + This method is deprecated and will be removed in version 1.0. Permissions are more complex than when + `get_token_permission` was first introduced. OAuth and fine-grain tokens allows for more detailed permissions. + If you need to know the permissions associated with a token, please use `whoami` and check the `'auth'` key. + + + + For more details about tokens, please refer to https://huggingface.co/docs/hub/security-tokens#what-are-user-access-tokens. + + Args: + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Literal["read", "write", "fineGrained", None]`: Permission granted by the token ("read" or "write"). Returns `None` if no + token passed, if token is invalid or if role is not returned by the server. This typically happens when the token is an OAuth token. + """ + try: + return self.whoami(token=token)["auth"]["accessToken"]["role"] + except (LocalTokenNotFoundError, HTTPError, KeyError): + return None + + def get_model_tags(self) -> Dict: + """ + List all valid model tags as a nested namespace object + """ + path = f"{self.endpoint}/api/models-tags-by-type" + r = get_session().get(path) + hf_raise_for_status(r) + return r.json() + + def get_dataset_tags(self) -> Dict: + """ + List all valid dataset tags as a nested namespace object. + """ + path = f"{self.endpoint}/api/datasets-tags-by-type" + r = get_session().get(path) + hf_raise_for_status(r) + return r.json() + + @validate_hf_hub_args + def list_models( + self, + *, + # Search-query parameter + filter: Union[str, Iterable[str], None] = None, + author: Optional[str] = None, + gated: Optional[bool] = None, + inference: Optional[Literal["cold", "frozen", "warm"]] = None, + library: Optional[Union[str, List[str]]] = None, + language: Optional[Union[str, List[str]]] = None, + model_name: Optional[str] = None, + task: Optional[Union[str, List[str]]] = None, + trained_dataset: Optional[Union[str, List[str]]] = None, + tags: Optional[Union[str, List[str]]] = None, + search: Optional[str] = None, + pipeline_tag: Optional[str] = None, + emissions_thresholds: Optional[Tuple[float, float]] = None, + # Sorting and pagination parameters + sort: Union[Literal["last_modified"], str, None] = None, + direction: Optional[Literal[-1]] = None, + limit: Optional[int] = None, + # Additional data to fetch + expand: Optional[List[ExpandModelProperty_T]] = None, + full: Optional[bool] = None, + cardData: bool = False, + fetch_config: bool = False, + token: Union[bool, str, None] = None, + ) -> Iterable[ModelInfo]: + """ + List models hosted on the Huggingface Hub, given some filters. + + Args: + filter (`str` or `Iterable[str]`, *optional*): + A string or list of string to filter models on the Hub. + author (`str`, *optional*): + A string which identify the author (user or organization) of the + returned models. + gated (`bool`, *optional*): + A boolean to filter models on the Hub that are gated or not. By default, all models are returned. + If `gated=True` is passed, only gated models are returned. + If `gated=False` is passed, only non-gated models are returned. + inference (`Literal["cold", "frozen", "warm"]`, *optional*): + A string to filter models on the Hub by their state on the Inference API. + Warm models are available for immediate use. Cold models will be loaded on first inference call. + Frozen models are not available in Inference API. + library (`str` or `List`, *optional*): + A string or list of strings of foundational libraries models were + originally trained from, such as pytorch, tensorflow, or allennlp. + language (`str` or `List`, *optional*): + A string or list of strings of languages, both by name and country + code, such as "en" or "English" + model_name (`str`, *optional*): + A string that contain complete or partial names for models on the + Hub, such as "bert" or "bert-base-cased" + task (`str` or `List`, *optional*): + A string or list of strings of tasks models were designed for, such + as: "fill-mask" or "automatic-speech-recognition" + trained_dataset (`str` or `List`, *optional*): + A string tag or a list of string tags of the trained dataset for a + model on the Hub. + tags (`str` or `List`, *optional*): + A string tag or a list of tags to filter models on the Hub by, such + as `text-generation` or `spacy`. + search (`str`, *optional*): + A string that will be contained in the returned model ids. + pipeline_tag (`str`, *optional*): + A string pipeline tag to filter models on the Hub by, such as `summarization`. + emissions_thresholds (`Tuple`, *optional*): + A tuple of two ints or floats representing a minimum and maximum + carbon footprint to filter the resulting models with in grams. + sort (`Literal["last_modified"]` or `str`, *optional*): + The key with which to sort the resulting models. Possible values are "last_modified", "trending_score", + "created_at", "downloads" and "likes". + direction (`Literal[-1]` or `int`, *optional*): + Direction in which to sort. The value `-1` sorts by descending + order while all other values sort by ascending order. + limit (`int`, *optional*): + The limit on the number of models fetched. Leaving this option + to `None` fetches all models. + expand (`List[ExpandModelProperty_T]`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `full`, `cardData` or `fetch_config` are passed. + Possible values are `"author"`, `"baseModels"`, `"cardData"`, `"childrenModelCount"`, `"config"`, `"createdAt"`, `"disabled"`, `"downloads"`, `"downloadsAllTime"`, `"gated"`, `"gguf"`, `"inference"`, `"inferenceProviderMapping"`, `"lastModified"`, `"library_name"`, `"likes"`, `"mask_token"`, `"model-index"`, `"pipeline_tag"`, `"private"`, `"safetensors"`, `"sha"`, `"siblings"`, `"spaces"`, `"tags"`, `"transformersInfo"`, `"trendingScore"`, `"widgetData"`, `"usedStorage"`, `"resourceGroup"` and `"xetEnabled"`. + full (`bool`, *optional*): + Whether to fetch all model data, including the `last_modified`, + the `sha`, the files and the `tags`. This is set to `True` by + default when using a filter. + cardData (`bool`, *optional*): + Whether to grab the metadata for the model as well. Can contain + useful information such as carbon emissions, metrics, and + datasets trained on. + fetch_config (`bool`, *optional*): + Whether to fetch the model configs as well. This is not included + in `full` due to its size. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + + Returns: + `Iterable[ModelInfo]`: an iterable of [`huggingface_hub.hf_api.ModelInfo`] objects. + + Example usage with the `filter` argument: + + ```python + >>> from huggingface_hub import HfApi + + >>> api = HfApi() + + # List all models + >>> api.list_models() + + # List only the text classification models + >>> api.list_models(filter="text-classification") + + # List only models from the AllenNLP library + >>> api.list_models(filter="allennlp") + ``` + + Example usage with the `search` argument: + + ```python + >>> from huggingface_hub import HfApi + + >>> api = HfApi() + + # List all models with "bert" in their name + >>> api.list_models(search="bert") + + # List all models with "bert" in their name made by google + >>> api.list_models(search="bert", author="google") + ``` + """ + if expand and (full or cardData or fetch_config): + raise ValueError("`expand` cannot be used if `full`, `cardData` or `fetch_config` are passed.") + + if emissions_thresholds is not None and cardData is None: + raise ValueError("`emissions_thresholds` were passed without setting `cardData=True`.") + + path = f"{self.endpoint}/api/models" + headers = self._build_hf_headers(token=token) + params: Dict[str, Any] = {} + + # Build the filter list + filter_list: List[str] = [] + if filter: + filter_list.extend([filter] if isinstance(filter, str) else filter) + if library: + filter_list.extend([library] if isinstance(library, str) else library) + if task: + filter_list.extend([task] if isinstance(task, str) else task) + if trained_dataset: + if isinstance(trained_dataset, str): + trained_dataset = [trained_dataset] + for dataset in trained_dataset: + if not dataset.startswith("dataset:"): + dataset = f"dataset:{dataset}" + filter_list.append(dataset) + if language: + filter_list.extend([language] if isinstance(language, str) else language) + if tags: + filter_list.extend([tags] if isinstance(tags, str) else tags) + if len(filter_list) > 0: + params["filter"] = filter_list + + # Handle other query params + if author: + params["author"] = author + if gated is not None: + params["gated"] = gated + if inference is not None: + params["inference"] = inference + if pipeline_tag: + params["pipeline_tag"] = pipeline_tag + search_list = [] + if model_name: + search_list.append(model_name) + if search: + search_list.append(search) + if len(search_list) > 0: + params["search"] = search_list + if sort is not None: + params["sort"] = ( + "lastModified" + if sort == "last_modified" + else "trendingScore" + if sort == "trending_score" + else "createdAt" + if sort == "created_at" + else sort + ) + if direction is not None: + params["direction"] = direction + if limit is not None: + params["limit"] = limit + + # Request additional data + if full: + params["full"] = True + if fetch_config: + params["config"] = True + if cardData: + params["cardData"] = True + if expand: + params["expand"] = expand + + # `items` is a generator + items = paginate(path, params=params, headers=headers) + if limit is not None: + items = islice(items, limit) # Do not iterate over all pages + for item in items: + if "siblings" not in item: + item["siblings"] = None + model_info = ModelInfo(**item) + if emissions_thresholds is None or _is_emission_within_threshold(model_info, *emissions_thresholds): + yield model_info + + @validate_hf_hub_args + def list_datasets( + self, + *, + # Search-query parameter + filter: Union[str, Iterable[str], None] = None, + author: Optional[str] = None, + benchmark: Optional[Union[str, List[str]]] = None, + dataset_name: Optional[str] = None, + gated: Optional[bool] = None, + language_creators: Optional[Union[str, List[str]]] = None, + language: Optional[Union[str, List[str]]] = None, + multilinguality: Optional[Union[str, List[str]]] = None, + size_categories: Optional[Union[str, List[str]]] = None, + tags: Optional[Union[str, List[str]]] = None, + task_categories: Optional[Union[str, List[str]]] = None, + task_ids: Optional[Union[str, List[str]]] = None, + search: Optional[str] = None, + # Sorting and pagination parameters + sort: Optional[Union[Literal["last_modified"], str]] = None, + direction: Optional[Literal[-1]] = None, + limit: Optional[int] = None, + # Additional data to fetch + expand: Optional[List[ExpandDatasetProperty_T]] = None, + full: Optional[bool] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[DatasetInfo]: + """ + List datasets hosted on the Huggingface Hub, given some filters. + + Args: + filter (`str` or `Iterable[str]`, *optional*): + A string or list of string to filter datasets on the hub. + author (`str`, *optional*): + A string which identify the author of the returned datasets. + benchmark (`str` or `List`, *optional*): + A string or list of strings that can be used to identify datasets on + the Hub by their official benchmark. + dataset_name (`str`, *optional*): + A string or list of strings that can be used to identify datasets on + the Hub by its name, such as `SQAC` or `wikineural` + gated (`bool`, *optional*): + A boolean to filter datasets on the Hub that are gated or not. By default, all datasets are returned. + If `gated=True` is passed, only gated datasets are returned. + If `gated=False` is passed, only non-gated datasets are returned. + language_creators (`str` or `List`, *optional*): + A string or list of strings that can be used to identify datasets on + the Hub with how the data was curated, such as `crowdsourced` or + `machine_generated`. + language (`str` or `List`, *optional*): + A string or list of strings representing a two-character language to + filter datasets by on the Hub. + multilinguality (`str` or `List`, *optional*): + A string or list of strings representing a filter for datasets that + contain multiple languages. + size_categories (`str` or `List`, *optional*): + A string or list of strings that can be used to identify datasets on + the Hub by the size of the dataset such as `100K>> from huggingface_hub import HfApi + + >>> api = HfApi() + + # List all datasets + >>> api.list_datasets() + + + # List only the text classification datasets + >>> api.list_datasets(filter="task_categories:text-classification") + + + # List only the datasets in russian for language modeling + >>> api.list_datasets( + ... filter=("language:ru", "task_ids:language-modeling") + ... ) + + # List FiftyOne datasets (identified by the tag "fiftyone" in dataset card) + >>> api.list_datasets(tags="fiftyone") + ``` + + Example usage with the `search` argument: + + ```python + >>> from huggingface_hub import HfApi + + >>> api = HfApi() + + # List all datasets with "text" in their name + >>> api.list_datasets(search="text") + + # List all datasets with "text" in their name made by google + >>> api.list_datasets(search="text", author="google") + ``` + """ + if expand and full: + raise ValueError("`expand` cannot be used if `full` is passed.") + + path = f"{self.endpoint}/api/datasets" + headers = self._build_hf_headers(token=token) + params: Dict[str, Any] = {} + + # Build `filter` list + filter_list = [] + if filter is not None: + if isinstance(filter, str): + filter_list.append(filter) + else: + filter_list.extend(filter) + for key, value in ( + ("benchmark", benchmark), + ("language_creators", language_creators), + ("language", language), + ("multilinguality", multilinguality), + ("size_categories", size_categories), + ("task_categories", task_categories), + ("task_ids", task_ids), + ): + if value: + if isinstance(value, str): + value = [value] + for value_item in value: + if not value_item.startswith(f"{key}:"): + data = f"{key}:{value_item}" + filter_list.append(data) + if tags is not None: + filter_list.extend([tags] if isinstance(tags, str) else tags) + if len(filter_list) > 0: + params["filter"] = filter_list + + # Handle other query params + if author: + params["author"] = author + if gated is not None: + params["gated"] = gated + search_list = [] + if dataset_name: + search_list.append(dataset_name) + if search: + search_list.append(search) + if len(search_list) > 0: + params["search"] = search_list + if sort is not None: + params["sort"] = ( + "lastModified" + if sort == "last_modified" + else "trendingScore" + if sort == "trending_score" + else "createdAt" + if sort == "created_at" + else sort + ) + if direction is not None: + params["direction"] = direction + if limit is not None: + params["limit"] = limit + + # Request additional data + if expand: + params["expand"] = expand + if full: + params["full"] = True + + items = paginate(path, params=params, headers=headers) + if limit is not None: + items = islice(items, limit) # Do not iterate over all pages + for item in items: + if "siblings" not in item: + item["siblings"] = None + yield DatasetInfo(**item) + + @validate_hf_hub_args + def list_spaces( + self, + *, + # Search-query parameter + filter: Union[str, Iterable[str], None] = None, + author: Optional[str] = None, + search: Optional[str] = None, + datasets: Union[str, Iterable[str], None] = None, + models: Union[str, Iterable[str], None] = None, + linked: bool = False, + # Sorting and pagination parameters + sort: Union[Literal["last_modified"], str, None] = None, + direction: Optional[Literal[-1]] = None, + limit: Optional[int] = None, + # Additional data to fetch + expand: Optional[List[ExpandSpaceProperty_T]] = None, + full: Optional[bool] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[SpaceInfo]: + """ + List spaces hosted on the Huggingface Hub, given some filters. + + Args: + filter (`str` or `Iterable`, *optional*): + A string tag or list of tags that can be used to identify Spaces on the Hub. + author (`str`, *optional*): + A string which identify the author of the returned Spaces. + search (`str`, *optional*): + A string that will be contained in the returned Spaces. + datasets (`str` or `Iterable`, *optional*): + Whether to return Spaces that make use of a dataset. + The name of a specific dataset can be passed as a string. + models (`str` or `Iterable`, *optional*): + Whether to return Spaces that make use of a model. + The name of a specific model can be passed as a string. + linked (`bool`, *optional*): + Whether to return Spaces that make use of either a model or a dataset. + sort (`Literal["last_modified"]` or `str`, *optional*): + The key with which to sort the resulting models. Possible values are "last_modified", "trending_score", + "created_at" and "likes". + direction (`Literal[-1]` or `int`, *optional*): + Direction in which to sort. The value `-1` sorts by descending + order while all other values sort by ascending order. + limit (`int`, *optional*): + The limit on the number of Spaces fetched. Leaving this option + to `None` fetches all Spaces. + expand (`List[ExpandSpaceProperty_T]`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `full` is passed. + Possible values are `"author"`, `"cardData"`, `"datasets"`, `"disabled"`, `"lastModified"`, `"createdAt"`, `"likes"`, `"models"`, `"private"`, `"runtime"`, `"sdk"`, `"siblings"`, `"sha"`, `"subdomain"`, `"tags"`, `"trendingScore"`, `"usedStorage"`, `"resourceGroup"` and `"xetEnabled"`. + full (`bool`, *optional*): + Whether to fetch all Spaces data, including the `last_modified`, `siblings` + and `card_data` fields. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[SpaceInfo]`: an iterable of [`huggingface_hub.hf_api.SpaceInfo`] objects. + """ + if expand and full: + raise ValueError("`expand` cannot be used if `full` is passed.") + + path = f"{self.endpoint}/api/spaces" + headers = self._build_hf_headers(token=token) + params: Dict[str, Any] = {} + if filter is not None: + params["filter"] = filter + if author is not None: + params["author"] = author + if search is not None: + params["search"] = search + if sort is not None: + params["sort"] = ( + "lastModified" + if sort == "last_modified" + else "trendingScore" + if sort == "trending_score" + else "createdAt" + if sort == "created_at" + else sort + ) + if direction is not None: + params["direction"] = direction + if limit is not None: + params["limit"] = limit + if linked: + params["linked"] = True + if datasets is not None: + params["datasets"] = datasets + if models is not None: + params["models"] = models + + # Request additional data + if expand: + params["expand"] = expand + if full: + params["full"] = True + + items = paginate(path, params=params, headers=headers) + if limit is not None: + items = islice(items, limit) # Do not iterate over all pages + for item in items: + if "siblings" not in item: + item["siblings"] = None + yield SpaceInfo(**item) + + @validate_hf_hub_args + def unlike( + self, + repo_id: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> None: + """ + Unlike a given repo on the Hub (e.g. remove from favorite list). + + To prevent spam usage, it is not possible to `like` a repository from a script. + + See also [`list_liked_repos`]. + + Args: + repo_id (`str`): + The repository to unlike. Example: `"user/my-cool-model"`. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if unliking a dataset or space, `None` or + `"model"` if unliking a model. Default is `None`. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + + Example: + ```python + >>> from huggingface_hub import list_liked_repos, unlike + >>> "gpt2" in list_liked_repos().models # we assume you have already liked gpt2 + True + >>> unlike("gpt2") + >>> "gpt2" in list_liked_repos().models + False + ``` + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + response = get_session().delete( + url=f"{self.endpoint}/api/{repo_type}s/{repo_id}/like", headers=self._build_hf_headers(token=token) + ) + hf_raise_for_status(response) + + @validate_hf_hub_args + def list_liked_repos( + self, + user: Optional[str] = None, + *, + token: Union[bool, str, None] = None, + ) -> UserLikes: + """ + List all public repos liked by a user on huggingface.co. + + This list is public so token is optional. If `user` is not passed, it defaults to + the logged in user. + + See also [`unlike`]. + + Args: + user (`str`, *optional*): + Name of the user for which you want to fetch the likes. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`UserLikes`]: object containing the user name and 3 lists of repo ids (1 for + models, 1 for datasets and 1 for Spaces). + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `user` is not passed and no token found (either from argument or from machine). + + Example: + ```python + >>> from huggingface_hub import list_liked_repos + + >>> likes = list_liked_repos("julien-c") + + >>> likes.user + "julien-c" + + >>> likes.models + ["osanseviero/streamlit_1.15", "Xhaheen/ChatGPT_HF", ...] + ``` + """ + # User is either provided explicitly or retrieved from current token. + if user is None: + me = self.whoami(token=token) + if me["type"] == "user": + user = me["name"] + else: + raise ValueError( + "Cannot list liked repos. You must provide a 'user' as input or be logged in as a user." + ) + + path = f"{self.endpoint}/api/users/{user}/likes" + headers = self._build_hf_headers(token=token) + + likes = list(paginate(path, params={}, headers=headers)) + # Looping over a list of items similar to: + # { + # 'createdAt': '2021-09-09T21:53:27.000Z', + # 'repo': { + # 'name': 'PaddlePaddle/PaddleOCR', + # 'type': 'space' + # } + # } + # Let's loop 3 times over the received list. Less efficient but more straightforward to read. + return UserLikes( + user=user, + total=len(likes), + models=[like["repo"]["name"] for like in likes if like["repo"]["type"] == "model"], + datasets=[like["repo"]["name"] for like in likes if like["repo"]["type"] == "dataset"], + spaces=[like["repo"]["name"] for like in likes if like["repo"]["type"] == "space"], + ) + + @validate_hf_hub_args + def list_repo_likers( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[User]: + """ + List all users who liked a given repo on the hugging Face Hub. + + See also [`list_liked_repos`]. + + Args: + repo_id (`str`): + The repository to retrieve . Example: `"user/my-cool-model"`. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + + Returns: + `Iterable[User]`: an iterable of [`huggingface_hub.hf_api.User`] objects. + """ + + # Construct the API endpoint + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + path = f"{self.endpoint}/api/{repo_type}s/{repo_id}/likers" + for liker in paginate(path, params={}, headers=self._build_hf_headers(token=token)): + yield User(username=liker["user"], fullname=liker["fullname"], avatar_url=liker["avatarUrl"]) + + @validate_hf_hub_args + def model_info( + self, + repo_id: str, + *, + revision: Optional[str] = None, + timeout: Optional[float] = None, + securityStatus: Optional[bool] = None, + files_metadata: bool = False, + expand: Optional[List[ExpandModelProperty_T]] = None, + token: Union[bool, str, None] = None, + ) -> ModelInfo: + """ + Get info on one specific model on huggingface.co + + Model can be private if you pass an acceptable token or are logged in. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`, *optional*): + The revision of the model repository from which to get the + information. + timeout (`float`, *optional*): + Whether to set a timeout for the request to the Hub. + securityStatus (`bool`, *optional*): + Whether to retrieve the security status from the model + repository as well. The security status will be returned in the `security_repo_status` field. + files_metadata (`bool`, *optional*): + Whether or not to retrieve metadata for files in the repository + (size, LFS metadata, etc). Defaults to `False`. + expand (`List[ExpandModelProperty_T]`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `securityStatus` or `files_metadata` are passed. + Possible values are `"author"`, `"baseModels"`, `"cardData"`, `"childrenModelCount"`, `"config"`, `"createdAt"`, `"disabled"`, `"downloads"`, `"downloadsAllTime"`, `"gated"`, `"gguf"`, `"inference"`, `"inferenceProviderMapping"`, `"lastModified"`, `"library_name"`, `"likes"`, `"mask_token"`, `"model-index"`, `"pipeline_tag"`, `"private"`, `"safetensors"`, `"sha"`, `"siblings"`, `"spaces"`, `"tags"`, `"transformersInfo"`, `"trendingScore"`, `"widgetData"`, `"usedStorage"`, `"resourceGroup"` and `"xetEnabled"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`huggingface_hub.hf_api.ModelInfo`]: The model repository information. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + + + """ + if expand and (securityStatus or files_metadata): + raise ValueError("`expand` cannot be used if `securityStatus` or `files_metadata` are set.") + + headers = self._build_hf_headers(token=token) + path = ( + f"{self.endpoint}/api/models/{repo_id}" + if revision is None + else (f"{self.endpoint}/api/models/{repo_id}/revision/{quote(revision, safe='')}") + ) + params: Dict = {} + if securityStatus: + params["securityStatus"] = True + if files_metadata: + params["blobs"] = True + if expand: + params["expand"] = expand + r = get_session().get(path, headers=headers, timeout=timeout, params=params) + hf_raise_for_status(r) + data = r.json() + return ModelInfo(**data) + + @validate_hf_hub_args + def dataset_info( + self, + repo_id: str, + *, + revision: Optional[str] = None, + timeout: Optional[float] = None, + files_metadata: bool = False, + expand: Optional[List[ExpandDatasetProperty_T]] = None, + token: Union[bool, str, None] = None, + ) -> DatasetInfo: + """ + Get info on one specific dataset on huggingface.co. + + Dataset can be private if you pass an acceptable token. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`, *optional*): + The revision of the dataset repository from which to get the + information. + timeout (`float`, *optional*): + Whether to set a timeout for the request to the Hub. + files_metadata (`bool`, *optional*): + Whether or not to retrieve metadata for files in the repository + (size, LFS metadata, etc). Defaults to `False`. + expand (`List[ExpandDatasetProperty_T]`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `files_metadata` is passed. + Possible values are `"author"`, `"cardData"`, `"citation"`, `"createdAt"`, `"disabled"`, `"description"`, `"downloads"`, `"downloadsAllTime"`, `"gated"`, `"lastModified"`, `"likes"`, `"paperswithcode_id"`, `"private"`, `"siblings"`, `"sha"`, `"tags"`, `"trendingScore"`,`"usedStorage"`, `"resourceGroup"` and `"xetEnabled"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`hf_api.DatasetInfo`]: The dataset repository information. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + + + """ + if expand and files_metadata: + raise ValueError("`expand` cannot be used if `files_metadata` is set.") + + headers = self._build_hf_headers(token=token) + path = ( + f"{self.endpoint}/api/datasets/{repo_id}" + if revision is None + else (f"{self.endpoint}/api/datasets/{repo_id}/revision/{quote(revision, safe='')}") + ) + params: Dict = {} + if files_metadata: + params["blobs"] = True + if expand: + params["expand"] = expand + + r = get_session().get(path, headers=headers, timeout=timeout, params=params) + hf_raise_for_status(r) + data = r.json() + return DatasetInfo(**data) + + @validate_hf_hub_args + def space_info( + self, + repo_id: str, + *, + revision: Optional[str] = None, + timeout: Optional[float] = None, + files_metadata: bool = False, + expand: Optional[List[ExpandSpaceProperty_T]] = None, + token: Union[bool, str, None] = None, + ) -> SpaceInfo: + """ + Get info on one specific Space on huggingface.co. + + Space can be private if you pass an acceptable token. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`, *optional*): + The revision of the space repository from which to get the + information. + timeout (`float`, *optional*): + Whether to set a timeout for the request to the Hub. + files_metadata (`bool`, *optional*): + Whether or not to retrieve metadata for files in the repository + (size, LFS metadata, etc). Defaults to `False`. + expand (`List[ExpandSpaceProperty_T]`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `full` is passed. + Possible values are `"author"`, `"cardData"`, `"createdAt"`, `"datasets"`, `"disabled"`, `"lastModified"`, `"likes"`, `"models"`, `"private"`, `"runtime"`, `"sdk"`, `"siblings"`, `"sha"`, `"subdomain"`, `"tags"`, `"trendingScore"`, `"usedStorage"`, `"resourceGroup"` and `"xetEnabled"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`~hf_api.SpaceInfo`]: The space repository information. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + + + """ + if expand and files_metadata: + raise ValueError("`expand` cannot be used if `files_metadata` is set.") + + headers = self._build_hf_headers(token=token) + path = ( + f"{self.endpoint}/api/spaces/{repo_id}" + if revision is None + else (f"{self.endpoint}/api/spaces/{repo_id}/revision/{quote(revision, safe='')}") + ) + params: Dict = {} + if files_metadata: + params["blobs"] = True + if expand: + params["expand"] = expand + + r = get_session().get(path, headers=headers, timeout=timeout, params=params) + hf_raise_for_status(r) + data = r.json() + return SpaceInfo(**data) + + @validate_hf_hub_args + def repo_info( + self, + repo_id: str, + *, + revision: Optional[str] = None, + repo_type: Optional[str] = None, + timeout: Optional[float] = None, + files_metadata: bool = False, + expand: Optional[Union[ExpandModelProperty_T, ExpandDatasetProperty_T, ExpandSpaceProperty_T]] = None, + token: Union[bool, str, None] = None, + ) -> Union[ModelInfo, DatasetInfo, SpaceInfo]: + """ + Get the info object for a given repo of a given type. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`, *optional*): + The revision of the repository from which to get the + information. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if getting repository info from a dataset or a space, + `None` or `"model"` if getting repository info from a model. Default is `None`. + timeout (`float`, *optional*): + Whether to set a timeout for the request to the Hub. + expand (`ExpandModelProperty_T` or `ExpandDatasetProperty_T` or `ExpandSpaceProperty_T`, *optional*): + List properties to return in the response. When used, only the properties in the list will be returned. + This parameter cannot be used if `files_metadata` is passed. + For an exhaustive list of available properties, check out [`model_info`], [`dataset_info`] or [`space_info`]. + files_metadata (`bool`, *optional*): + Whether or not to retrieve metadata for files in the repository + (size, LFS metadata, etc). Defaults to `False`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Union[SpaceInfo, DatasetInfo, ModelInfo]`: The repository information, as a + [`huggingface_hub.hf_api.DatasetInfo`], [`huggingface_hub.hf_api.ModelInfo`] + or [`huggingface_hub.hf_api.SpaceInfo`] object. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + + + """ + if repo_type is None or repo_type == "model": + method = self.model_info + elif repo_type == "dataset": + method = self.dataset_info # type: ignore + elif repo_type == "space": + method = self.space_info # type: ignore + else: + raise ValueError("Unsupported repo type.") + return method( + repo_id, + revision=revision, + token=token, + timeout=timeout, + expand=expand, # type: ignore[arg-type] + files_metadata=files_metadata, + ) + + @validate_hf_hub_args + def repo_exists( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> bool: + """ + Checks if a repository exists on the Hugging Face Hub. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if getting repository info from a dataset or a space, + `None` or `"model"` if getting repository info from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + True if the repository exists, False otherwise. + + Examples: + ```py + >>> from huggingface_hub import repo_exists + >>> repo_exists("google/gemma-7b") + True + >>> repo_exists("google/not-a-repo") + False + ``` + """ + try: + self.repo_info(repo_id=repo_id, repo_type=repo_type, token=token) + return True + except GatedRepoError: + return True # we don't have access but it exists + except RepositoryNotFoundError: + return False + + @validate_hf_hub_args + def revision_exists( + self, + repo_id: str, + revision: str, + *, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> bool: + """ + Checks if a specific revision exists on a repo on the Hugging Face Hub. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`): + The revision of the repository to check. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if getting repository info from a dataset or a space, + `None` or `"model"` if getting repository info from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + True if the repository and the revision exists, False otherwise. + + Examples: + ```py + >>> from huggingface_hub import revision_exists + >>> revision_exists("google/gemma-7b", "float16") + True + >>> revision_exists("google/gemma-7b", "not-a-revision") + False + ``` + """ + try: + self.repo_info(repo_id=repo_id, revision=revision, repo_type=repo_type, token=token) + return True + except RevisionNotFoundError: + return False + except RepositoryNotFoundError: + return False + + @validate_hf_hub_args + def file_exists( + self, + repo_id: str, + filename: str, + *, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> bool: + """ + Checks if a file exists in a repository on the Hugging Face Hub. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + filename (`str`): + The name of the file to check, for example: + `"config.json"` + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if getting repository info from a dataset or a space, + `None` or `"model"` if getting repository info from a model. Default is `None`. + revision (`str`, *optional*): + The revision of the repository from which to get the information. Defaults to `"main"` branch. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + True if the file exists, False otherwise. + + Examples: + ```py + >>> from huggingface_hub import file_exists + >>> file_exists("bigcode/starcoder", "config.json") + True + >>> file_exists("bigcode/starcoder", "not-a-file") + False + >>> file_exists("bigcode/not-a-repo", "config.json") + False + ``` + """ + url = hf_hub_url( + repo_id=repo_id, repo_type=repo_type, revision=revision, filename=filename, endpoint=self.endpoint + ) + try: + if token is None: + token = self.token + get_hf_file_metadata(url, token=token) + return True + except GatedRepoError: # raise specifically on gated repo + raise + except (RepositoryNotFoundError, EntryNotFoundError, RevisionNotFoundError): + return False + + @validate_hf_hub_args + def list_repo_files( + self, + repo_id: str, + *, + revision: Optional[str] = None, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> List[str]: + """ + Get the list of files in a given repo. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + revision (`str`, *optional*): + The revision of the repository from which to get the information. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or space, `None` or `"model"` if uploading to + a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[str]`: the list of files in a given repository. + """ + return [ + f.rfilename + for f in self.list_repo_tree( + repo_id=repo_id, recursive=True, revision=revision, repo_type=repo_type, token=token + ) + if isinstance(f, RepoFile) + ] + + @validate_hf_hub_args + def list_repo_tree( + self, + repo_id: str, + path_in_repo: Optional[str] = None, + *, + recursive: bool = False, + expand: bool = False, + revision: Optional[str] = None, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> Iterable[Union[RepoFile, RepoFolder]]: + """ + List a repo tree's files and folders and get information about them. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + path_in_repo (`str`, *optional*): + Relative path of the tree (folder) in the repo, for example: + `"checkpoints/1fec34a/results"`. Will default to the root tree (folder) of the repository. + recursive (`bool`, *optional*, defaults to `False`): + Whether to list tree's files and folders recursively. + expand (`bool`, *optional*, defaults to `False`): + Whether to fetch more information about the tree's files and folders (e.g. last commit and files' security scan results). This + operation is more expensive for the server so only 50 results are returned per page (instead of 1000). + As pagination is implemented in `huggingface_hub`, this is transparent for you except for the time it + takes to get the results. + revision (`str`, *optional*): + The revision of the repository from which to get the tree. Defaults to `"main"` branch. + repo_type (`str`, *optional*): + The type of the repository from which to get the tree (`"model"`, `"dataset"` or `"space"`. + Defaults to `"model"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[Union[RepoFile, RepoFolder]]`: + The information about the tree's files and folders, as an iterable of [`RepoFile`] and [`RepoFolder`] objects. The order of the files and folders is + not guaranteed. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private but not authenticated or repo + does not exist. + [`~utils.RevisionNotFoundError`]: + If revision is not found (error 404) on the repo. + [`~utils.EntryNotFoundError`]: + If the tree (folder) does not exist (error 404) on the repo. + + Examples: + + Get information about a repo's tree. + ```py + >>> from huggingface_hub import list_repo_tree + >>> repo_tree = list_repo_tree("lysandre/arxiv-nlp") + >>> repo_tree + + >>> list(repo_tree) + [ + RepoFile(path='.gitattributes', size=391, blob_id='ae8c63daedbd4206d7d40126955d4e6ab1c80f8f', lfs=None, last_commit=None, security=None), + RepoFile(path='README.md', size=391, blob_id='43bd404b159de6fba7c2f4d3264347668d43af25', lfs=None, last_commit=None, security=None), + RepoFile(path='config.json', size=554, blob_id='2f9618c3a19b9a61add74f70bfb121335aeef666', lfs=None, last_commit=None, security=None), + RepoFile( + path='flax_model.msgpack', size=497764107, blob_id='8095a62ccb4d806da7666fcda07467e2d150218e', + lfs={'size': 497764107, 'sha256': 'd88b0d6a6ff9c3f8151f9d3228f57092aaea997f09af009eefd7373a77b5abb9', 'pointer_size': 134}, last_commit=None, security=None + ), + RepoFile(path='merges.txt', size=456318, blob_id='226b0752cac7789c48f0cb3ec53eda48b7be36cc', lfs=None, last_commit=None, security=None), + RepoFile( + path='pytorch_model.bin', size=548123560, blob_id='64eaa9c526867e404b68f2c5d66fd78e27026523', + lfs={'size': 548123560, 'sha256': '9be78edb5b928eba33aa88f431551348f7466ba9f5ef3daf1d552398722a5436', 'pointer_size': 134}, last_commit=None, security=None + ), + RepoFile(path='vocab.json', size=898669, blob_id='b00361fece0387ca34b4b8b8539ed830d644dbeb', lfs=None, last_commit=None, security=None)] + ] + ``` + + Get even more information about a repo's tree (last commit and files' security scan results) + ```py + >>> from huggingface_hub import list_repo_tree + >>> repo_tree = list_repo_tree("prompthero/openjourney-v4", expand=True) + >>> list(repo_tree) + [ + RepoFolder( + path='feature_extractor', + tree_id='aa536c4ea18073388b5b0bc791057a7296a00398', + last_commit={ + 'oid': '47b62b20b20e06b9de610e840282b7e6c3d51190', + 'title': 'Upload diffusers weights (#48)', + 'date': datetime.datetime(2023, 3, 21, 9, 5, 27, tzinfo=datetime.timezone.utc) + } + ), + RepoFolder( + path='safety_checker', + tree_id='65aef9d787e5557373fdf714d6c34d4fcdd70440', + last_commit={ + 'oid': '47b62b20b20e06b9de610e840282b7e6c3d51190', + 'title': 'Upload diffusers weights (#48)', + 'date': datetime.datetime(2023, 3, 21, 9, 5, 27, tzinfo=datetime.timezone.utc) + } + ), + RepoFile( + path='model_index.json', + size=582, + blob_id='d3d7c1e8c3e78eeb1640b8e2041ee256e24c9ee1', + lfs=None, + last_commit={ + 'oid': 'b195ed2d503f3eb29637050a886d77bd81d35f0e', + 'title': 'Fix deprecation warning by changing `CLIPFeatureExtractor` to `CLIPImageProcessor`. (#54)', + 'date': datetime.datetime(2023, 5, 15, 21, 41, 59, tzinfo=datetime.timezone.utc) + }, + security={ + 'safe': True, + 'av_scan': {'virusFound': False, 'virusNames': None}, + 'pickle_import_scan': None + } + ) + ... + ] + ``` + """ + repo_type = repo_type or constants.REPO_TYPE_MODEL + revision = quote(revision, safe="") if revision is not None else constants.DEFAULT_REVISION + headers = self._build_hf_headers(token=token) + + encoded_path_in_repo = "/" + quote(path_in_repo, safe="") if path_in_repo else "" + tree_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/tree/{revision}{encoded_path_in_repo}" + for path_info in paginate(path=tree_url, headers=headers, params={"recursive": recursive, "expand": expand}): + yield (RepoFile(**path_info) if path_info["type"] == "file" else RepoFolder(**path_info)) + + @validate_hf_hub_args + def list_repo_refs( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + include_pull_requests: bool = False, + token: Union[str, bool, None] = None, + ) -> GitRefs: + """ + Get the list of refs of a given repo (both tags and branches). + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if listing refs from a dataset or a Space, + `None` or `"model"` if listing from a model. Default is `None`. + include_pull_requests (`bool`, *optional*): + Whether to include refs from pull requests in the list. Defaults to `False`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> api.list_repo_refs("gpt2") + GitRefs(branches=[GitRefInfo(name='main', ref='refs/heads/main', target_commit='e7da7f221d5bf496a48136c0cd264e630fe9fcc8')], converts=[], tags=[]) + + >>> api.list_repo_refs("bigcode/the-stack", repo_type='dataset') + GitRefs( + branches=[ + GitRefInfo(name='main', ref='refs/heads/main', target_commit='18edc1591d9ce72aa82f56c4431b3c969b210ae3'), + GitRefInfo(name='v1.1.a1', ref='refs/heads/v1.1.a1', target_commit='f9826b862d1567f3822d3d25649b0d6d22ace714') + ], + converts=[], + tags=[ + GitRefInfo(name='v1.0', ref='refs/tags/v1.0', target_commit='c37a8cd1e382064d8aced5e05543c5f7753834da') + ] + ) + ``` + + Returns: + [`GitRefs`]: object containing all information about branches and tags for a + repo on the Hub. + """ + repo_type = repo_type or constants.REPO_TYPE_MODEL + response = get_session().get( + f"{self.endpoint}/api/{repo_type}s/{repo_id}/refs", + headers=self._build_hf_headers(token=token), + params={"include_prs": 1} if include_pull_requests else {}, + ) + hf_raise_for_status(response) + data = response.json() + + def _format_as_git_ref_info(item: Dict) -> GitRefInfo: + return GitRefInfo(name=item["name"], ref=item["ref"], target_commit=item["targetCommit"]) + + return GitRefs( + branches=[_format_as_git_ref_info(item) for item in data["branches"]], + converts=[_format_as_git_ref_info(item) for item in data["converts"]], + tags=[_format_as_git_ref_info(item) for item in data["tags"]], + pull_requests=[_format_as_git_ref_info(item) for item in data["pullRequests"]] + if include_pull_requests + else None, + ) + + @validate_hf_hub_args + def list_repo_commits( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + revision: Optional[str] = None, + formatted: bool = False, + ) -> List[GitCommitInfo]: + """ + Get the list of commits of a given revision for a repo on the Hub. + + Commits are sorted by date (last commit first). + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if listing commits from a dataset or a Space, `None` or `"model"` if + listing from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + formatted (`bool`): + Whether to return the HTML-formatted title and description of the commits. Defaults to False. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + + # Commits are sorted by date (last commit first) + >>> initial_commit = api.list_repo_commits("gpt2")[-1] + + # Initial commit is always a system commit containing the `.gitattributes` file. + >>> initial_commit + GitCommitInfo( + commit_id='9b865efde13a30c13e0a33e536cf3e4a5a9d71d8', + authors=['system'], + created_at=datetime.datetime(2019, 2, 18, 10, 36, 15, tzinfo=datetime.timezone.utc), + title='initial commit', + message='', + formatted_title=None, + formatted_message=None + ) + + # Create an empty branch by deriving from initial commit + >>> api.create_branch("gpt2", "new_empty_branch", revision=initial_commit.commit_id) + ``` + + Returns: + List[[`GitCommitInfo`]]: list of objects containing information about the commits for a repo on the Hub. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private but not authenticated or repo + does not exist. + [`~utils.RevisionNotFoundError`]: + If revision is not found (error 404) on the repo. + """ + repo_type = repo_type or constants.REPO_TYPE_MODEL + revision = quote(revision, safe="") if revision is not None else constants.DEFAULT_REVISION + + # Paginate over results and return the list of commits. + return [ + GitCommitInfo( + commit_id=item["id"], + authors=[author["user"] for author in item["authors"]], + created_at=parse_datetime(item["date"]), + title=item["title"], + message=item["message"], + formatted_title=item.get("formatted", {}).get("title"), + formatted_message=item.get("formatted", {}).get("message"), + ) + for item in paginate( + f"{self.endpoint}/api/{repo_type}s/{repo_id}/commits/{revision}", + headers=self._build_hf_headers(token=token), + params={"expand[]": "formatted"} if formatted else {}, + ) + ] + + @validate_hf_hub_args + def get_paths_info( + self, + repo_id: str, + paths: Union[List[str], str], + *, + expand: bool = False, + revision: Optional[str] = None, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> List[Union[RepoFile, RepoFolder]]: + """ + Get information about a repo's paths. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + paths (`Union[List[str], str]`, *optional*): + The paths to get information about. If a path do not exist, it is ignored without raising + an exception. + expand (`bool`, *optional*, defaults to `False`): + Whether to fetch more information about the paths (e.g. last commit and files' security scan results). This + operation is more expensive for the server so only 50 results are returned per page (instead of 1000). + As pagination is implemented in `huggingface_hub`, this is transparent for you except for the time it + takes to get the results. + revision (`str`, *optional*): + The revision of the repository from which to get the information. Defaults to `"main"` branch. + repo_type (`str`, *optional*): + The type of the repository from which to get the information (`"model"`, `"dataset"` or `"space"`. + Defaults to `"model"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[Union[RepoFile, RepoFolder]]`: + The information about the paths, as a list of [`RepoFile`] and [`RepoFolder`] objects. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private but not authenticated or repo + does not exist. + [`~utils.RevisionNotFoundError`]: + If revision is not found (error 404) on the repo. + + Example: + ```py + >>> from huggingface_hub import get_paths_info + >>> paths_info = get_paths_info("allenai/c4", ["README.md", "en"], repo_type="dataset") + >>> paths_info + [ + RepoFile(path='README.md', size=2379, blob_id='f84cb4c97182890fc1dbdeaf1a6a468fd27b4fff', lfs=None, last_commit=None, security=None), + RepoFolder(path='en', tree_id='dc943c4c40f53d02b31ced1defa7e5f438d5862e', last_commit=None) + ] + ``` + """ + repo_type = repo_type or constants.REPO_TYPE_MODEL + revision = quote(revision, safe="") if revision is not None else constants.DEFAULT_REVISION + headers = self._build_hf_headers(token=token) + + response = get_session().post( + f"{self.endpoint}/api/{repo_type}s/{repo_id}/paths-info/{revision}", + data={ + "paths": paths if isinstance(paths, list) else [paths], + "expand": expand, + }, + headers=headers, + ) + hf_raise_for_status(response) + paths_info = response.json() + return [ + RepoFile(**path_info) if path_info["type"] == "file" else RepoFolder(**path_info) + for path_info in paths_info + ] + + @validate_hf_hub_args + def super_squash_history( + self, + repo_id: str, + *, + branch: Optional[str] = None, + commit_message: Optional[str] = None, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ) -> None: + """Squash commit history on a branch for a repo on the Hub. + + Squashing the repo history is useful when you know you'll make hundreds of commits and you don't want to + clutter the history. Squashing commits can only be performed from the head of a branch. + + + + Once squashed, the commit history cannot be retrieved. This is a non-revertible operation. + + + + + + Once the history of a branch has been squashed, it is not possible to merge it back into another branch since + their history will have diverged. + + + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + branch (`str`, *optional*): + The branch to squash. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The commit message to use for the squashed commit. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if listing commits from a dataset or a Space, `None` or `"model"` if + listing from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private but not authenticated or repo + does not exist. + [`~utils.RevisionNotFoundError`]: + If the branch to squash cannot be found. + [`~utils.BadRequestError`]: + If invalid reference for a branch. You cannot squash history on tags. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + + # Create repo + >>> repo_id = api.create_repo("test-squash").repo_id + + # Make a lot of commits. + >>> api.upload_file(repo_id=repo_id, path_in_repo="file.txt", path_or_fileobj=b"content") + >>> api.upload_file(repo_id=repo_id, path_in_repo="lfs.bin", path_or_fileobj=b"content") + >>> api.upload_file(repo_id=repo_id, path_in_repo="file.txt", path_or_fileobj=b"another_content") + + # Squash history + >>> api.super_squash_history(repo_id=repo_id) + ``` + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + if repo_type not in constants.REPO_TYPES: + raise ValueError("Invalid repo type") + if branch is None: + branch = constants.DEFAULT_REVISION + + # Prepare request + url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/super-squash/{quote(branch, safe='')}" + headers = self._build_hf_headers(token=token) + commit_message = commit_message or f"Super-squash branch '{branch}' using huggingface_hub" + + # Super-squash + response = get_session().post(url=url, headers=headers, json={"message": commit_message}) + hf_raise_for_status(response) + + @validate_hf_hub_args + def list_lfs_files( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[LFSFileInfo]: + """ + List all LFS files in a repo on the Hub. + + This is primarily useful to count how much storage a repo is using and to eventually clean up large files + with [`permanently_delete_lfs_files`]. Note that this would be a permanent action that will affect all commits + referencing this deleted files and that cannot be undone. + + Args: + repo_id (`str`): + The repository for which you are listing LFS files. + repo_type (`str`, *optional*): + Type of repository. Set to `"dataset"` or `"space"` if listing from a dataset or space, `None` or + `"model"` if listing from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[LFSFileInfo]`: An iterator of [`LFSFileInfo`] objects. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> lfs_files = api.list_lfs_files("username/my-cool-repo") + + # Filter files files to delete based on a combination of `filename`, `pushed_at`, `ref` or `size`. + # e.g. select only LFS files in the "checkpoints" folder + >>> lfs_files_to_delete = (lfs_file for lfs_file in lfs_files if lfs_file.filename.startswith("checkpoints/")) + + # Permanently delete LFS files + >>> api.permanently_delete_lfs_files("username/my-cool-repo", lfs_files_to_delete) + ``` + """ + # Prepare request + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/lfs-files" + headers = self._build_hf_headers(token=token) + + # Paginate over LFS items + for item in paginate(url, params={}, headers=headers): + yield LFSFileInfo(**item) + + @validate_hf_hub_args + def permanently_delete_lfs_files( + self, + repo_id: str, + lfs_files: Iterable[LFSFileInfo], + *, + rewrite_history: bool = True, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> None: + """ + Permanently delete LFS files from a repo on the Hub. + + + + This is a permanent action that will affect all commits referencing the deleted files and might corrupt your + repository. This is a non-revertible operation. Use it only if you know what you are doing. + + + + Args: + repo_id (`str`): + The repository for which you are listing LFS files. + lfs_files (`Iterable[LFSFileInfo]`): + An iterable of [`LFSFileInfo`] items to permanently delete from the repo. Use [`list_lfs_files`] to list + all LFS files from a repo. + rewrite_history (`bool`, *optional*, default to `True`): + Whether to rewrite repository history to remove file pointers referencing the deleted LFS files (recommended). + repo_type (`str`, *optional*): + Type of repository. Set to `"dataset"` or `"space"` if listing from a dataset or space, `None` or + `"model"` if listing from a model. Default is `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Example: + ```py + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> lfs_files = api.list_lfs_files("username/my-cool-repo") + + # Filter files files to delete based on a combination of `filename`, `pushed_at`, `ref` or `size`. + # e.g. select only LFS files in the "checkpoints" folder + >>> lfs_files_to_delete = (lfs_file for lfs_file in lfs_files if lfs_file.filename.startswith("checkpoints/")) + + # Permanently delete LFS files + >>> api.permanently_delete_lfs_files("username/my-cool-repo", lfs_files_to_delete) + ``` + """ + # Prepare request + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/lfs-files/batch" + headers = self._build_hf_headers(token=token) + + # Delete LFS items by batches of 1000 + for batch in chunk_iterable(lfs_files, 1000): + shas = [item.file_oid for item in batch] + if len(shas) == 0: + return + payload = { + "deletions": { + "sha": shas, + "rewriteHistory": rewrite_history, + } + } + response = get_session().post(url, headers=headers, json=payload) + hf_raise_for_status(response) + + @validate_hf_hub_args + def create_repo( + self, + repo_id: str, + *, + token: Union[str, bool, None] = None, + private: Optional[bool] = None, + repo_type: Optional[str] = None, + exist_ok: bool = False, + resource_group_id: Optional[str] = None, + space_sdk: Optional[str] = None, + space_hardware: Optional[SpaceHardware] = None, + space_storage: Optional[SpaceStorage] = None, + space_sleep_time: Optional[int] = None, + space_secrets: Optional[List[Dict[str, str]]] = None, + space_variables: Optional[List[Dict[str, str]]] = None, + ) -> RepoUrl: + """Create an empty repo on the HuggingFace Hub. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + exist_ok (`bool`, *optional*, defaults to `False`): + If `True`, do not raise an error if repo already exists. + resource_group_id (`str`, *optional*): + Resource group in which to create the repo. Resource groups is only available for organizations and + allow to define which members of the organization can access the resource. The ID of a resource group + can be found in the URL of the resource's page on the Hub (e.g. `"66670e5163145ca562cb1988"`). + To learn more about resource groups, see https://huggingface.co/docs/hub/en/security-resource-groups. + space_sdk (`str`, *optional*): + Choice of SDK to use if repo_type is "space". Can be "streamlit", "gradio", "docker", or "static". + space_hardware (`SpaceHardware` or `str`, *optional*): + Choice of Hardware if repo_type is "space". See [`SpaceHardware`] for a complete list. + space_storage (`SpaceStorage` or `str`, *optional*): + Choice of persistent storage tier. Example: `"small"`. See [`SpaceStorage`] for a complete list. + space_sleep_time (`int`, *optional*): + Number of seconds of inactivity to wait before a Space is put to sleep. Set to `-1` if you don't want + your Space to sleep (default behavior for upgraded hardware). For free hardware, you can't configure + the sleep time (value is fixed to 48 hours of inactivity). + See https://huggingface.co/docs/hub/spaces-gpus#sleep-time for more details. + space_secrets (`List[Dict[str, str]]`, *optional*): + A list of secret keys to set in your Space. Each item is in the form `{"key": ..., "value": ..., "description": ...}` where description is optional. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets. + space_variables (`List[Dict[str, str]]`, *optional*): + A list of public environment variables to set in your Space. Each item is in the form `{"key": ..., "value": ..., "description": ...}` where description is optional. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets-and-environment-variables. + + Returns: + [`RepoUrl`]: URL to the newly created repo. Value is a subclass of `str` containing + attributes like `endpoint`, `repo_type` and `repo_id`. + """ + organization, name = repo_id.split("/") if "/" in repo_id else (None, repo_id) + + path = f"{self.endpoint}/api/repos/create" + + if repo_type not in constants.REPO_TYPES: + raise ValueError("Invalid repo type") + + json: Dict[str, Any] = {"name": name, "organization": organization} + if private is not None: + json["private"] = private + if repo_type is not None: + json["type"] = repo_type + if repo_type == "space": + if space_sdk is None: + raise ValueError( + "No space_sdk provided. `create_repo` expects space_sdk to be one" + f" of {constants.SPACES_SDK_TYPES} when repo_type is 'space'`" + ) + if space_sdk not in constants.SPACES_SDK_TYPES: + raise ValueError(f"Invalid space_sdk. Please choose one of {constants.SPACES_SDK_TYPES}.") + json["sdk"] = space_sdk + + if space_sdk is not None and repo_type != "space": + warnings.warn("Ignoring provided space_sdk because repo_type is not 'space'.") + + function_args = [ + "space_hardware", + "space_storage", + "space_sleep_time", + "space_secrets", + "space_variables", + ] + json_keys = ["hardware", "storageTier", "sleepTimeSeconds", "secrets", "variables"] + values = [space_hardware, space_storage, space_sleep_time, space_secrets, space_variables] + + if repo_type == "space": + json.update({k: v for k, v in zip(json_keys, values) if v is not None}) + else: + provided_space_args = [key for key, value in zip(function_args, values) if value is not None] + + if provided_space_args: + warnings.warn(f"Ignoring provided {', '.join(provided_space_args)} because repo_type is not 'space'.") + + if getattr(self, "_lfsmultipartthresh", None): + # Testing purposes only. + # See https://github.com/huggingface/huggingface_hub/pull/733/files#r820604472 + json["lfsmultipartthresh"] = self._lfsmultipartthresh # type: ignore + + if resource_group_id is not None: + json["resourceGroupId"] = resource_group_id + + headers = self._build_hf_headers(token=token) + while True: + r = get_session().post(path, headers=headers, json=json) + if r.status_code == 409 and "Cannot create repo: another conflicting operation is in progress" in r.text: + # Since https://github.com/huggingface/moon-landing/pull/7272 (private repo), it is not possible to + # concurrently create repos on the Hub for a same user. This is rarely an issue, except when running + # tests. To avoid any inconvenience, we retry to create the repo for this specific error. + # NOTE: This could have being fixed directly in the tests but adding it here should fixed CIs for all + # dependent libraries. + # NOTE: If a fix is implemented server-side, we should be able to remove this retry mechanism. + logger.debug("Create repo failed due to a concurrency issue. Retrying...") + continue + break + + try: + hf_raise_for_status(r) + except HTTPError as err: + if exist_ok and err.response.status_code == 409: + # Repo already exists and `exist_ok=True` + pass + elif exist_ok and err.response.status_code == 403: + # No write permission on the namespace but repo might already exist + try: + self.repo_info(repo_id=repo_id, repo_type=repo_type, token=token) + if repo_type is None or repo_type == constants.REPO_TYPE_MODEL: + return RepoUrl(f"{self.endpoint}/{repo_id}") + return RepoUrl(f"{self.endpoint}/{repo_type}/{repo_id}") + except HfHubHTTPError: + raise err + else: + raise + + d = r.json() + return RepoUrl(d["url"], endpoint=self.endpoint) + + @validate_hf_hub_args + def delete_repo( + self, + repo_id: str, + *, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + missing_ok: bool = False, + ) -> None: + """ + Delete a repo from the HuggingFace Hub. CAUTION: this is irreversible. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. + missing_ok (`bool`, *optional*, defaults to `False`): + If `True`, do not raise an error if repo does not exist. + + Raises: + [`~utils.RepositoryNotFoundError`] + If the repository to delete from cannot be found and `missing_ok` is set to False (default). + """ + organization, name = repo_id.split("/") if "/" in repo_id else (None, repo_id) + + path = f"{self.endpoint}/api/repos/delete" + + if repo_type not in constants.REPO_TYPES: + raise ValueError("Invalid repo type") + + json = {"name": name, "organization": organization} + if repo_type is not None: + json["type"] = repo_type + + headers = self._build_hf_headers(token=token) + r = get_session().delete(path, headers=headers, json=json) + try: + hf_raise_for_status(r) + except RepositoryNotFoundError: + if not missing_ok: + raise + + @_deprecate_method(version="0.32", message="Please use `update_repo_settings` instead.") + @validate_hf_hub_args + def update_repo_visibility( + self, + repo_id: str, + private: bool = False, + *, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + ) -> Dict[str, bool]: + """Update the visibility setting of a repository. + + Deprecated. Use `update_repo_settings` instead. + + Args: + repo_id (`str`, *optional*): + A namespace (user or an organization) and a repo name separated by a `/`. + private (`bool`, *optional*, defaults to `False`): + Whether the repository should be private. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + + Returns: + The HTTP response in json. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL # default repo type + + r = get_session().put( + url=f"{self.endpoint}/api/{repo_type}s/{repo_id}/settings", + headers=self._build_hf_headers(token=token), + json={"private": private}, + ) + hf_raise_for_status(r) + return r.json() + + @validate_hf_hub_args + def update_repo_settings( + self, + repo_id: str, + *, + gated: Optional[Literal["auto", "manual", False]] = None, + private: Optional[bool] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + xet_enabled: Optional[bool] = None, + ) -> None: + """ + Update the settings of a repository, including gated access and visibility. + + To give more control over how repos are used, the Hub allows repo authors to enable + access requests for their repos, and also to set the visibility of the repo to private. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a /. + gated (`Literal["auto", "manual", False]`, *optional*): + The gated status for the repository. If set to `None` (default), the `gated` setting of the repository won't be updated. + * "auto": The repository is gated, and access requests are automatically approved or denied based on predefined criteria. + * "manual": The repository is gated, and access requests require manual approval. + * False : The repository is not gated, and anyone can access it. + private (`bool`, *optional*): + Whether the repository should be private. + token (`Union[str, bool, None]`, *optional*): + A valid user access token (string). Defaults to the locally saved token, + which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass False. + repo_type (`str`, *optional*): + The type of the repository to update settings from (`"model"`, `"dataset"` or `"space"`). + Defaults to `"model"`. + xet_enabled (`bool`, *optional*): + Whether the repository should be enabled for Xet Storage. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If gated is not one of "auto", "manual", or False. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If repo_type is not one of the values in constants.REPO_TYPES. + [`~utils.HfHubHTTPError`]: + If the request to the Hugging Face Hub API fails. + [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + """ + + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL # default repo type + + # Prepare the JSON payload for the PUT request + payload: Dict = {} + + if gated is not None: + if gated not in ["auto", "manual", False]: + raise ValueError(f"Invalid gated status, must be one of 'auto', 'manual', or False. Got '{gated}'.") + payload["gated"] = gated + + if private is not None: + payload["private"] = private + + if xet_enabled is not None: + payload["xetEnabled"] = xet_enabled + + if len(payload) == 0: + raise ValueError("At least one setting must be updated.") + + # Build headers + headers = self._build_hf_headers(token=token) + + r = get_session().put( + url=f"{self.endpoint}/api/{repo_type}s/{repo_id}/settings", + headers=headers, + json=payload, + ) + hf_raise_for_status(r) + + def move_repo( + self, + from_id: str, + to_id: str, + *, + repo_type: Optional[str] = None, + token: Union[str, bool, None] = None, + ): + """ + Moving a repository from namespace1/repo_name1 to namespace2/repo_name2 + + Note there are certain limitations. For more information about moving + repositories, please see + https://hf.co/docs/hub/repositories-settings#renaming-or-transferring-a-repo. + + Args: + from_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. Original repository identifier. + to_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. Final repository identifier. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + + + Raises the following errors: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + if len(from_id.split("/")) != 2: + raise ValueError(f"Invalid repo_id: {from_id}. It should have a namespace (:namespace:/:repo_name:)") + + if len(to_id.split("/")) != 2: + raise ValueError(f"Invalid repo_id: {to_id}. It should have a namespace (:namespace:/:repo_name:)") + + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL # Hub won't accept `None`. + + json = {"fromRepo": from_id, "toRepo": to_id, "type": repo_type} + + path = f"{self.endpoint}/api/repos/move" + headers = self._build_hf_headers(token=token) + r = get_session().post(path, headers=headers, json=json) + try: + hf_raise_for_status(r) + except HfHubHTTPError as e: + e.append_to_message( + "\nFor additional documentation please see" + " https://hf.co/docs/hub/repositories-settings#renaming-or-transferring-a-repo." + ) + raise + + @overload + def create_commit( # type: ignore + self, + repo_id: str, + operations: Iterable[CommitOperation], + *, + commit_message: str, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + num_threads: int = 5, + parent_commit: Optional[str] = None, + run_as_future: Literal[False] = ..., + ) -> CommitInfo: ... + + @overload + def create_commit( + self, + repo_id: str, + operations: Iterable[CommitOperation], + *, + commit_message: str, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + num_threads: int = 5, + parent_commit: Optional[str] = None, + run_as_future: Literal[True] = ..., + ) -> Future[CommitInfo]: ... + + @validate_hf_hub_args + @future_compatible + def create_commit( + self, + repo_id: str, + operations: Iterable[CommitOperation], + *, + commit_message: str, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + num_threads: int = 5, + parent_commit: Optional[str] = None, + run_as_future: bool = False, + ) -> Union[CommitInfo, Future[CommitInfo]]: + """ + Creates a commit in the given repo, deleting & uploading files as needed. + + + + The input list of `CommitOperation` will be mutated during the commit process. Do not reuse the same objects + for multiple commits. + + + + + + `create_commit` assumes that the repo already exists on the Hub. If you get a + Client error 404, please make sure you are authenticated and that `repo_id` and + `repo_type` are set correctly. If repo does not exist, create it first using + [`~hf_api.create_repo`]. + + + + + + `create_commit` is limited to 25k LFS files and a 1GB payload for regular files. + + + + Args: + repo_id (`str`): + The repository in which the commit will be created, for example: + `"username/custom_transformers"` + + operations (`Iterable` of [`~hf_api.CommitOperation`]): + An iterable of operations to include in the commit, either: + + - [`~hf_api.CommitOperationAdd`] to upload a file + - [`~hf_api.CommitOperationDelete`] to delete a file + - [`~hf_api.CommitOperationCopy`] to copy a file + + Operation objects will be mutated to include information relative to the upload. Do not reuse the + same objects for multiple commits. + + commit_message (`str`): + The summary (first line) of the commit that will be created. + + commit_description (`str`, *optional*): + The description of the commit that will be created + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. + If `revision` is not set, PR is opened against the `"main"` branch. If + `revision` is set and is a branch, PR is opened against this branch. If + `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + + num_threads (`int`, *optional*): + Number of concurrent threads for uploading files. Defaults to 5. + Setting it to 2 means at most 2 files will be uploaded concurrently. + + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. + Shorthands (7 first characters) are also supported. If specified and `create_pr` is `False`, + the commit will fail if `revision` does not point to `parent_commit`. If specified and `create_pr` + is `True`, the pull request will be created from `parent_commit`. Specifying `parent_commit` + ensures the repo has not changed before committing the changes, and can be especially useful + if the repo is updated / committed to concurrently. + run_as_future (`bool`, *optional*): + Whether or not to run this method in the background. Background jobs are run sequentially without + blocking the main thread. Passing `run_as_future=True` will return a [Future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) + object. Defaults to `False`. + + Returns: + [`CommitInfo`] or `Future`: + Instance of [`CommitInfo`] containing information about the newly created commit (commit hash, commit + url, pr url, commit message,...). If `run_as_future=True` is passed, returns a Future object which will + contain the result when executed. + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If commit message is empty. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If parent commit is not a valid commit OID. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If a README.md file with an invalid metadata section is committed. In this case, the commit will fail + early, before trying to upload any file. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `create_pr` is `True` and revision is neither `None` nor `"main"`. + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + """ + if parent_commit is not None and not constants.REGEX_COMMIT_OID.fullmatch(parent_commit): + raise ValueError( + f"`parent_commit` is not a valid commit OID. It must match the following regex: {constants.REGEX_COMMIT_OID}" + ) + + if commit_message is None or len(commit_message) == 0: + raise ValueError("`commit_message` can't be empty, please pass a value.") + + commit_description = commit_description if commit_description is not None else "" + repo_type = repo_type if repo_type is not None else constants.REPO_TYPE_MODEL + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + unquoted_revision = revision or constants.DEFAULT_REVISION + revision = quote(unquoted_revision, safe="") + create_pr = create_pr if create_pr is not None else False + + headers = self._build_hf_headers(token=token) + + operations = list(operations) + additions = [op for op in operations if isinstance(op, CommitOperationAdd)] + copies = [op for op in operations if isinstance(op, CommitOperationCopy)] + nb_additions = len(additions) + nb_copies = len(copies) + nb_deletions = len(operations) - nb_additions - nb_copies + + for addition in additions: + if addition._is_committed: + raise ValueError( + f"CommitOperationAdd {addition} has already being committed and cannot be reused. Please create a" + " new CommitOperationAdd object if you want to create a new commit." + ) + + if repo_type != "dataset": + for addition in additions: + if addition.path_in_repo.endswith((".arrow", ".parquet")): + warnings.warn( + f"It seems that you are about to commit a data file ({addition.path_in_repo}) to a {repo_type}" + " repository. You are sure this is intended? If you are trying to upload a dataset, please" + " set `repo_type='dataset'` or `--repo-type=dataset` in a CLI." + ) + + logger.debug( + f"About to commit to the hub: {len(additions)} addition(s), {len(copies)} copie(s) and" + f" {nb_deletions} deletion(s)." + ) + + # If updating a README.md file, make sure the metadata format is valid + # It's better to fail early than to fail after all the files have been uploaded. + for addition in additions: + if addition.path_in_repo == "README.md": + with addition.as_file() as file: + content = file.read().decode() + self._validate_yaml(content, repo_type=repo_type, token=token) + # Skip other additions after `README.md` has been processed + break + + # If updating twice the same file or update then delete a file in a single commit + _warn_on_overwriting_operations(operations) + + self.preupload_lfs_files( + repo_id=repo_id, + additions=additions, + token=token, + repo_type=repo_type, + revision=unquoted_revision, # first-class methods take unquoted revision + create_pr=create_pr, + num_threads=num_threads, + free_memory=False, # do not remove `CommitOperationAdd.path_or_fileobj` on LFS files for "normal" users + ) + + files_to_copy = _fetch_files_to_copy( + copies=copies, + repo_type=repo_type, + repo_id=repo_id, + headers=headers, + revision=unquoted_revision, + endpoint=self.endpoint, + ) + # Remove no-op operations (files that have not changed) + operations_without_no_op = [] + for operation in operations: + if ( + isinstance(operation, CommitOperationAdd) + and operation._remote_oid is not None + and operation._remote_oid == operation._local_oid + ): + # File already exists on the Hub and has not changed: we can skip it. + logger.debug(f"Skipping upload for '{operation.path_in_repo}' as the file has not changed.") + continue + if ( + isinstance(operation, CommitOperationCopy) + and operation._dest_oid is not None + and operation._dest_oid == operation._src_oid + ): + # Source and destination files are identical - skip + logger.debug( + f"Skipping copy for '{operation.src_path_in_repo}' -> '{operation.path_in_repo}' as the content of the source file is the same as the destination file." + ) + continue + operations_without_no_op.append(operation) + if len(operations) != len(operations_without_no_op): + logger.info( + f"Removing {len(operations) - len(operations_without_no_op)} file(s) from commit that have not changed." + ) + + # Return early if empty commit + if len(operations_without_no_op) == 0: + logger.warning("No files have been modified since last commit. Skipping to prevent empty commit.") + + # Get latest commit info + try: + info = self.repo_info(repo_id=repo_id, repo_type=repo_type, revision=unquoted_revision, token=token) + except RepositoryNotFoundError as e: + e.append_to_message(_CREATE_COMMIT_NO_REPO_ERROR_MESSAGE) + raise + + # Return commit info based on latest commit + url_prefix = self.endpoint + if repo_type is not None and repo_type != constants.REPO_TYPE_MODEL: + url_prefix = f"{url_prefix}/{repo_type}s" + return CommitInfo( + commit_url=f"{url_prefix}/{repo_id}/commit/{info.sha}", + commit_message=commit_message, + commit_description=commit_description, + oid=info.sha, # type: ignore[arg-type] + ) + + commit_payload = _prepare_commit_payload( + operations=operations, + files_to_copy=files_to_copy, + commit_message=commit_message, + commit_description=commit_description, + parent_commit=parent_commit, + ) + commit_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/commit/{revision}" + + def _payload_as_ndjson() -> Iterable[bytes]: + for item in commit_payload: + yield json.dumps(item).encode() + yield b"\n" + + headers = { + # See https://github.com/huggingface/huggingface_hub/issues/1085#issuecomment-1265208073 + "Content-Type": "application/x-ndjson", + **headers, + } + data = b"".join(_payload_as_ndjson()) + params = {"create_pr": "1"} if create_pr else None + + try: + commit_resp = get_session().post(url=commit_url, headers=headers, data=data, params=params) + hf_raise_for_status(commit_resp, endpoint_name="commit") + except RepositoryNotFoundError as e: + e.append_to_message(_CREATE_COMMIT_NO_REPO_ERROR_MESSAGE) + raise + except EntryNotFoundError as e: + if nb_deletions > 0 and "A file with this name doesn't exist" in str(e): + e.append_to_message( + "\nMake sure to differentiate file and folder paths in delete" + " operations with a trailing '/' or using `is_folder=True/False`." + ) + raise + + # Mark additions as committed (cannot be reused in another commit) + for addition in additions: + addition._is_committed = True + + commit_data = commit_resp.json() + return CommitInfo( + commit_url=commit_data["commitUrl"], + commit_message=commit_message, + commit_description=commit_description, + oid=commit_data["commitOid"], + pr_url=commit_data["pullRequestUrl"] if create_pr else None, + ) + + def preupload_lfs_files( + self, + repo_id: str, + additions: Iterable[CommitOperationAdd], + *, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + num_threads: int = 5, + free_memory: bool = True, + gitignore_content: Optional[str] = None, + ): + """Pre-upload LFS files to S3 in preparation on a future commit. + + This method is useful if you are generating the files to upload on-the-fly and you don't want to store them + in memory before uploading them all at once. + + + + This is a power-user method. You shouldn't need to call it directly to make a normal commit. + Use [`create_commit`] directly instead. + + + + + + Commit operations will be mutated during the process. In particular, the attached `path_or_fileobj` will be + removed after the upload to save memory (and replaced by an empty `bytes` object). Do not reuse the same + objects except to pass them to [`create_commit`]. If you don't want to remove the attached content from the + commit operation object, pass `free_memory=False`. + + + + Args: + repo_id (`str`): + The repository in which you will commit the files, for example: `"username/custom_transformers"`. + + operations (`Iterable` of [`CommitOperationAdd`]): + The list of files to upload. Warning: the objects in this list will be mutated to include information + relative to the upload. Do not reuse the same objects for multiple commits. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + The type of repository to upload to (e.g. `"model"` -default-, `"dataset"` or `"space"`). + + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + + create_pr (`boolean`, *optional*): + Whether or not you plan to create a Pull Request with that commit. Defaults to `False`. + + num_threads (`int`, *optional*): + Number of concurrent threads for uploading files. Defaults to 5. + Setting it to 2 means at most 2 files will be uploaded concurrently. + + gitignore_content (`str`, *optional*): + The content of the `.gitignore` file to know which files should be ignored. The order of priority + is to first check if `gitignore_content` is passed, then check if the `.gitignore` file is present + in the list of files to commit and finally default to the `.gitignore` file already hosted on the Hub + (if any). + + Example: + ```py + >>> from huggingface_hub import CommitOperationAdd, preupload_lfs_files, create_commit, create_repo + + >>> repo_id = create_repo("test_preupload").repo_id + + # Generate and preupload LFS files one by one + >>> operations = [] # List of all `CommitOperationAdd` objects that will be generated + >>> for i in range(5): + ... content = ... # generate binary content + ... addition = CommitOperationAdd(path_in_repo=f"shard_{i}_of_5.bin", path_or_fileobj=content) + ... preupload_lfs_files(repo_id, additions=[addition]) # upload + free memory + ... operations.append(addition) + + # Create commit + >>> create_commit(repo_id, operations=operations, commit_message="Commit all shards") + ``` + """ + repo_type = repo_type if repo_type is not None else constants.REPO_TYPE_MODEL + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + revision = quote(revision, safe="") if revision is not None else constants.DEFAULT_REVISION + create_pr = create_pr if create_pr is not None else False + headers = self._build_hf_headers(token=token) + + # Check if a `gitignore` file is being committed to the Hub. + additions = list(additions) + if gitignore_content is None: + for addition in additions: + if addition.path_in_repo == ".gitignore": + with addition.as_file() as f: + gitignore_content = f.read().decode() + break + + # Filter out already uploaded files + new_additions = [addition for addition in additions if not addition._is_uploaded] + + # Check which new files are LFS + try: + _fetch_upload_modes( + additions=new_additions, + repo_type=repo_type, + repo_id=repo_id, + headers=headers, + revision=revision, + endpoint=self.endpoint, + create_pr=create_pr or False, + gitignore_content=gitignore_content, + ) + except RepositoryNotFoundError as e: + e.append_to_message(_CREATE_COMMIT_NO_REPO_ERROR_MESSAGE) + raise + + # Filter out regular files + new_lfs_additions = [addition for addition in new_additions if addition._upload_mode == "lfs"] + + # Filter out files listed in .gitignore + new_lfs_additions_to_upload = [] + for addition in new_lfs_additions: + if addition._should_ignore: + logger.debug(f"Skipping upload for LFS file '{addition.path_in_repo}' (ignored by gitignore file).") + else: + new_lfs_additions_to_upload.append(addition) + if len(new_lfs_additions) != len(new_lfs_additions_to_upload): + logger.info( + f"Skipped upload for {len(new_lfs_additions) - len(new_lfs_additions_to_upload)} LFS file(s) " + "(ignored by gitignore file)." + ) + # Prepare upload parameters + upload_kwargs = { + "additions": new_lfs_additions_to_upload, + "repo_type": repo_type, + "repo_id": repo_id, + "headers": headers, + "endpoint": self.endpoint, + # If `create_pr`, we don't want to check user permission on the revision as users with read permission + # should still be able to create PRs even if they don't have write permission on the target branch of the + # PR (i.e. `revision`). + "revision": revision if not create_pr else None, + } + # Upload files using Xet protocol if all of the following are true: + # - xet is enabled for the repo, + # - the files are provided as str or paths objects, + # - the library is installed. + # Otherwise, default back to LFS. + xet_enabled = self.repo_info( + repo_id=repo_id, + repo_type=repo_type, + revision=unquote(revision) if revision is not None else revision, + expand="xetEnabled", + token=token, + ).xet_enabled + has_buffered_io_data = any( + isinstance(addition.path_or_fileobj, io.BufferedIOBase) for addition in new_lfs_additions_to_upload + ) + if xet_enabled and not has_buffered_io_data and is_xet_available(): + logger.info("Uploading files using Xet Storage..") + _upload_xet_files(**upload_kwargs, create_pr=create_pr) # type: ignore [arg-type] + else: + if xet_enabled and is_xet_available(): + if has_buffered_io_data: + logger.warning( + "Uploading files as a binary IO buffer is not supported by Xet Storage. " + "Falling back to HTTP upload." + ) + _upload_lfs_files(**upload_kwargs, num_threads=num_threads) # type: ignore [arg-type] + for addition in new_lfs_additions_to_upload: + addition._is_uploaded = True + if free_memory: + addition.path_or_fileobj = b"" + + @overload + def upload_file( # type: ignore + self, + *, + path_or_fileobj: Union[str, Path, bytes, BinaryIO], + path_in_repo: str, + repo_id: str, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + run_as_future: Literal[False] = ..., + ) -> CommitInfo: ... + + @overload + def upload_file( + self, + *, + path_or_fileobj: Union[str, Path, bytes, BinaryIO], + path_in_repo: str, + repo_id: str, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + run_as_future: Literal[True] = ..., + ) -> Future[CommitInfo]: ... + + @validate_hf_hub_args + @future_compatible + def upload_file( + self, + *, + path_or_fileobj: Union[str, Path, bytes, BinaryIO], + path_in_repo: str, + repo_id: str, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + run_as_future: bool = False, + ) -> Union[CommitInfo, Future[CommitInfo]]: + """ + Upload a local file (up to 50 GB) to the given repo. The upload is done + through a HTTP post request, and doesn't require git or git-lfs to be + installed. + + Args: + path_or_fileobj (`str`, `Path`, `bytes`, or `IO`): + Path to a file on the local machine or binary data stream / + fileobj / buffer. + path_in_repo (`str`): + Relative filepath in the repo, for example: + `"checkpoints/1fec34a/weights.bin"` + repo_id (`str`): + The repository to which the file will be uploaded, for example: + `"username/custom_transformers"` + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit + commit_description (`str` *optional*) + The description of the generated commit + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. + If `revision` is not set, PR is opened against the `"main"` branch. If + `revision` is set and is a branch, PR is opened against this branch. If + `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + run_as_future (`bool`, *optional*): + Whether or not to run this method in the background. Background jobs are run sequentially without + blocking the main thread. Passing `run_as_future=True` will return a [Future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) + object. Defaults to `False`. + + + Returns: + [`CommitInfo`] or `Future`: + Instance of [`CommitInfo`] containing information about the newly created commit (commit hash, commit + url, pr url, commit message,...). If `run_as_future=True` is passed, returns a Future object which will + contain the result when executed. + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + + + + + + `upload_file` assumes that the repo already exists on the Hub. If you get a + Client error 404, please make sure you are authenticated and that `repo_id` and + `repo_type` are set correctly. If repo does not exist, create it first using + [`~hf_api.create_repo`]. + + + + Example: + + ```python + >>> from huggingface_hub import upload_file + + >>> with open("./local/filepath", "rb") as fobj: + ... upload_file( + ... path_or_fileobj=fileobj, + ... path_in_repo="remote/file/path.h5", + ... repo_id="username/my-dataset", + ... repo_type="dataset", + ... token="my_token", + ... ) + "https://huggingface.co/datasets/username/my-dataset/blob/main/remote/file/path.h5" + + >>> upload_file( + ... path_or_fileobj=".\\\\local\\\\file\\\\path", + ... path_in_repo="remote/file/path.h5", + ... repo_id="username/my-model", + ... token="my_token", + ... ) + "https://huggingface.co/username/my-model/blob/main/remote/file/path.h5" + + >>> upload_file( + ... path_or_fileobj=".\\\\local\\\\file\\\\path", + ... path_in_repo="remote/file/path.h5", + ... repo_id="username/my-model", + ... token="my_token", + ... create_pr=True, + ... ) + "https://huggingface.co/username/my-model/blob/refs%2Fpr%2F1/remote/file/path.h5" + ``` + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + + commit_message = ( + commit_message if commit_message is not None else f"Upload {path_in_repo} with huggingface_hub" + ) + operation = CommitOperationAdd( + path_or_fileobj=path_or_fileobj, + path_in_repo=path_in_repo, + ) + + commit_info = self.create_commit( + repo_id=repo_id, + repo_type=repo_type, + operations=[operation], + commit_message=commit_message, + commit_description=commit_description, + token=token, + revision=revision, + create_pr=create_pr, + parent_commit=parent_commit, + ) + + if commit_info.pr_url is not None: + revision = quote(_parse_revision_from_pr_url(commit_info.pr_url), safe="") + if repo_type in constants.REPO_TYPES_URL_PREFIXES: + repo_id = constants.REPO_TYPES_URL_PREFIXES[repo_type] + repo_id + revision = revision if revision is not None else constants.DEFAULT_REVISION + + return CommitInfo( + commit_url=commit_info.commit_url, + commit_message=commit_info.commit_message, + commit_description=commit_info.commit_description, + oid=commit_info.oid, + pr_url=commit_info.pr_url, + # Similar to `hf_hub_url` but it's "blob" instead of "resolve" + # TODO: remove this in v1.0 + _url=f"{self.endpoint}/{repo_id}/blob/{revision}/{path_in_repo}", + ) + + @overload + def upload_folder( # type: ignore + self, + *, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + run_as_future: Literal[False] = ..., + ) -> CommitInfo: ... + + @overload + def upload_folder( # type: ignore + self, + *, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + run_as_future: Literal[True] = ..., + ) -> Future[CommitInfo]: ... + + @validate_hf_hub_args + @future_compatible + def upload_folder( + self, + *, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + run_as_future: bool = False, + ) -> Union[CommitInfo, Future[CommitInfo]]: + """ + Upload a local folder to the given repo. The upload is done through a HTTP requests, and doesn't require git or + git-lfs to be installed. + + The structure of the folder will be preserved. Files with the same name already present in the repository will + be overwritten. Others will be left untouched. + + Use the `allow_patterns` and `ignore_patterns` arguments to specify which files to upload. These parameters + accept either a single pattern or a list of patterns. Patterns are Standard Wildcards (globbing patterns) as + documented [here](https://tldp.org/LDP/GNU-Linux-Tools-Summary/html/x11655.htm). If both `allow_patterns` and + `ignore_patterns` are provided, both constraints apply. By default, all files from the folder are uploaded. + + Use the `delete_patterns` argument to specify remote files you want to delete. Input type is the same as for + `allow_patterns` (see above). If `path_in_repo` is also provided, the patterns are matched against paths + relative to this folder. For example, `upload_folder(..., path_in_repo="experiment", delete_patterns="logs/*")` + will delete any remote file under `./experiment/logs/`. Note that the `.gitattributes` file will not be deleted + even if it matches the patterns. + + Any `.git/` folder present in any subdirectory will be ignored. However, please be aware that the `.gitignore` + file is not taken into account. + + Uses `HfApi.create_commit` under the hood. + + Args: + repo_id (`str`): + The repository to which the file will be uploaded, for example: + `"username/custom_transformers"` + folder_path (`str` or `Path`): + Path to the folder to upload on the local file system + path_in_repo (`str`, *optional*): + Relative path of the directory in the repo, for example: + `"checkpoints/1fec34a/results"`. Will default to the root folder of the repository. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit. Defaults to: + `f"Upload {path_in_repo} with huggingface_hub"` + commit_description (`str` *optional*): + The description of the generated commit + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. If `revision` is not + set, PR is opened against the `"main"` branch. If `revision` is set and is a branch, PR is opened + against this branch. If `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are uploaded. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not uploaded. + delete_patterns (`List[str]` or `str`, *optional*): + If provided, remote files matching any of the patterns will be deleted from the repo while committing + new files. This is useful if you don't know which files have already been uploaded. + Note: to avoid discrepancies the `.gitattributes` file is not deleted even if it matches the pattern. + run_as_future (`bool`, *optional*): + Whether or not to run this method in the background. Background jobs are run sequentially without + blocking the main thread. Passing `run_as_future=True` will return a [Future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) + object. Defaults to `False`. + + Returns: + [`CommitInfo`] or `Future`: + Instance of [`CommitInfo`] containing information about the newly created commit (commit hash, commit + url, pr url, commit message,...). If `run_as_future=True` is passed, returns a Future object which will + contain the result when executed. + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + + + + + + `upload_folder` assumes that the repo already exists on the Hub. If you get a Client error 404, please make + sure you are authenticated and that `repo_id` and `repo_type` are set correctly. If repo does not exist, create + it first using [`~hf_api.create_repo`]. + + + + + + When dealing with a large folder (thousands of files or hundreds of GB), we recommend using [`~hf_api.upload_large_folder`] instead. + + + + Example: + + ```python + # Upload checkpoints folder except the log files + >>> upload_folder( + ... folder_path="local/checkpoints", + ... path_in_repo="remote/experiment/checkpoints", + ... repo_id="username/my-dataset", + ... repo_type="datasets", + ... token="my_token", + ... ignore_patterns="**/logs/*.txt", + ... ) + # "https://huggingface.co/datasets/username/my-dataset/tree/main/remote/experiment/checkpoints" + + # Upload checkpoints folder including logs while deleting existing logs from the repo + # Useful if you don't know exactly which log files have already being pushed + >>> upload_folder( + ... folder_path="local/checkpoints", + ... path_in_repo="remote/experiment/checkpoints", + ... repo_id="username/my-dataset", + ... repo_type="datasets", + ... token="my_token", + ... delete_patterns="**/logs/*.txt", + ... ) + "https://huggingface.co/datasets/username/my-dataset/tree/main/remote/experiment/checkpoints" + + # Upload checkpoints folder while creating a PR + >>> upload_folder( + ... folder_path="local/checkpoints", + ... path_in_repo="remote/experiment/checkpoints", + ... repo_id="username/my-dataset", + ... repo_type="datasets", + ... token="my_token", + ... create_pr=True, + ... ) + "https://huggingface.co/datasets/username/my-dataset/tree/refs%2Fpr%2F1/remote/experiment/checkpoints" + + ``` + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + + # By default, upload folder to the root directory in repo. + if path_in_repo is None: + path_in_repo = "" + + # Do not upload .git folder + if ignore_patterns is None: + ignore_patterns = [] + elif isinstance(ignore_patterns, str): + ignore_patterns = [ignore_patterns] + ignore_patterns += DEFAULT_IGNORE_PATTERNS + + delete_operations = self._prepare_folder_deletions( + repo_id=repo_id, + repo_type=repo_type, + revision=constants.DEFAULT_REVISION if create_pr else revision, + token=token, + path_in_repo=path_in_repo, + delete_patterns=delete_patterns, + ) + add_operations = self._prepare_upload_folder_additions( + folder_path, + path_in_repo, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + token=token, + repo_type=repo_type, + ) + + # Optimize operations: if some files will be overwritten, we don't need to delete them first + if len(add_operations) > 0: + added_paths = set(op.path_in_repo for op in add_operations) + delete_operations = [ + delete_op for delete_op in delete_operations if delete_op.path_in_repo not in added_paths + ] + commit_operations = delete_operations + add_operations + + commit_message = commit_message or "Upload folder using huggingface_hub" + + commit_info = self.create_commit( + repo_type=repo_type, + repo_id=repo_id, + operations=commit_operations, + commit_message=commit_message, + commit_description=commit_description, + token=token, + revision=revision, + create_pr=create_pr, + parent_commit=parent_commit, + ) + + # Create url to uploaded folder (for legacy return value) + if create_pr and commit_info.pr_url is not None: + revision = quote(_parse_revision_from_pr_url(commit_info.pr_url), safe="") + if repo_type in constants.REPO_TYPES_URL_PREFIXES: + repo_id = constants.REPO_TYPES_URL_PREFIXES[repo_type] + repo_id + revision = revision if revision is not None else constants.DEFAULT_REVISION + + return CommitInfo( + commit_url=commit_info.commit_url, + commit_message=commit_info.commit_message, + commit_description=commit_info.commit_description, + oid=commit_info.oid, + pr_url=commit_info.pr_url, + # Similar to `hf_hub_url` but it's "tree" instead of "resolve" + # TODO: remove this in v1.0 + _url=f"{self.endpoint}/{repo_id}/tree/{revision}/{path_in_repo}", + ) + + @validate_hf_hub_args + def delete_file( + self, + path_in_repo: str, + repo_id: str, + *, + token: Union[str, bool, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + ) -> CommitInfo: + """ + Deletes a file in the given repo. + + Args: + path_in_repo (`str`): + Relative filepath in the repo, for example: + `"checkpoints/1fec34a/weights.bin"` + repo_id (`str`): + The repository from which the file will be deleted, for example: + `"username/custom_transformers"` + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if the file is in a dataset or + space, `None` or `"model"` if in a model. Default is `None`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit. Defaults to + `f"Delete {path_in_repo} with huggingface_hub"`. + commit_description (`str` *optional*) + The description of the generated commit + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. + If `revision` is not set, PR is opened against the `"main"` branch. If + `revision` is set and is a branch, PR is opened against this branch. If + `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + - [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + - [`~utils.EntryNotFoundError`] + If the file to download cannot be found. + + + + """ + commit_message = ( + commit_message if commit_message is not None else f"Delete {path_in_repo} with huggingface_hub" + ) + + operations = [CommitOperationDelete(path_in_repo=path_in_repo)] + + return self.create_commit( + repo_id=repo_id, + repo_type=repo_type, + token=token, + operations=operations, + revision=revision, + commit_message=commit_message, + commit_description=commit_description, + create_pr=create_pr, + parent_commit=parent_commit, + ) + + @validate_hf_hub_args + def delete_files( + self, + repo_id: str, + delete_patterns: List[str], + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + ) -> CommitInfo: + """ + Delete files from a repository on the Hub. + + If a folder path is provided, the entire folder is deleted as well as + all files it contained. + + Args: + repo_id (`str`): + The repository from which the folder will be deleted, for example: + `"username/custom_transformers"` + delete_patterns (`List[str]`): + List of files or folders to delete. Each string can either be + a file path, a folder path or a Unix shell-style wildcard. + E.g. `["file.txt", "folder/", "data/*.parquet"]` + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + to the stored token. + repo_type (`str`, *optional*): + Type of the repo to delete files from. Can be `"model"`, + `"dataset"` or `"space"`. Defaults to `"model"`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The summary (first line) of the generated commit. Defaults to + `f"Delete files using huggingface_hub"`. + commit_description (`str` *optional*) + The description of the generated commit. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. + If `revision` is not set, PR is opened against the `"main"` branch. If + `revision` is set and is a branch, PR is opened against this branch. If + `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + """ + operations = self._prepare_folder_deletions( + repo_id=repo_id, repo_type=repo_type, delete_patterns=delete_patterns, path_in_repo="", revision=revision + ) + + if commit_message is None: + commit_message = f"Delete files {' '.join(delete_patterns)} with huggingface_hub" + + return self.create_commit( + repo_id=repo_id, + repo_type=repo_type, + token=token, + operations=operations, + revision=revision, + commit_message=commit_message, + commit_description=commit_description, + create_pr=create_pr, + parent_commit=parent_commit, + ) + + @validate_hf_hub_args + def delete_folder( + self, + path_in_repo: str, + repo_id: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + ) -> CommitInfo: + """ + Deletes a folder in the given repo. + + Simple wrapper around [`create_commit`] method. + + Args: + path_in_repo (`str`): + Relative folder path in the repo, for example: `"checkpoints/1fec34a"`. + repo_id (`str`): + The repository from which the folder will be deleted, for example: + `"username/custom_transformers"` + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + to the stored token. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if the folder is in a dataset or + space, `None` or `"model"` if in a model. Default is `None`. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit. Defaults to + `f"Delete folder {path_in_repo} with huggingface_hub"`. + commit_description (`str` *optional*) + The description of the generated commit. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request with that commit. Defaults to `False`. + If `revision` is not set, PR is opened against the `"main"` branch. If + `revision` is set and is a branch, PR is opened against this branch. If + `revision` is set and is not a branch name (example: a commit oid), an + `RevisionNotFoundError` is returned by the server. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + """ + return self.create_commit( + repo_id=repo_id, + repo_type=repo_type, + token=token, + operations=[CommitOperationDelete(path_in_repo=path_in_repo, is_folder=True)], + revision=revision, + commit_message=( + commit_message if commit_message is not None else f"Delete folder {path_in_repo} with huggingface_hub" + ), + commit_description=commit_description, + create_pr=create_pr, + parent_commit=parent_commit, + ) + + def upload_large_folder( + self, + repo_id: str, + folder_path: Union[str, Path], + *, + repo_type: str, # Repo type is required! + revision: Optional[str] = None, + private: Optional[bool] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + num_workers: Optional[int] = None, + print_report: bool = True, + print_report_every: int = 60, + ) -> None: + """Upload a large folder to the Hub in the most resilient way possible. + + Several workers are started to upload files in an optimized way. Before being committed to a repo, files must be + hashed and be pre-uploaded if they are LFS files. Workers will perform these tasks for each file in the folder. + At each step, some metadata information about the upload process is saved in the folder under `.cache/.huggingface/` + to be able to resume the process if interrupted. The whole process might result in several commits. + + Args: + repo_id (`str`): + The repository to which the file will be uploaded. + E.g. `"HuggingFaceTB/smollm-corpus"`. + folder_path (`str` or `Path`): + Path to the folder to upload on the local file system. + repo_type (`str`): + Type of the repository. Must be one of `"model"`, `"dataset"` or `"space"`. + Unlike in all other `HfApi` methods, `repo_type` is explicitly required here. This is to avoid + any mistake when uploading a large folder to the Hub, and therefore prevent from having to re-upload + everything. + revision (`str`, `optional`): + The branch to commit to. If not provided, the `main` branch will be used. + private (`bool`, `optional`): + Whether the repository should be private. + If `None` (default), the repo will be public unless the organization's default is private. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are uploaded. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not uploaded. + num_workers (`int`, *optional*): + Number of workers to start. Defaults to `os.cpu_count() - 2` (minimum 2). + A higher number of workers may speed up the process if your machine allows it. However, on machines with a + slower connection, it is recommended to keep the number of workers low to ensure better resumability. + Indeed, partially uploaded files will have to be completely re-uploaded if the process is interrupted. + print_report (`bool`, *optional*): + Whether to print a report of the upload progress. Defaults to True. + Report is printed to `sys.stdout` every X seconds (60 by defaults) and overwrites the previous report. + print_report_every (`int`, *optional*): + Frequency at which the report is printed. Defaults to 60 seconds. + + + + A few things to keep in mind: + - Repository limits still apply: https://huggingface.co/docs/hub/repositories-recommendations + - Do not start several processes in parallel. + - You can interrupt and resume the process at any time. + - Do not upload the same folder to several repositories. If you need to do so, you must delete the local `.cache/.huggingface/` folder first. + + + + + + While being much more robust to upload large folders, `upload_large_folder` is more limited than [`upload_folder`] feature-wise. In practice: + - you cannot set a custom `path_in_repo`. If you want to upload to a subfolder, you need to set the proper structure locally. + - you cannot set a custom `commit_message` and `commit_description` since multiple commits are created. + - you cannot delete from the repo while uploading. Please make a separate commit first. + - you cannot create a PR directly. Please create a PR first (from the UI or using [`create_pull_request`]) and then commit to it by passing `revision`. + + + + **Technical details:** + + `upload_large_folder` process is as follow: + 1. (Check parameters and setup.) + 2. Create repo if missing. + 3. List local files to upload. + 4. Start workers. Workers can perform the following tasks: + - Hash a file. + - Get upload mode (regular or LFS) for a list of files. + - Pre-upload an LFS file. + - Commit a bunch of files. + Once a worker finishes a task, it will move on to the next task based on the priority list (see below) until + all files are uploaded and committed. + 5. While workers are up, regularly print a report to sys.stdout. + + Order of priority: + 1. Commit if more than 5 minutes since last commit attempt (and at least 1 file). + 2. Commit if at least 150 files are ready to commit. + 3. Get upload mode if at least 10 files have been hashed. + 4. Pre-upload LFS file if at least 1 file and no worker is pre-uploading. + 5. Hash file if at least 1 file and no worker is hashing. + 6. Get upload mode if at least 1 file and no worker is getting upload mode. + 7. Pre-upload LFS file if at least 1 file (exception: if hf_transfer is enabled, only 1 worker can preupload LFS at a time). + 8. Hash file if at least 1 file to hash. + 9. Get upload mode if at least 1 file to get upload mode. + 10. Commit if at least 1 file to commit and at least 1 min since last commit attempt. + 11. Commit if at least 1 file to commit and all other queues are empty. + + Special rules: + - If `hf_transfer` is enabled, only 1 LFS uploader at a time. Otherwise the CPU would be bloated by `hf_transfer`. + - Only one worker can commit at a time. + - If no tasks are available, the worker waits for 10 seconds before checking again. + """ + return upload_large_folder_internal( + self, + repo_id=repo_id, + folder_path=folder_path, + repo_type=repo_type, + revision=revision, + private=private, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + num_workers=num_workers, + print_report=print_report, + print_report_every=print_report_every, + ) + + @validate_hf_hub_args + def get_hf_file_metadata( + self, + *, + url: str, + token: Union[bool, str, None] = None, + proxies: Optional[Dict] = None, + timeout: Optional[float] = constants.DEFAULT_REQUEST_TIMEOUT, + ) -> HfFileMetadata: + """Fetch metadata of a file versioned on the Hub for a given url. + + Args: + url (`str`): + File url, for example returned by [`hf_hub_url`]. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to `requests.request`. + timeout (`float`, *optional*, defaults to 10): + How many seconds to wait for the server to send metadata before giving up. + + Returns: + A [`HfFileMetadata`] object containing metadata such as location, etag, size and commit_hash. + """ + if token is None: + # Cannot do `token = token or self.token` as token can be `False`. + token = self.token + + return get_hf_file_metadata( + url=url, + token=token, + proxies=proxies, + timeout=timeout, + library_name=self.library_name, + library_version=self.library_version, + user_agent=self.user_agent, + ) + + @validate_hf_hub_args + def hf_hub_download( + self, + repo_id: str, + filename: str, + *, + subfolder: Optional[str] = None, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + cache_dir: Union[str, Path, None] = None, + local_dir: Union[str, Path, None] = None, + force_download: bool = False, + proxies: Optional[Dict] = None, + etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT, + token: Union[bool, str, None] = None, + local_files_only: bool = False, + # Deprecated args + resume_download: Optional[bool] = None, + force_filename: Optional[str] = None, + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + ) -> str: + """Download a given file if it's not already present in the local cache. + + The new cache file layout looks like this: + - The cache directory contains one subfolder per repo_id (namespaced by repo type) + - inside each repo folder: + - refs is a list of the latest known revision => commit_hash pairs + - blobs contains the actual file blobs (identified by their git-sha or sha256, depending on + whether they're LFS files or not) + - snapshots contains one subfolder per commit, each "commit" contains the subset of the files + that have been resolved at that particular commit. Each filename is a symlink to the blob + at that particular commit. + + ``` + [ 96] . + └── [ 160] models--julien-c--EsperBERTo-small + ├── [ 160] blobs + │ ├── [321M] 403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + │ ├── [ 398] 7cb18dc9bafbfcf74629a4b760af1b160957a83e + │ └── [1.4K] d7edf6bd2a681fb0175f7735299831ee1b22b812 + ├── [ 96] refs + │ └── [ 40] main + └── [ 128] snapshots + ├── [ 128] 2439f60ef33a0d46d85da5001d52aeda5b00ce9f + │ ├── [ 52] README.md -> ../../blobs/d7edf6bd2a681fb0175f7735299831ee1b22b812 + │ └── [ 76] pytorch_model.bin -> ../../blobs/403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + └── [ 128] bbc77c8132af1cc5cf678da3f1ddf2de43606d48 + ├── [ 52] README.md -> ../../blobs/7cb18dc9bafbfcf74629a4b760af1b160957a83e + └── [ 76] pytorch_model.bin -> ../../blobs/403450e234d65943a7dcf7e05a771ce3c92faa84dd07db4ac20f592037a1e4bd + ``` + + If `local_dir` is provided, the file structure from the repo will be replicated in this location. When using this + option, the `cache_dir` will not be used and a `.cache/huggingface/` folder will be created at the root of `local_dir` + to store some metadata related to the downloaded files. While this mechanism is not as robust as the main + cache-system, it's optimized for regularly pulling the latest version of a repository. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + filename (`str`): + The name of the file in the repo. + subfolder (`str`, *optional*): + An optional value corresponding to a folder inside the repository. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if downloading from a dataset or space, + `None` or `"model"` if downloading from a model. Default is `None`. + revision (`str`, *optional*): + An optional Git revision id which can be a branch name, a tag, or a + commit hash. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_dir (`str` or `Path`, *optional*): + If provided, the downloaded file will be placed under this directory. + force_download (`bool`, *optional*, defaults to `False`): + Whether the file should be downloaded even if it already exists in + the local cache. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to + `requests.request`. + etag_timeout (`float`, *optional*, defaults to `10`): + When fetching ETag, how many seconds to wait for the server to send + data before giving up which is passed to `requests.request`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the + local cached file if it exists. + + Returns: + `str`: Local path of file or if networking is off, last version of file cached on disk. + + Raises: + [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + [`~utils.EntryNotFoundError`] + If the file to download cannot be found. + [`~utils.LocalEntryNotFoundError`] + If network is disabled or unavailable and file is not found in cache. + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `token=True` but the token cannot be found. + [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) + If ETag cannot be determined. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If some parameter value is invalid. + """ + from .file_download import hf_hub_download + + if token is None: + # Cannot do `token = token or self.token` as token can be `False`. + token = self.token + + return hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + repo_type=repo_type, + revision=revision, + endpoint=self.endpoint, + library_name=self.library_name, + library_version=self.library_version, + cache_dir=cache_dir, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + user_agent=self.user_agent, + force_download=force_download, + force_filename=force_filename, + proxies=proxies, + etag_timeout=etag_timeout, + resume_download=resume_download, + token=token, + headers=self.headers, + local_files_only=local_files_only, + ) + + @validate_hf_hub_args + def snapshot_download( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + cache_dir: Union[str, Path, None] = None, + local_dir: Union[str, Path, None] = None, + proxies: Optional[Dict] = None, + etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT, + force_download: bool = False, + token: Union[bool, str, None] = None, + local_files_only: bool = False, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + max_workers: int = 8, + tqdm_class: Optional[base_tqdm] = None, + # Deprecated args + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + resume_download: Optional[bool] = None, + ) -> str: + """Download repo files. + + Download a whole snapshot of a repo's files at the specified revision. This is useful when you want all files from + a repo, because you don't know which ones you will need a priori. All files are nested inside a folder in order + to keep their actual filename relative to that folder. You can also filter which files to download using + `allow_patterns` and `ignore_patterns`. + + If `local_dir` is provided, the file structure from the repo will be replicated in this location. When using this + option, the `cache_dir` will not be used and a `.cache/huggingface/` folder will be created at the root of `local_dir` + to store some metadata related to the downloaded files.While this mechanism is not as robust as the main + cache-system, it's optimized for regularly pulling the latest version of a repository. + + An alternative would be to clone the repo but this requires git and git-lfs to be installed and properly + configured. It is also not possible to filter which files to download when cloning a repository using git. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if downloading from a dataset or space, + `None` or `"model"` if downloading from a model. Default is `None`. + revision (`str`, *optional*): + An optional Git revision id which can be a branch name, a tag, or a + commit hash. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_dir (`str` or `Path`, *optional*): + If provided, the downloaded files will be placed under this directory. + proxies (`dict`, *optional*): + Dictionary mapping protocol to the URL of the proxy passed to + `requests.request`. + etag_timeout (`float`, *optional*, defaults to `10`): + When fetching ETag, how many seconds to wait for the server to send + data before giving up which is passed to `requests.request`. + force_download (`bool`, *optional*, defaults to `False`): + Whether the file should be downloaded even if it already exists in the local cache. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the + local cached file if it exists. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are downloaded. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not downloaded. + max_workers (`int`, *optional*): + Number of concurrent threads to download files (1 thread = 1 file download). + Defaults to 8. + tqdm_class (`tqdm`, *optional*): + If provided, overwrites the default behavior for the progress bar. Passed + argument must inherit from `tqdm.auto.tqdm` or at least mimic its behavior. + Note that the `tqdm_class` is not passed to each individual download. + Defaults to the custom HF progress bar that can be disabled by setting + `HF_HUB_DISABLE_PROGRESS_BARS` environment variable. + + Returns: + `str`: folder path of the repo snapshot. + + Raises: + [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + [`~utils.RevisionNotFoundError`] + If the revision to download from cannot be found. + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `token=True` and the token cannot be found. + [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if + ETag cannot be determined. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid. + """ + from ._snapshot_download import snapshot_download + + if token is None: + # Cannot do `token = token or self.token` as token can be `False`. + token = self.token + + return snapshot_download( + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + endpoint=self.endpoint, + cache_dir=cache_dir, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + library_name=self.library_name, + library_version=self.library_version, + user_agent=self.user_agent, + proxies=proxies, + etag_timeout=etag_timeout, + resume_download=resume_download, + force_download=force_download, + token=token, + local_files_only=local_files_only, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + max_workers=max_workers, + tqdm_class=tqdm_class, + ) + + def get_safetensors_metadata( + self, + repo_id: str, + *, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> SafetensorsRepoMetadata: + """ + Parse metadata for a safetensors repo on the Hub. + + We first check if the repo has a single safetensors file or a sharded safetensors repo. If it's a single + safetensors file, we parse the metadata from this file. If it's a sharded safetensors repo, we parse the + metadata from the index file and then parse the metadata from each shard. + + To parse metadata from a single safetensors file, use [`parse_safetensors_file_metadata`]. + + For more details regarding the safetensors format, check out https://huggingface.co/docs/safetensors/index#format. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if the file is in a dataset or space, `None` or `"model"` if in a + model. Default is `None`. + revision (`str`, *optional*): + The git revision to fetch the file from. Can be a branch name, a tag, or a commit hash. Defaults to the + head of the `"main"` branch. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`SafetensorsRepoMetadata`]: information related to safetensors repo. + + Raises: + [`NotASafetensorsRepoError`] + If the repo is not a safetensors repo i.e. doesn't have either a + `model.safetensors` or a `model.safetensors.index.json` file. + [`SafetensorsParsingError`] + If a safetensors file header couldn't be parsed correctly. + + Example: + ```py + # Parse repo with single weights file + >>> metadata = get_safetensors_metadata("bigscience/bloomz-560m") + >>> metadata + SafetensorsRepoMetadata( + metadata=None, + sharded=False, + weight_map={'h.0.input_layernorm.bias': 'model.safetensors', ...}, + files_metadata={'model.safetensors': SafetensorsFileMetadata(...)} + ) + >>> metadata.files_metadata["model.safetensors"].metadata + {'format': 'pt'} + + # Parse repo with sharded model + >>> metadata = get_safetensors_metadata("bigscience/bloom") + Parse safetensors files: 100%|██████████████████████████████████████████| 72/72 [00:12<00:00, 5.78it/s] + >>> metadata + SafetensorsRepoMetadata(metadata={'total_size': 352494542848}, sharded=True, weight_map={...}, files_metadata={...}) + >>> len(metadata.files_metadata) + 72 # All safetensors files have been fetched + + # Parse repo with sharded model + >>> get_safetensors_metadata("runwayml/stable-diffusion-v1-5") + NotASafetensorsRepoError: 'runwayml/stable-diffusion-v1-5' is not a safetensors repo. Couldn't find 'model.safetensors.index.json' or 'model.safetensors' files. + ``` + """ + if self.file_exists( # Single safetensors file => non-sharded model + repo_id=repo_id, + filename=constants.SAFETENSORS_SINGLE_FILE, + repo_type=repo_type, + revision=revision, + token=token, + ): + file_metadata = self.parse_safetensors_file_metadata( + repo_id=repo_id, + filename=constants.SAFETENSORS_SINGLE_FILE, + repo_type=repo_type, + revision=revision, + token=token, + ) + return SafetensorsRepoMetadata( + metadata=None, + sharded=False, + weight_map={ + tensor_name: constants.SAFETENSORS_SINGLE_FILE for tensor_name in file_metadata.tensors.keys() + }, + files_metadata={constants.SAFETENSORS_SINGLE_FILE: file_metadata}, + ) + elif self.file_exists( # Multiple safetensors files => sharded with index + repo_id=repo_id, + filename=constants.SAFETENSORS_INDEX_FILE, + repo_type=repo_type, + revision=revision, + token=token, + ): + # Fetch index + index_file = self.hf_hub_download( + repo_id=repo_id, + filename=constants.SAFETENSORS_INDEX_FILE, + repo_type=repo_type, + revision=revision, + token=token, + ) + with open(index_file) as f: + index = json.load(f) + + weight_map = index.get("weight_map", {}) + + # Fetch metadata per shard + files_metadata = {} + + def _parse(filename: str) -> None: + files_metadata[filename] = self.parse_safetensors_file_metadata( + repo_id=repo_id, filename=filename, repo_type=repo_type, revision=revision, token=token + ) + + thread_map( + _parse, + set(weight_map.values()), + desc="Parse safetensors files", + tqdm_class=hf_tqdm, + ) + + return SafetensorsRepoMetadata( + metadata=index.get("metadata", None), + sharded=True, + weight_map=weight_map, + files_metadata=files_metadata, + ) + else: + # Not a safetensors repo + raise NotASafetensorsRepoError( + f"'{repo_id}' is not a safetensors repo. Couldn't find '{constants.SAFETENSORS_INDEX_FILE}' or '{constants.SAFETENSORS_SINGLE_FILE}' files." + ) + + def parse_safetensors_file_metadata( + self, + repo_id: str, + filename: str, + *, + repo_type: Optional[str] = None, + revision: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> SafetensorsFileMetadata: + """ + Parse metadata from a safetensors file on the Hub. + + To parse metadata from all safetensors files in a repo at once, use [`get_safetensors_metadata`]. + + For more details regarding the safetensors format, check out https://huggingface.co/docs/safetensors/index#format. + + Args: + repo_id (`str`): + A user or an organization name and a repo name separated by a `/`. + filename (`str`): + The name of the file in the repo. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if the file is in a dataset or space, `None` or `"model"` if in a + model. Default is `None`. + revision (`str`, *optional*): + The git revision to fetch the file from. Can be a branch name, a tag, or a commit hash. Defaults to the + head of the `"main"` branch. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`SafetensorsFileMetadata`]: information related to a safetensors file. + + Raises: + [`NotASafetensorsRepoError`]: + If the repo is not a safetensors repo i.e. doesn't have either a + `model.safetensors` or a `model.safetensors.index.json` file. + [`SafetensorsParsingError`]: + If a safetensors file header couldn't be parsed correctly. + """ + url = hf_hub_url( + repo_id=repo_id, filename=filename, repo_type=repo_type, revision=revision, endpoint=self.endpoint + ) + _headers = self._build_hf_headers(token=token) + + # 1. Fetch first 100kb + # Empirically, 97% of safetensors files have a metadata size < 100kb (over the top 1000 models on the Hub). + # We assume fetching 100kb is faster than making 2 GET requests. Therefore we always fetch the first 100kb to + # avoid the 2nd GET in most cases. + # See https://github.com/huggingface/huggingface_hub/pull/1855#discussion_r1404286419. + response = get_session().get(url, headers={**_headers, "range": "bytes=0-100000"}) + hf_raise_for_status(response) + + # 2. Parse metadata size + metadata_size = struct.unpack(" constants.SAFETENSORS_MAX_HEADER_LENGTH: + raise SafetensorsParsingError( + f"Failed to parse safetensors header for '{filename}' (repo '{repo_id}', revision " + f"'{revision or constants.DEFAULT_REVISION}'): safetensors header is too big. Maximum supported size is " + f"{constants.SAFETENSORS_MAX_HEADER_LENGTH} bytes (got {metadata_size})." + ) + + # 3.a. Get metadata from payload + if metadata_size <= 100000: + metadata_as_bytes = response.content[8 : 8 + metadata_size] + else: # 3.b. Request full metadata + response = get_session().get(url, headers={**_headers, "range": f"bytes=8-{metadata_size + 7}"}) + hf_raise_for_status(response) + metadata_as_bytes = response.content + + # 4. Parse json header + try: + metadata_as_dict = json.loads(metadata_as_bytes.decode(errors="ignore")) + except json.JSONDecodeError as e: + raise SafetensorsParsingError( + f"Failed to parse safetensors header for '{filename}' (repo '{repo_id}', revision " + f"'{revision or constants.DEFAULT_REVISION}'): header is not json-encoded string. Please make sure this is a " + "correctly formatted safetensors file." + ) from e + + try: + return SafetensorsFileMetadata( + metadata=metadata_as_dict.get("__metadata__", {}), + tensors={ + key: TensorInfo( + dtype=tensor["dtype"], + shape=tensor["shape"], + data_offsets=tuple(tensor["data_offsets"]), # type: ignore + ) + for key, tensor in metadata_as_dict.items() + if key != "__metadata__" + }, + ) + except (KeyError, IndexError) as e: + raise SafetensorsParsingError( + f"Failed to parse safetensors header for '{filename}' (repo '{repo_id}', revision " + f"'{revision or constants.DEFAULT_REVISION}'): header format not recognized. Please make sure this is a correctly" + " formatted safetensors file." + ) from e + + @validate_hf_hub_args + def create_branch( + self, + repo_id: str, + *, + branch: str, + revision: Optional[str] = None, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + exist_ok: bool = False, + ) -> None: + """ + Create a new branch for a repo on the Hub, starting from the specified revision (defaults to `main`). + To find a revision suiting your needs, you can use [`list_repo_refs`] or [`list_repo_commits`]. + + Args: + repo_id (`str`): + The repository in which the branch will be created. + Example: `"user/my-cool-model"`. + + branch (`str`): + The name of the branch to create. + + revision (`str`, *optional*): + The git revision to create the branch from. It can be a branch name or + the OID/SHA of a commit, as a hexadecimal string. Defaults to the head + of the `"main"` branch. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if creating a branch on a dataset or + space, `None` or `"model"` if tagging a model. Default is `None`. + + exist_ok (`bool`, *optional*, defaults to `False`): + If `True`, do not raise an error if branch already exists. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + [`~utils.BadRequestError`]: + If invalid reference for a branch. Ex: `refs/pr/5` or 'refs/foo/bar'. + [`~utils.HfHubHTTPError`]: + If the branch already exists on the repo (error 409) and `exist_ok` is + set to `False`. + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + branch = quote(branch, safe="") + + # Prepare request + branch_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/branch/{branch}" + headers = self._build_hf_headers(token=token) + payload = {} + if revision is not None: + payload["startingPoint"] = revision + + # Create branch + response = get_session().post(url=branch_url, headers=headers, json=payload) + try: + hf_raise_for_status(response) + except HfHubHTTPError as e: + if exist_ok and e.response.status_code == 409: + return + elif exist_ok and e.response.status_code == 403: + # No write permission on the namespace but branch might already exist + try: + refs = self.list_repo_refs(repo_id=repo_id, repo_type=repo_type, token=token) + for branch_ref in refs.branches: + if branch_ref.name == branch: + return # Branch already exists => do not raise + except HfHubHTTPError: + pass # We raise the original error if the branch does not exist + raise + + @validate_hf_hub_args + def delete_branch( + self, + repo_id: str, + *, + branch: str, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> None: + """ + Delete a branch from a repo on the Hub. + + Args: + repo_id (`str`): + The repository in which a branch will be deleted. + Example: `"user/my-cool-model"`. + + branch (`str`): + The name of the branch to delete. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if creating a branch on a dataset or + space, `None` or `"model"` if tagging a model. Default is `None`. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + [`~utils.HfHubHTTPError`]: + If trying to delete a protected branch. Ex: `main` cannot be deleted. + [`~utils.HfHubHTTPError`]: + If trying to delete a branch that does not exist. + + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + branch = quote(branch, safe="") + + # Prepare request + branch_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/branch/{branch}" + headers = self._build_hf_headers(token=token) + + # Delete branch + response = get_session().delete(url=branch_url, headers=headers) + hf_raise_for_status(response) + + @validate_hf_hub_args + def create_tag( + self, + repo_id: str, + *, + tag: str, + tag_message: Optional[str] = None, + revision: Optional[str] = None, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + exist_ok: bool = False, + ) -> None: + """ + Tag a given commit of a repo on the Hub. + + Args: + repo_id (`str`): + The repository in which a commit will be tagged. + Example: `"user/my-cool-model"`. + + tag (`str`): + The name of the tag to create. + + tag_message (`str`, *optional*): + The description of the tag to create. + + revision (`str`, *optional*): + The git revision to tag. It can be a branch name or the OID/SHA of a + commit, as a hexadecimal string. Shorthands (7 first characters) are + also supported. Defaults to the head of the `"main"` branch. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if tagging a dataset or + space, `None` or `"model"` if tagging a model. Default is + `None`. + + exist_ok (`bool`, *optional*, defaults to `False`): + If `True`, do not raise an error if tag already exists. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + [`~utils.RevisionNotFoundError`]: + If revision is not found (error 404) on the repo. + [`~utils.HfHubHTTPError`]: + If the branch already exists on the repo (error 409) and `exist_ok` is + set to `False`. + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + revision = quote(revision, safe="") if revision is not None else constants.DEFAULT_REVISION + + # Prepare request + tag_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/tag/{revision}" + headers = self._build_hf_headers(token=token) + payload = {"tag": tag} + if tag_message is not None: + payload["message"] = tag_message + + # Tag + response = get_session().post(url=tag_url, headers=headers, json=payload) + try: + hf_raise_for_status(response) + except HfHubHTTPError as e: + if not (e.response.status_code == 409 and exist_ok): + raise + + @validate_hf_hub_args + def delete_tag( + self, + repo_id: str, + *, + tag: str, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> None: + """ + Delete a tag from a repo on the Hub. + + Args: + repo_id (`str`): + The repository in which a tag will be deleted. + Example: `"user/my-cool-model"`. + + tag (`str`): + The name of the tag to delete. + + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if tagging a dataset or space, `None` or + `"model"` if tagging a model. Default is `None`. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If repository is not found (error 404): wrong repo_id/repo_type, private + but not authenticated or repo does not exist. + [`~utils.RevisionNotFoundError`]: + If tag is not found. + """ + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + tag = quote(tag, safe="") + + # Prepare request + tag_url = f"{self.endpoint}/api/{repo_type}s/{repo_id}/tag/{tag}" + headers = self._build_hf_headers(token=token) + + # Un-tag + response = get_session().delete(url=tag_url, headers=headers) + hf_raise_for_status(response) + + @validate_hf_hub_args + def get_full_repo_name( + self, + model_id: str, + *, + organization: Optional[str] = None, + token: Union[bool, str, None] = None, + ): + """ + Returns the repository name for a given model ID and optional + organization. + + Args: + model_id (`str`): + The name of the model. + organization (`str`, *optional*): + If passed, the repository name will be in the organization + namespace instead of the user namespace. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `str`: The repository name in the user's namespace + ({username}/{model_id}) if no organization is passed, and under the + organization namespace ({organization}/{model_id}) otherwise. + """ + if organization is None: + if "/" in model_id: + username = model_id.split("/")[0] + else: + username = self.whoami(token=token)["name"] # type: ignore + return f"{username}/{model_id}" + else: + return f"{organization}/{model_id}" + + @validate_hf_hub_args + def get_repo_discussions( + self, + repo_id: str, + *, + author: Optional[str] = None, + discussion_type: Optional[constants.DiscussionTypeFilter] = None, + discussion_status: Optional[constants.DiscussionStatusFilter] = None, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Iterator[Discussion]: + """ + Fetches Discussions and Pull Requests for the given repo. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + author (`str`, *optional*): + Pass a value to filter by discussion author. `None` means no filter. + Default is `None`. + discussion_type (`str`, *optional*): + Set to `"pull_request"` to fetch only pull requests, `"discussion"` + to fetch only discussions. Set to `"all"` or `None` to fetch both. + Default is `None`. + discussion_status (`str`, *optional*): + Set to `"open"` (respectively `"closed"`) to fetch only open + (respectively closed) discussions. Set to `"all"` or `None` + to fetch both. + Default is `None`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if fetching from a dataset or + space, `None` or `"model"` if fetching from a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterator[Discussion]`: An iterator of [`Discussion`] objects. + + Example: + Collecting all discussions of a repo in a list: + + ```python + >>> from huggingface_hub import get_repo_discussions + >>> discussions_list = list(get_repo_discussions(repo_id="bert-base-uncased")) + ``` + + Iterating over discussions of a repo: + + ```python + >>> from huggingface_hub import get_repo_discussions + >>> for discussion in get_repo_discussions(repo_id="bert-base-uncased"): + ... print(discussion.num, discussion.title) + ``` + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + if discussion_type is not None and discussion_type not in constants.DISCUSSION_TYPES: + raise ValueError(f"Invalid discussion_type, must be one of {constants.DISCUSSION_TYPES}") + if discussion_status is not None and discussion_status not in constants.DISCUSSION_STATUS: + raise ValueError(f"Invalid discussion_status, must be one of {constants.DISCUSSION_STATUS}") + + headers = self._build_hf_headers(token=token) + path = f"{self.endpoint}/api/{repo_type}s/{repo_id}/discussions" + + params: Dict[str, Union[str, int]] = {} + if discussion_type is not None: + params["type"] = discussion_type + if discussion_status is not None: + params["status"] = discussion_status + if author is not None: + params["author"] = author + + def _fetch_discussion_page(page_index: int): + params["p"] = page_index + resp = get_session().get(path, headers=headers, params=params) + hf_raise_for_status(resp) + paginated_discussions = resp.json() + total = paginated_discussions["count"] + start = paginated_discussions["start"] + discussions = paginated_discussions["discussions"] + has_next = (start + len(discussions)) < total + return discussions, has_next + + has_next, page_index = True, 0 + + while has_next: + discussions, has_next = _fetch_discussion_page(page_index=page_index) + for discussion in discussions: + yield Discussion( + title=discussion["title"], + num=discussion["num"], + author=discussion.get("author", {}).get("name", "deleted"), + created_at=parse_datetime(discussion["createdAt"]), + status=discussion["status"], + repo_id=discussion["repo"]["name"], + repo_type=discussion["repo"]["type"], + is_pull_request=discussion["isPullRequest"], + endpoint=self.endpoint, + ) + page_index = page_index + 1 + + @validate_hf_hub_args + def get_discussion_details( + self, + repo_id: str, + discussion_num: int, + *, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> DiscussionWithDetails: + """Fetches a Discussion's / Pull Request 's details from the Hub. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: [`DiscussionWithDetails`] + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + if not isinstance(discussion_num, int) or discussion_num <= 0: + raise ValueError("Invalid discussion_num, must be a positive integer") + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + + path = f"{self.endpoint}/api/{repo_type}s/{repo_id}/discussions/{discussion_num}" + headers = self._build_hf_headers(token=token) + resp = get_session().get(path, params={"diff": "1"}, headers=headers) + hf_raise_for_status(resp) + + discussion_details = resp.json() + is_pull_request = discussion_details["isPullRequest"] + + target_branch = discussion_details["changes"]["base"] if is_pull_request else None + conflicting_files = discussion_details["filesWithConflicts"] if is_pull_request else None + merge_commit_oid = discussion_details["changes"].get("mergeCommitId", None) if is_pull_request else None + + return DiscussionWithDetails( + title=discussion_details["title"], + num=discussion_details["num"], + author=discussion_details.get("author", {}).get("name", "deleted"), + created_at=parse_datetime(discussion_details["createdAt"]), + status=discussion_details["status"], + repo_id=discussion_details["repo"]["name"], + repo_type=discussion_details["repo"]["type"], + is_pull_request=discussion_details["isPullRequest"], + events=[deserialize_event(evt) for evt in discussion_details["events"]], + conflicting_files=conflicting_files, + target_branch=target_branch, + merge_commit_oid=merge_commit_oid, + diff=discussion_details.get("diff"), + endpoint=self.endpoint, + ) + + @validate_hf_hub_args + def create_discussion( + self, + repo_id: str, + title: str, + *, + token: Union[bool, str, None] = None, + description: Optional[str] = None, + repo_type: Optional[str] = None, + pull_request: bool = False, + ) -> DiscussionWithDetails: + """Creates a Discussion or Pull Request. + + Pull Requests created programmatically will be in `"draft"` status. + + Creating a Pull Request with changes can also be done at once with [`HfApi.create_commit`]. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + title (`str`): + The title of the discussion. It can be up to 200 characters long, + and must be at least 3 characters long. Leading and trailing whitespaces + will be stripped. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + description (`str`, *optional*): + An optional description for the Pull Request. + Defaults to `"Discussion opened with the huggingface_hub Python library"` + pull_request (`bool`, *optional*): + Whether to create a Pull Request or discussion. If `True`, creates a Pull Request. + If `False`, creates a discussion. Defaults to `False`. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + + Returns: [`DiscussionWithDetails`] + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + + if description is not None: + description = description.strip() + description = ( + description + if description + else ( + f"{'Pull Request' if pull_request else 'Discussion'} opened with the" + " [huggingface_hub Python" + " library](https://huggingface.co/docs/huggingface_hub)" + ) + ) + + headers = self._build_hf_headers(token=token) + resp = get_session().post( + f"{self.endpoint}/api/{repo_type}s/{repo_id}/discussions", + json={ + "title": title.strip(), + "description": description, + "pullRequest": pull_request, + }, + headers=headers, + ) + hf_raise_for_status(resp) + num = resp.json()["num"] + return self.get_discussion_details( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=num, + token=token, + ) + + @validate_hf_hub_args + def create_pull_request( + self, + repo_id: str, + title: str, + *, + token: Union[bool, str, None] = None, + description: Optional[str] = None, + repo_type: Optional[str] = None, + ) -> DiscussionWithDetails: + """Creates a Pull Request . Pull Requests created programmatically will be in `"draft"` status. + + Creating a Pull Request with changes can also be done at once with [`HfApi.create_commit`]; + + This is a wrapper around [`HfApi.create_discussion`]. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + title (`str`): + The title of the discussion. It can be up to 200 characters long, + and must be at least 3 characters long. Leading and trailing whitespaces + will be stripped. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + description (`str`, *optional*): + An optional description for the Pull Request. + Defaults to `"Discussion opened with the huggingface_hub Python library"` + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + + Returns: [`DiscussionWithDetails`] + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + """ + return self.create_discussion( + repo_id=repo_id, + title=title, + token=token, + description=description, + repo_type=repo_type, + pull_request=True, + ) + + def _post_discussion_changes( + self, + *, + repo_id: str, + discussion_num: int, + resource: str, + body: Optional[dict] = None, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> requests.Response: + """Internal utility to POST changes to a Discussion or Pull Request""" + if not isinstance(discussion_num, int) or discussion_num <= 0: + raise ValueError("Invalid discussion_num, must be a positive integer") + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + repo_id = f"{repo_type}s/{repo_id}" + + path = f"{self.endpoint}/api/{repo_id}/discussions/{discussion_num}/{resource}" + + headers = self._build_hf_headers(token=token) + resp = requests.post(path, headers=headers, json=body) + hf_raise_for_status(resp) + return resp + + @validate_hf_hub_args + def comment_discussion( + self, + repo_id: str, + discussion_num: int, + comment: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> DiscussionComment: + """Creates a new comment on the given Discussion. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + comment (`str`): + The content of the comment to create. Comments support markdown formatting. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionComment`]: the newly created comment + + + Examples: + ```python + + >>> comment = \"\"\" + ... Hello @otheruser! + ... + ... # This is a title + ... + ... **This is bold**, *this is italic* and ~this is strikethrough~ + ... And [this](http://url) is a link + ... \"\"\" + + >>> HfApi().comment_discussion( + ... repo_id="username/repo_name", + ... discussion_num=34 + ... comment=comment + ... ) + # DiscussionComment(id='deadbeef0000000', type='comment', ...) + + ``` + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + resp = self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource="comment", + body={"comment": comment}, + ) + return deserialize_event(resp.json()["newMessage"]) # type: ignore + + @validate_hf_hub_args + def rename_discussion( + self, + repo_id: str, + discussion_num: int, + new_title: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> DiscussionTitleChange: + """Renames a Discussion. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + new_title (`str`): + The new title for the discussion + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionTitleChange`]: the title change event + + + Examples: + ```python + >>> new_title = "New title, fixing a typo" + >>> HfApi().rename_discussion( + ... repo_id="username/repo_name", + ... discussion_num=34 + ... new_title=new_title + ... ) + # DiscussionTitleChange(id='deadbeef0000000', type='title-change', ...) + + ``` + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + resp = self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource="title", + body={"title": new_title}, + ) + return deserialize_event(resp.json()["newTitle"]) # type: ignore + + @validate_hf_hub_args + def change_discussion_status( + self, + repo_id: str, + discussion_num: int, + new_status: Literal["open", "closed"], + *, + token: Union[bool, str, None] = None, + comment: Optional[str] = None, + repo_type: Optional[str] = None, + ) -> DiscussionStatusChange: + """Closes or re-opens a Discussion or Pull Request. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + new_status (`str`): + The new status for the discussion, either `"open"` or `"closed"`. + comment (`str`, *optional*): + An optional comment to post with the status change. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionStatusChange`]: the status change event + + + Examples: + ```python + >>> new_title = "New title, fixing a typo" + >>> HfApi().rename_discussion( + ... repo_id="username/repo_name", + ... discussion_num=34 + ... new_title=new_title + ... ) + # DiscussionStatusChange(id='deadbeef0000000', type='status-change', ...) + + ``` + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + if new_status not in ["open", "closed"]: + raise ValueError("Invalid status, valid statuses are: 'open' and 'closed'") + body: Dict[str, str] = {"status": new_status} + if comment and comment.strip(): + body["comment"] = comment.strip() + resp = self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource="status", + body=body, + ) + return deserialize_event(resp.json()["newStatus"]) # type: ignore + + @validate_hf_hub_args + def merge_pull_request( + self, + repo_id: str, + discussion_num: int, + *, + token: Union[bool, str, None] = None, + comment: Optional[str] = None, + repo_type: Optional[str] = None, + ): + """Merges a Pull Request. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + comment (`str`, *optional*): + An optional comment to post with the status change. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionStatusChange`]: the status change event + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource="merge", + body={"comment": comment.strip()} if comment and comment.strip() else None, + ) + + @validate_hf_hub_args + def edit_discussion_comment( + self, + repo_id: str, + discussion_num: int, + comment_id: str, + new_content: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> DiscussionComment: + """Edits a comment on a Discussion / Pull Request. + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + comment_id (`str`): + The ID of the comment to edit. + new_content (`str`): + The new content of the comment. Comments support markdown formatting. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionComment`]: the edited comment + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + resp = self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource=f"comment/{comment_id.lower()}/edit", + body={"content": new_content}, + ) + return deserialize_event(resp.json()["updatedComment"]) # type: ignore + + @validate_hf_hub_args + def hide_discussion_comment( + self, + repo_id: str, + discussion_num: int, + comment_id: str, + *, + token: Union[bool, str, None] = None, + repo_type: Optional[str] = None, + ) -> DiscussionComment: + """Hides a comment on a Discussion / Pull Request. + + + Hidden comments' content cannot be retrieved anymore. Hiding a comment is irreversible. + + + Args: + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + discussion_num (`int`): + The number of the Discussion or Pull Request . Must be a strictly positive integer. + comment_id (`str`): + The ID of the comment to edit. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if uploading to a dataset or + space, `None` or `"model"` if uploading to a model. Default is + `None`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`DiscussionComment`]: the hidden comment + + + + Raises the following errors: + + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the HuggingFace API returned an error + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if some parameter value is invalid + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + + + """ + warnings.warn( + "Hidden comments' content cannot be retrieved anymore. Hiding a comment is irreversible.", + UserWarning, + ) + resp = self._post_discussion_changes( + repo_id=repo_id, + repo_type=repo_type, + discussion_num=discussion_num, + token=token, + resource=f"comment/{comment_id.lower()}/hide", + ) + return deserialize_event(resp.json()["updatedComment"]) # type: ignore + + @validate_hf_hub_args + def add_space_secret( + self, + repo_id: str, + key: str, + value: str, + *, + description: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> None: + """Adds or updates a secret in a Space. + + Secrets allow to set secret keys or tokens to a Space without hardcoding them. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets. + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + key (`str`): + Secret key. Example: `"GITHUB_API_KEY"` + value (`str`): + Secret value. Example: `"your_github_api_key"`. + description (`str`, *optional*): + Secret description. Example: `"Github API key to access the Github API"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + payload = {"key": key, "value": value} + if description is not None: + payload["description"] = description + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/secrets", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(r) + + @validate_hf_hub_args + def delete_space_secret(self, repo_id: str, key: str, *, token: Union[bool, str, None] = None) -> None: + """Deletes a secret from a Space. + + Secrets allow to set secret keys or tokens to a Space without hardcoding them. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets. + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + key (`str`): + Secret key. Example: `"GITHUB_API_KEY"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + r = get_session().delete( + f"{self.endpoint}/api/spaces/{repo_id}/secrets", + headers=self._build_hf_headers(token=token), + json={"key": key}, + ) + hf_raise_for_status(r) + + @validate_hf_hub_args + def get_space_variables(self, repo_id: str, *, token: Union[bool, str, None] = None) -> Dict[str, SpaceVariable]: + """Gets all variables from a Space. + + Variables allow to set environment variables to a Space without hardcoding them. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets-and-environment-variables + + Args: + repo_id (`str`): + ID of the repo to query. Example: `"bigcode/in-the-stack"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + r = get_session().get( + f"{self.endpoint}/api/spaces/{repo_id}/variables", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(r) + return {k: SpaceVariable(k, v) for k, v in r.json().items()} + + @validate_hf_hub_args + def add_space_variable( + self, + repo_id: str, + key: str, + value: str, + *, + description: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Dict[str, SpaceVariable]: + """Adds or updates a variable in a Space. + + Variables allow to set environment variables to a Space without hardcoding them. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets-and-environment-variables + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + key (`str`): + Variable key. Example: `"MODEL_REPO_ID"` + value (`str`): + Variable value. Example: `"the_model_repo_id"`. + description (`str`): + Description of the variable. Example: `"Model Repo ID of the implemented model"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + payload = {"key": key, "value": value} + if description is not None: + payload["description"] = description + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/variables", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(r) + return {k: SpaceVariable(k, v) for k, v in r.json().items()} + + @validate_hf_hub_args + def delete_space_variable( + self, repo_id: str, key: str, *, token: Union[bool, str, None] = None + ) -> Dict[str, SpaceVariable]: + """Deletes a variable from a Space. + + Variables allow to set environment variables to a Space without hardcoding them. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets-and-environment-variables + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + key (`str`): + Variable key. Example: `"MODEL_REPO_ID"` + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + r = get_session().delete( + f"{self.endpoint}/api/spaces/{repo_id}/variables", + headers=self._build_hf_headers(token=token), + json={"key": key}, + ) + hf_raise_for_status(r) + return {k: SpaceVariable(k, v) for k, v in r.json().items()} + + @validate_hf_hub_args + def get_space_runtime(self, repo_id: str, *, token: Union[bool, str, None] = None) -> SpaceRuntime: + """Gets runtime information about a Space. + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + Returns: + [`SpaceRuntime`]: Runtime information about a Space including Space stage and hardware. + """ + r = get_session().get( + f"{self.endpoint}/api/spaces/{repo_id}/runtime", headers=self._build_hf_headers(token=token) + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + @validate_hf_hub_args + def request_space_hardware( + self, + repo_id: str, + hardware: SpaceHardware, + *, + token: Union[bool, str, None] = None, + sleep_time: Optional[int] = None, + ) -> SpaceRuntime: + """Request new hardware for a Space. + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + hardware (`str` or [`SpaceHardware`]): + Hardware on which to run the Space. Example: `"t4-medium"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + sleep_time (`int`, *optional*): + Number of seconds of inactivity to wait before a Space is put to sleep. Set to `-1` if you don't want + your Space to sleep (default behavior for upgraded hardware). For free hardware, you can't configure + the sleep time (value is fixed to 48 hours of inactivity). + See https://huggingface.co/docs/hub/spaces-gpus#sleep-time for more details. + Returns: + [`SpaceRuntime`]: Runtime information about a Space including Space stage and hardware. + + + + It is also possible to request hardware directly when creating the Space repo! See [`create_repo`] for details. + + + """ + if sleep_time is not None and hardware == SpaceHardware.CPU_BASIC: + warnings.warn( + "If your Space runs on the default 'cpu-basic' hardware, it will go to sleep if inactive for more" + " than 48 hours. This value is not configurable. If you don't want your Space to deactivate or if" + " you want to set a custom sleep time, you need to upgrade to a paid Hardware.", + UserWarning, + ) + payload: Dict[str, Any] = {"flavor": hardware} + if sleep_time is not None: + payload["sleepTimeSeconds"] = sleep_time + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/hardware", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + @validate_hf_hub_args + def set_space_sleep_time( + self, repo_id: str, sleep_time: int, *, token: Union[bool, str, None] = None + ) -> SpaceRuntime: + """Set a custom sleep time for a Space running on upgraded hardware.. + + Your Space will go to sleep after X seconds of inactivity. You are not billed when your Space is in "sleep" + mode. If a new visitor lands on your Space, it will "wake it up". Only upgraded hardware can have a + configurable sleep time. To know more about the sleep stage, please refer to + https://huggingface.co/docs/hub/spaces-gpus#sleep-time. + + Args: + repo_id (`str`): + ID of the repo to update. Example: `"bigcode/in-the-stack"`. + sleep_time (`int`, *optional*): + Number of seconds of inactivity to wait before a Space is put to sleep. Set to `-1` if you don't want + your Space to pause (default behavior for upgraded hardware). For free hardware, you can't configure + the sleep time (value is fixed to 48 hours of inactivity). + See https://huggingface.co/docs/hub/spaces-gpus#sleep-time for more details. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + Returns: + [`SpaceRuntime`]: Runtime information about a Space including Space stage and hardware. + + + + It is also possible to set a custom sleep time when requesting hardware with [`request_space_hardware`]. + + + """ + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/sleeptime", + headers=self._build_hf_headers(token=token), + json={"seconds": sleep_time}, + ) + hf_raise_for_status(r) + runtime = SpaceRuntime(r.json()) + + hardware = runtime.requested_hardware or runtime.hardware + if hardware == SpaceHardware.CPU_BASIC: + warnings.warn( + "If your Space runs on the default 'cpu-basic' hardware, it will go to sleep if inactive for more" + " than 48 hours. This value is not configurable. If you don't want your Space to deactivate or if" + " you want to set a custom sleep time, you need to upgrade to a paid Hardware.", + UserWarning, + ) + return runtime + + @validate_hf_hub_args + def pause_space(self, repo_id: str, *, token: Union[bool, str, None] = None) -> SpaceRuntime: + """Pause your Space. + + A paused Space stops executing until manually restarted by its owner. This is different from the sleeping + state in which free Spaces go after 48h of inactivity. Paused time is not billed to your account, no matter the + hardware you've selected. To restart your Space, use [`restart_space`] and go to your Space settings page. + + For more details, please visit [the docs](https://huggingface.co/docs/hub/spaces-gpus#pause). + + Args: + repo_id (`str`): + ID of the Space to pause. Example: `"Salesforce/BLIP2"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`SpaceRuntime`]: Runtime information about your Space including `stage=PAUSED` and requested hardware. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If your Space is not found (error 404). Most probably wrong repo_id or your space is private but you + are not authenticated. + [`~utils.HfHubHTTPError`]: + 403 Forbidden: only the owner of a Space can pause it. If you want to manage a Space that you don't + own, either ask the owner by opening a Discussion or duplicate the Space. + [`~utils.BadRequestError`]: + If your Space is a static Space. Static Spaces are always running and never billed. If you want to hide + a static Space, you can set it to private. + """ + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/pause", headers=self._build_hf_headers(token=token) + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + @validate_hf_hub_args + def restart_space( + self, repo_id: str, *, token: Union[bool, str, None] = None, factory_reboot: bool = False + ) -> SpaceRuntime: + """Restart your Space. + + This is the only way to programmatically restart a Space if you've put it on Pause (see [`pause_space`]). You + must be the owner of the Space to restart it. If you are using an upgraded hardware, your account will be + billed as soon as the Space is restarted. You can trigger a restart no matter the current state of a Space. + + For more details, please visit [the docs](https://huggingface.co/docs/hub/spaces-gpus#pause). + + Args: + repo_id (`str`): + ID of the Space to restart. Example: `"Salesforce/BLIP2"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + factory_reboot (`bool`, *optional*): + If `True`, the Space will be rebuilt from scratch without caching any requirements. + + Returns: + [`SpaceRuntime`]: Runtime information about your Space. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If your Space is not found (error 404). Most probably wrong repo_id or your space is private but you + are not authenticated. + [`~utils.HfHubHTTPError`]: + 403 Forbidden: only the owner of a Space can restart it. If you want to restart a Space that you don't + own, either ask the owner by opening a Discussion or duplicate the Space. + [`~utils.BadRequestError`]: + If your Space is a static Space. Static Spaces are always running and never billed. If you want to hide + a static Space, you can set it to private. + """ + params = {} + if factory_reboot: + params["factory"] = "true" + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/restart", headers=self._build_hf_headers(token=token), params=params + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + @validate_hf_hub_args + def duplicate_space( + self, + from_id: str, + to_id: Optional[str] = None, + *, + private: Optional[bool] = None, + token: Union[bool, str, None] = None, + exist_ok: bool = False, + hardware: Optional[SpaceHardware] = None, + storage: Optional[SpaceStorage] = None, + sleep_time: Optional[int] = None, + secrets: Optional[List[Dict[str, str]]] = None, + variables: Optional[List[Dict[str, str]]] = None, + ) -> RepoUrl: + """Duplicate a Space. + + Programmatically duplicate a Space. The new Space will be created in your account and will be in the same state + as the original Space (running or paused). You can duplicate a Space no matter the current state of a Space. + + Args: + from_id (`str`): + ID of the Space to duplicate. Example: `"pharma/CLIP-Interrogator"`. + to_id (`str`, *optional*): + ID of the new Space. Example: `"dog/CLIP-Interrogator"`. If not provided, the new Space will have the same + name as the original Space, but in your account. + private (`bool`, *optional*): + Whether the new Space should be private or not. Defaults to the same privacy as the original Space. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + exist_ok (`bool`, *optional*, defaults to `False`): + If `True`, do not raise an error if repo already exists. + hardware (`SpaceHardware` or `str`, *optional*): + Choice of Hardware. Example: `"t4-medium"`. See [`SpaceHardware`] for a complete list. + storage (`SpaceStorage` or `str`, *optional*): + Choice of persistent storage tier. Example: `"small"`. See [`SpaceStorage`] for a complete list. + sleep_time (`int`, *optional*): + Number of seconds of inactivity to wait before a Space is put to sleep. Set to `-1` if you don't want + your Space to sleep (default behavior for upgraded hardware). For free hardware, you can't configure + the sleep time (value is fixed to 48 hours of inactivity). + See https://huggingface.co/docs/hub/spaces-gpus#sleep-time for more details. + secrets (`List[Dict[str, str]]`, *optional*): + A list of secret keys to set in your Space. Each item is in the form `{"key": ..., "value": ..., "description": ...}` where description is optional. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets. + variables (`List[Dict[str, str]]`, *optional*): + A list of public environment variables to set in your Space. Each item is in the form `{"key": ..., "value": ..., "description": ...}` where description is optional. + For more details, see https://huggingface.co/docs/hub/spaces-overview#managing-secrets-and-environment-variables. + + Returns: + [`RepoUrl`]: URL to the newly created repo. Value is a subclass of `str` containing + attributes like `endpoint`, `repo_type` and `repo_id`. + + Raises: + [`~utils.RepositoryNotFoundError`]: + If one of `from_id` or `to_id` cannot be found. This may be because it doesn't exist, + or because it is set to `private` and you do not have access. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + If the HuggingFace API returned an error + + Example: + ```python + >>> from huggingface_hub import duplicate_space + + # Duplicate a Space to your account + >>> duplicate_space("multimodalart/dreambooth-training") + RepoUrl('https://huggingface.co/spaces/nateraw/dreambooth-training',...) + + # Can set custom destination id and visibility flag. + >>> duplicate_space("multimodalart/dreambooth-training", to_id="my-dreambooth", private=True) + RepoUrl('https://huggingface.co/spaces/nateraw/my-dreambooth',...) + ``` + """ + # Parse to_id if provided + parsed_to_id = RepoUrl(to_id) if to_id is not None else None + + # Infer target repo_id + to_namespace = ( # set namespace manually or default to username + parsed_to_id.namespace + if parsed_to_id is not None and parsed_to_id.namespace is not None + else self.whoami(token)["name"] + ) + to_repo_name = parsed_to_id.repo_name if to_id is not None else RepoUrl(from_id).repo_name # type: ignore + + # repository must be a valid repo_id (namespace/repo_name). + payload: Dict[str, Any] = {"repository": f"{to_namespace}/{to_repo_name}"} + + keys = ["private", "hardware", "storageTier", "sleepTimeSeconds", "secrets", "variables"] + values = [private, hardware, storage, sleep_time, secrets, variables] + payload.update({k: v for k, v in zip(keys, values) if v is not None}) + + if sleep_time is not None and hardware == SpaceHardware.CPU_BASIC: + warnings.warn( + "If your Space runs on the default 'cpu-basic' hardware, it will go to sleep if inactive for more" + " than 48 hours. This value is not configurable. If you don't want your Space to deactivate or if" + " you want to set a custom sleep time, you need to upgrade to a paid Hardware.", + UserWarning, + ) + + r = get_session().post( + f"{self.endpoint}/api/spaces/{from_id}/duplicate", + headers=self._build_hf_headers(token=token), + json=payload, + ) + + try: + hf_raise_for_status(r) + except HTTPError as err: + if exist_ok and err.response.status_code == 409: + # Repo already exists and `exist_ok=True` + pass + else: + raise + + return RepoUrl(r.json()["url"], endpoint=self.endpoint) + + @validate_hf_hub_args + def request_space_storage( + self, + repo_id: str, + storage: SpaceStorage, + *, + token: Union[bool, str, None] = None, + ) -> SpaceRuntime: + """Request persistent storage for a Space. + + Args: + repo_id (`str`): + ID of the Space to update. Example: `"open-llm-leaderboard/open_llm_leaderboard"`. + storage (`str` or [`SpaceStorage`]): + Storage tier. Either 'small', 'medium', or 'large'. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + Returns: + [`SpaceRuntime`]: Runtime information about a Space including Space stage and hardware. + + + + It is not possible to decrease persistent storage after its granted. To do so, you must delete it + via [`delete_space_storage`]. + + + """ + payload: Dict[str, SpaceStorage] = {"tier": storage} + r = get_session().post( + f"{self.endpoint}/api/spaces/{repo_id}/storage", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + @validate_hf_hub_args + def delete_space_storage( + self, + repo_id: str, + *, + token: Union[bool, str, None] = None, + ) -> SpaceRuntime: + """Delete persistent storage for a Space. + + Args: + repo_id (`str`): + ID of the Space to update. Example: `"open-llm-leaderboard/open_llm_leaderboard"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + Returns: + [`SpaceRuntime`]: Runtime information about a Space including Space stage and hardware. + Raises: + [`BadRequestError`] + If space has no persistent storage. + + """ + r = get_session().delete( + f"{self.endpoint}/api/spaces/{repo_id}/storage", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(r) + return SpaceRuntime(r.json()) + + ####################### + # Inference Endpoints # + ####################### + + def list_inference_endpoints( + self, namespace: Optional[str] = None, *, token: Union[bool, str, None] = None + ) -> List[InferenceEndpoint]: + """Lists all inference endpoints for the given namespace. + + Args: + namespace (`str`, *optional*): + The namespace to list endpoints for. Defaults to the current user. Set to `"*"` to list all endpoints + from all namespaces (i.e. personal namespace and all orgs the user belongs to). + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + List[`InferenceEndpoint`]: A list of all inference endpoints for the given namespace. + + Example: + ```python + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> api.list_inference_endpoints() + [InferenceEndpoint(name='my-endpoint', ...), ...] + ``` + """ + # Special case: list all endpoints for all namespaces the user has access to + if namespace == "*": + user = self.whoami(token=token) + + # List personal endpoints first + endpoints: List[InferenceEndpoint] = list_inference_endpoints(namespace=self._get_namespace(token=token)) + + # Then list endpoints for all orgs the user belongs to and ignore 401 errors (no billing or no access) + for org in user.get("orgs", []): + try: + endpoints += list_inference_endpoints(namespace=org["name"], token=token) + except HfHubHTTPError as error: + if error.response.status_code == 401: # Either no billing or user don't have access) + logger.debug("Cannot list Inference Endpoints for org '%s': %s", org["name"], error) + pass + + return endpoints + + # Normal case: list endpoints for a specific namespace + namespace = namespace or self._get_namespace(token=token) + + response = get_session().get( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + return [ + InferenceEndpoint.from_raw(endpoint, namespace=namespace, token=token) + for endpoint in response.json()["items"] + ] + + def create_inference_endpoint( + self, + name: str, + *, + repository: str, + framework: str, + accelerator: str, + instance_size: str, + instance_type: str, + region: str, + vendor: str, + account_id: Optional[str] = None, + min_replica: int = 0, + max_replica: int = 1, + scale_to_zero_timeout: int = 15, + revision: Optional[str] = None, + task: Optional[str] = None, + custom_image: Optional[Dict] = None, + env: Optional[Dict[str, str]] = None, + secrets: Optional[Dict[str, str]] = None, + type: InferenceEndpointType = InferenceEndpointType.PROTECTED, + domain: Optional[str] = None, + path: Optional[str] = None, + cache_http_responses: Optional[bool] = None, + tags: Optional[List[str]] = None, + namespace: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> InferenceEndpoint: + """Create a new Inference Endpoint. + + Args: + name (`str`): + The unique name for the new Inference Endpoint. + repository (`str`): + The name of the model repository associated with the Inference Endpoint (e.g. `"gpt2"`). + framework (`str`): + The machine learning framework used for the model (e.g. `"custom"`). + accelerator (`str`): + The hardware accelerator to be used for inference (e.g. `"cpu"`). + instance_size (`str`): + The size or type of the instance to be used for hosting the model (e.g. `"x4"`). + instance_type (`str`): + The cloud instance type where the Inference Endpoint will be deployed (e.g. `"intel-icl"`). + region (`str`): + The cloud region in which the Inference Endpoint will be created (e.g. `"us-east-1"`). + vendor (`str`): + The cloud provider or vendor where the Inference Endpoint will be hosted (e.g. `"aws"`). + account_id (`str`, *optional*): + The account ID used to link a VPC to a private Inference Endpoint (if applicable). + min_replica (`int`, *optional*): + The minimum number of replicas (instances) to keep running for the Inference Endpoint. Defaults to 0. + max_replica (`int`, *optional*): + The maximum number of replicas (instances) to scale to for the Inference Endpoint. Defaults to 1. + scale_to_zero_timeout (`int`, *optional*): + The duration in minutes before an inactive endpoint is scaled to zero. Defaults to 15. + revision (`str`, *optional*): + The specific model revision to deploy on the Inference Endpoint (e.g. `"6c0e6080953db56375760c0471a8c5f2929baf11"`). + task (`str`, *optional*): + The task on which to deploy the model (e.g. `"text-classification"`). + custom_image (`Dict`, *optional*): + A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an + Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). + env (`Dict[str, str]`, *optional*): + Non-secret environment variables to inject in the container environment. + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. + type ([`InferenceEndpointType]`, *optional*): + The type of the Inference Endpoint, which can be `"protected"` (default), `"public"` or `"private"`. + domain (`str`, *optional*): + The custom domain for the Inference Endpoint deployment, if setup the inference endpoint will be available at this domain (e.g. `"my-new-domain.cool-website.woof"`). + path (`str`, *optional*): + The custom path to the deployed model, should start with a `/` (e.g. `"/models/google-bert/bert-base-uncased"`). + cache_http_responses (`bool`, *optional*): + Whether to cache HTTP responses from the Inference Endpoint. Defaults to `False`. + tags (`List[str]`, *optional*): + A list of tags to associate with the Inference Endpoint. + namespace (`str`, *optional*): + The namespace where the Inference Endpoint will be created. Defaults to the current user's namespace. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the updated Inference Endpoint. + + Example: + ```python + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> endpoint = api.create_inference_endpoint( + ... "my-endpoint-name", + ... repository="gpt2", + ... framework="pytorch", + ... task="text-generation", + ... accelerator="cpu", + ... vendor="aws", + ... region="us-east-1", + ... type="protected", + ... instance_size="x2", + ... instance_type="intel-icl", + ... ) + >>> endpoint + InferenceEndpoint(name='my-endpoint-name', status="pending",...) + + # Run inference on the endpoint + >>> endpoint.client.text_generation(...) + "..." + ``` + + ```python + # Start an Inference Endpoint running Zephyr-7b-beta on TGI + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> endpoint = api.create_inference_endpoint( + ... "aws-zephyr-7b-beta-0486", + ... repository="HuggingFaceH4/zephyr-7b-beta", + ... framework="pytorch", + ... task="text-generation", + ... accelerator="gpu", + ... vendor="aws", + ... region="us-east-1", + ... type="protected", + ... instance_size="x1", + ... instance_type="nvidia-a10g", + ... env={ + ... "MAX_BATCH_PREFILL_TOKENS": "2048", + ... "MAX_INPUT_LENGTH": "1024", + ... "MAX_TOTAL_TOKENS": "1512", + ... "MODEL_ID": "/repository" + ... }, + ... custom_image={ + ... "health_route": "/health", + ... "url": "ghcr.io/huggingface/text-generation-inference:1.1.0", + ... }, + ... secrets={"MY_SECRET_KEY": "secret_value"}, + ... tags=["dev", "text-generation"], + ... ) + + ``` + """ + namespace = namespace or self._get_namespace(token=token) + + if custom_image is not None: + image = ( + custom_image + if next(iter(custom_image)) in constants.INFERENCE_ENDPOINT_IMAGE_KEYS + else {"custom": custom_image} + ) + else: + image = {"huggingface": {}} + + payload: Dict = { + "accountId": account_id, + "compute": { + "accelerator": accelerator, + "instanceSize": instance_size, + "instanceType": instance_type, + "scaling": { + "maxReplica": max_replica, + "minReplica": min_replica, + "scaleToZeroTimeout": scale_to_zero_timeout, + }, + }, + "model": { + "framework": framework, + "repository": repository, + "revision": revision, + "task": task, + "image": image, + }, + "name": name, + "provider": { + "region": region, + "vendor": vendor, + }, + "type": type, + } + if env: + payload["model"]["env"] = env + if secrets: + payload["model"]["secrets"] = secrets + if domain is not None or path is not None: + payload["route"] = {} + if domain is not None: + payload["route"]["domain"] = domain + if path is not None: + payload["route"]["path"] = path + if cache_http_responses is not None: + payload["cacheHttpResponses"] = cache_http_responses + if tags is not None: + payload["tags"] = tags + + response = get_session().post( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(response) + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + @experimental + @validate_hf_hub_args + def create_inference_endpoint_from_catalog( + self, + repo_id: str, + *, + name: Optional[str] = None, + token: Union[bool, str, None] = None, + namespace: Optional[str] = None, + ) -> InferenceEndpoint: + """Create a new Inference Endpoint from a model in the Hugging Face Inference Catalog. + + The goal of the Inference Catalog is to provide a curated list of models that are optimized for inference + and for which default configurations have been tested. See https://endpoints.huggingface.co/catalog for a list + of available models in the catalog. + + Args: + repo_id (`str`): + The ID of the model in the catalog to deploy as an Inference Endpoint. + name (`str`, *optional*): + The unique name for the new Inference Endpoint. If not provided, a random name will be generated. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + namespace (`str`, *optional*): + The namespace where the Inference Endpoint will be created. Defaults to the current user's namespace. + + Returns: + [`InferenceEndpoint`]: information about the new Inference Endpoint. + + + + `create_inference_endpoint_from_catalog` is experimental. Its API is subject to change in the future. Please provide feedback + if you have any suggestions or requests. + + + """ + token = token or self.token or get_token() + payload: Dict = { + "namespace": namespace or self._get_namespace(token=token), + "repoId": repo_id, + } + if name is not None: + payload["endpointName"] = name + + response = get_session().post( + f"{constants.INFERENCE_CATALOG_ENDPOINT}/deploy", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(response) + data = response.json()["endpoint"] + return InferenceEndpoint.from_raw(data, namespace=data["name"], token=token) + + @experimental + @validate_hf_hub_args + def list_inference_catalog(self, *, token: Union[bool, str, None] = None) -> List[str]: + """List models available in the Hugging Face Inference Catalog. + + The goal of the Inference Catalog is to provide a curated list of models that are optimized for inference + and for which default configurations have been tested. See https://endpoints.huggingface.co/catalog for a list + of available models in the catalog. + + Use [`create_inference_endpoint_from_catalog`] to deploy a model from the catalog. + + Args: + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + + Returns: + List[`str`]: A list of model IDs available in the catalog. + + + `list_inference_catalog` is experimental. Its API is subject to change in the future. Please provide feedback + if you have any suggestions or requests. + + + """ + response = get_session().get( + f"{constants.INFERENCE_CATALOG_ENDPOINT}/repo-list", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + return response.json()["models"] + + def get_inference_endpoint( + self, name: str, *, namespace: Optional[str] = None, token: Union[bool, str, None] = None + ) -> InferenceEndpoint: + """Get information about an Inference Endpoint. + + Args: + name (`str`): + The name of the Inference Endpoint to retrieve information about. + namespace (`str`, *optional*): + The namespace in which the Inference Endpoint is located. Defaults to the current user. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the requested Inference Endpoint. + + Example: + ```python + >>> from huggingface_hub import HfApi + >>> api = HfApi() + >>> endpoint = api.get_inference_endpoint("my-text-to-image") + >>> endpoint + InferenceEndpoint(name='my-text-to-image', ...) + + # Get status + >>> endpoint.status + 'running' + >>> endpoint.url + 'https://my-text-to-image.region.vendor.endpoints.huggingface.cloud' + + # Run inference + >>> endpoint.client.text_to_image(...) + ``` + """ + namespace = namespace or self._get_namespace(token=token) + + response = get_session().get( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + def update_inference_endpoint( + self, + name: str, + *, + # Compute update + accelerator: Optional[str] = None, + instance_size: Optional[str] = None, + instance_type: Optional[str] = None, + min_replica: Optional[int] = None, + max_replica: Optional[int] = None, + scale_to_zero_timeout: Optional[int] = None, + # Model update + repository: Optional[str] = None, + framework: Optional[str] = None, + revision: Optional[str] = None, + task: Optional[str] = None, + custom_image: Optional[Dict] = None, + env: Optional[Dict[str, str]] = None, + secrets: Optional[Dict[str, str]] = None, + # Route update + domain: Optional[str] = None, + path: Optional[str] = None, + # Other + cache_http_responses: Optional[bool] = None, + tags: Optional[List[str]] = None, + namespace: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> InferenceEndpoint: + """Update an Inference Endpoint. + + This method allows the update of either the compute configuration, the deployed model, the route, or any combination. + All arguments are optional but at least one must be provided. + + For convenience, you can also update an Inference Endpoint using [`InferenceEndpoint.update`]. + + Args: + name (`str`): + The name of the Inference Endpoint to update. + + accelerator (`str`, *optional*): + The hardware accelerator to be used for inference (e.g. `"cpu"`). + instance_size (`str`, *optional*): + The size or type of the instance to be used for hosting the model (e.g. `"x4"`). + instance_type (`str`, *optional*): + The cloud instance type where the Inference Endpoint will be deployed (e.g. `"intel-icl"`). + min_replica (`int`, *optional*): + The minimum number of replicas (instances) to keep running for the Inference Endpoint. + max_replica (`int`, *optional*): + The maximum number of replicas (instances) to scale to for the Inference Endpoint. + scale_to_zero_timeout (`int`, *optional*): + The duration in minutes before an inactive endpoint is scaled to zero. + + repository (`str`, *optional*): + The name of the model repository associated with the Inference Endpoint (e.g. `"gpt2"`). + framework (`str`, *optional*): + The machine learning framework used for the model (e.g. `"custom"`). + revision (`str`, *optional*): + The specific model revision to deploy on the Inference Endpoint (e.g. `"6c0e6080953db56375760c0471a8c5f2929baf11"`). + task (`str`, *optional*): + The task on which to deploy the model (e.g. `"text-classification"`). + custom_image (`Dict`, *optional*): + A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an + Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). + env (`Dict[str, str]`, *optional*): + Non-secret environment variables to inject in the container environment + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. + + domain (`str`, *optional*): + The custom domain for the Inference Endpoint deployment, if setup the inference endpoint will be available at this domain (e.g. `"my-new-domain.cool-website.woof"`). + path (`str`, *optional*): + The custom path to the deployed model, should start with a `/` (e.g. `"/models/google-bert/bert-base-uncased"`). + + cache_http_responses (`bool`, *optional*): + Whether to cache HTTP responses from the Inference Endpoint. + tags (`List[str]`, *optional*): + A list of tags to associate with the Inference Endpoint. + + namespace (`str`, *optional*): + The namespace where the Inference Endpoint will be updated. Defaults to the current user's namespace. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the updated Inference Endpoint. + """ + namespace = namespace or self._get_namespace(token=token) + + # Populate only the fields that are not None + payload: Dict = defaultdict(lambda: defaultdict(dict)) + if accelerator is not None: + payload["compute"]["accelerator"] = accelerator + if instance_size is not None: + payload["compute"]["instanceSize"] = instance_size + if instance_type is not None: + payload["compute"]["instanceType"] = instance_type + if max_replica is not None: + payload["compute"]["scaling"]["maxReplica"] = max_replica + if min_replica is not None: + payload["compute"]["scaling"]["minReplica"] = min_replica + if scale_to_zero_timeout is not None: + payload["compute"]["scaling"]["scaleToZeroTimeout"] = scale_to_zero_timeout + if repository is not None: + payload["model"]["repository"] = repository + if framework is not None: + payload["model"]["framework"] = framework + if revision is not None: + payload["model"]["revision"] = revision + if task is not None: + payload["model"]["task"] = task + if custom_image is not None: + payload["model"]["image"] = {"custom": custom_image} + if env is not None: + payload["model"]["env"] = env + if secrets is not None: + payload["model"]["secrets"] = secrets + if domain is not None: + payload["route"]["domain"] = domain + if path is not None: + payload["route"]["path"] = path + if cache_http_responses is not None: + payload["cacheHttpResponses"] = cache_http_responses + if tags is not None: + payload["tags"] = tags + + response = get_session().put( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(response) + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + def delete_inference_endpoint( + self, name: str, *, namespace: Optional[str] = None, token: Union[bool, str, None] = None + ) -> None: + """Delete an Inference Endpoint. + + This operation is not reversible. If you don't want to be charged for an Inference Endpoint, it is preferable + to pause it with [`pause_inference_endpoint`] or scale it to zero with [`scale_to_zero_inference_endpoint`]. + + For convenience, you can also delete an Inference Endpoint using [`InferenceEndpoint.delete`]. + + Args: + name (`str`): + The name of the Inference Endpoint to delete. + namespace (`str`, *optional*): + The namespace in which the Inference Endpoint is located. Defaults to the current user. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + """ + namespace = namespace or self._get_namespace(token=token) + response = get_session().delete( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + def pause_inference_endpoint( + self, name: str, *, namespace: Optional[str] = None, token: Union[bool, str, None] = None + ) -> InferenceEndpoint: + """Pause an Inference Endpoint. + + A paused Inference Endpoint will not be charged. It can be resumed at any time using [`resume_inference_endpoint`]. + This is different than scaling the Inference Endpoint to zero with [`scale_to_zero_inference_endpoint`], which + would be automatically restarted when a request is made to it. + + For convenience, you can also pause an Inference Endpoint using [`pause_inference_endpoint`]. + + Args: + name (`str`): + The name of the Inference Endpoint to pause. + namespace (`str`, *optional*): + The namespace in which the Inference Endpoint is located. Defaults to the current user. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the paused Inference Endpoint. + """ + namespace = namespace or self._get_namespace(token=token) + + response = get_session().post( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}/pause", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + def resume_inference_endpoint( + self, + name: str, + *, + namespace: Optional[str] = None, + running_ok: bool = True, + token: Union[bool, str, None] = None, + ) -> InferenceEndpoint: + """Resume an Inference Endpoint. + + For convenience, you can also resume an Inference Endpoint using [`InferenceEndpoint.resume`]. + + Args: + name (`str`): + The name of the Inference Endpoint to resume. + namespace (`str`, *optional*): + The namespace in which the Inference Endpoint is located. Defaults to the current user. + running_ok (`bool`, *optional*): + If `True`, the method will not raise an error if the Inference Endpoint is already running. Defaults to + `True`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the resumed Inference Endpoint. + """ + namespace = namespace or self._get_namespace(token=token) + + response = get_session().post( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}/resume", + headers=self._build_hf_headers(token=token), + ) + try: + hf_raise_for_status(response) + except HfHubHTTPError as error: + # If already running (and it's ok), then fetch current status and return + if running_ok and error.response.status_code == 400 and "already running" in error.response.text: + return self.get_inference_endpoint(name, namespace=namespace, token=token) + # Otherwise, raise the error + raise + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + def scale_to_zero_inference_endpoint( + self, name: str, *, namespace: Optional[str] = None, token: Union[bool, str, None] = None + ) -> InferenceEndpoint: + """Scale Inference Endpoint to zero. + + An Inference Endpoint scaled to zero will not be charged. It will be resume on the next request to it, with a + cold start delay. This is different than pausing the Inference Endpoint with [`pause_inference_endpoint`], which + would require a manual resume with [`resume_inference_endpoint`]. + + For convenience, you can also scale an Inference Endpoint to zero using [`InferenceEndpoint.scale_to_zero`]. + + Args: + name (`str`): + The name of the Inference Endpoint to scale to zero. + namespace (`str`, *optional*): + The namespace in which the Inference Endpoint is located. Defaults to the current user. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`InferenceEndpoint`]: information about the scaled-to-zero Inference Endpoint. + """ + namespace = namespace or self._get_namespace(token=token) + + response = get_session().post( + f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}/scale-to-zero", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + return InferenceEndpoint.from_raw(response.json(), namespace=namespace, token=token) + + def _get_namespace(self, token: Union[bool, str, None] = None) -> str: + """Get the default namespace for the current user.""" + me = self.whoami(token=token) + if me["type"] == "user": + return me["name"] + else: + raise ValueError( + "Cannot determine default namespace. You must provide a 'namespace' as input or be logged in as a" + " user." + ) + + ######################## + # Collection Endpoints # + ######################## + @validate_hf_hub_args + def list_collections( + self, + *, + owner: Union[List[str], str, None] = None, + item: Union[List[str], str, None] = None, + sort: Optional[Literal["lastModified", "trending", "upvotes"]] = None, + limit: Optional[int] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[Collection]: + """List collections on the Huggingface Hub, given some filters. + + + + When listing collections, the item list per collection is truncated to 4 items maximum. To retrieve all items + from a collection, you must use [`get_collection`]. + + + + Args: + owner (`List[str]` or `str`, *optional*): + Filter by owner's username. + item (`List[str]` or `str`, *optional*): + Filter collections containing a particular items. Example: `"models/teknium/OpenHermes-2.5-Mistral-7B"`, `"datasets/squad"` or `"papers/2311.12983"`. + sort (`Literal["lastModified", "trending", "upvotes"]`, *optional*): + Sort collections by last modified, trending or upvotes. + limit (`int`, *optional*): + Maximum number of collections to be returned. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[Collection]`: an iterable of [`Collection`] objects. + """ + # Construct the API endpoint + path = f"{self.endpoint}/api/collections" + headers = self._build_hf_headers(token=token) + params: Dict = {} + if owner is not None: + params.update({"owner": owner}) + if item is not None: + params.update({"item": item}) + if sort is not None: + params.update({"sort": sort}) + if limit is not None: + params.update({"limit": limit}) + + # Paginate over the results until limit is reached + items = paginate(path, headers=headers, params=params) + if limit is not None: + items = islice(items, limit) # Do not iterate over all pages + + # Parse as Collection and return + for position, collection_data in enumerate(items): + yield Collection(position=position, **collection_data) + + def get_collection(self, collection_slug: str, *, token: Union[bool, str, None] = None) -> Collection: + """Gets information about a Collection on the Hub. + + Args: + collection_slug (`str`): + Slug of the collection of the Hub. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: [`Collection`] + + Example: + + ```py + >>> from huggingface_hub import get_collection + >>> collection = get_collection("TheBloke/recent-models-64f9a55bb3115b4f513ec026") + >>> collection.title + 'Recent models' + >>> len(collection.items) + 37 + >>> collection.items[0] + CollectionItem( + item_object_id='651446103cd773a050bf64c2', + item_id='TheBloke/U-Amethyst-20B-AWQ', + item_type='model', + position=88, + note=None + ) + ``` + """ + r = get_session().get( + f"{self.endpoint}/api/collections/{collection_slug}", headers=self._build_hf_headers(token=token) + ) + hf_raise_for_status(r) + return Collection(**{**r.json(), "endpoint": self.endpoint}) + + def create_collection( + self, + title: str, + *, + namespace: Optional[str] = None, + description: Optional[str] = None, + private: bool = False, + exists_ok: bool = False, + token: Union[bool, str, None] = None, + ) -> Collection: + """Create a new Collection on the Hub. + + Args: + title (`str`): + Title of the collection to create. Example: `"Recent models"`. + namespace (`str`, *optional*): + Namespace of the collection to create (username or org). Will default to the owner name. + description (`str`, *optional*): + Description of the collection to create. + private (`bool`, *optional*): + Whether the collection should be private or not. Defaults to `False` (i.e. public collection). + exists_ok (`bool`, *optional*): + If `True`, do not raise an error if collection already exists. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: [`Collection`] + + Example: + + ```py + >>> from huggingface_hub import create_collection + >>> collection = create_collection( + ... title="ICCV 2023", + ... description="Portfolio of models, papers and demos I presented at ICCV 2023", + ... ) + >>> collection.slug + "username/iccv-2023-64f9a55bb3115b4f513ec026" + ``` + """ + if namespace is None: + namespace = self.whoami(token)["name"] + + payload = { + "title": title, + "namespace": namespace, + "private": private, + } + if description is not None: + payload["description"] = description + + r = get_session().post( + f"{self.endpoint}/api/collections", headers=self._build_hf_headers(token=token), json=payload + ) + try: + hf_raise_for_status(r) + except HTTPError as err: + if exists_ok and err.response.status_code == 409: + # Collection already exists and `exists_ok=True` + slug = r.json()["slug"] + return self.get_collection(slug, token=token) + else: + raise + return Collection(**{**r.json(), "endpoint": self.endpoint}) + + def update_collection_metadata( + self, + collection_slug: str, + *, + title: Optional[str] = None, + description: Optional[str] = None, + position: Optional[int] = None, + private: Optional[bool] = None, + theme: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Collection: + """Update metadata of a collection on the Hub. + + All arguments are optional. Only provided metadata will be updated. + + Args: + collection_slug (`str`): + Slug of the collection to update. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + title (`str`): + Title of the collection to update. + description (`str`, *optional*): + Description of the collection to update. + position (`int`, *optional*): + New position of the collection in the list of collections of the user. + private (`bool`, *optional*): + Whether the collection should be private or not. + theme (`str`, *optional*): + Theme of the collection on the Hub. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: [`Collection`] + + Example: + + ```py + >>> from huggingface_hub import update_collection_metadata + >>> collection = update_collection_metadata( + ... collection_slug="username/iccv-2023-64f9a55bb3115b4f513ec026", + ... title="ICCV Oct. 2023" + ... description="Portfolio of models, datasets, papers and demos I presented at ICCV Oct. 2023", + ... private=False, + ... theme="pink", + ... ) + >>> collection.slug + "username/iccv-oct-2023-64f9a55bb3115b4f513ec026" + # ^collection slug got updated but not the trailing ID + ``` + """ + payload = { + "position": position, + "private": private, + "theme": theme, + "title": title, + "description": description, + } + r = get_session().patch( + f"{self.endpoint}/api/collections/{collection_slug}", + headers=self._build_hf_headers(token=token), + # Only send not-none values to the API + json={key: value for key, value in payload.items() if value is not None}, + ) + hf_raise_for_status(r) + return Collection(**{**r.json()["data"], "endpoint": self.endpoint}) + + def delete_collection( + self, collection_slug: str, *, missing_ok: bool = False, token: Union[bool, str, None] = None + ) -> None: + """Delete a collection on the Hub. + + Args: + collection_slug (`str`): + Slug of the collection to delete. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + missing_ok (`bool`, *optional*): + If `True`, do not raise an error if collection doesn't exists. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Example: + + ```py + >>> from huggingface_hub import delete_collection + >>> collection = delete_collection("username/useless-collection-64f9a55bb3115b4f513ec026", missing_ok=True) + ``` + + + + This is a non-revertible action. A deleted collection cannot be restored. + + + """ + r = get_session().delete( + f"{self.endpoint}/api/collections/{collection_slug}", headers=self._build_hf_headers(token=token) + ) + try: + hf_raise_for_status(r) + except HTTPError as err: + if missing_ok and err.response.status_code == 404: + # Collection doesn't exists and `missing_ok=True` + return + else: + raise + + def add_collection_item( + self, + collection_slug: str, + item_id: str, + item_type: CollectionItemType_T, + *, + note: Optional[str] = None, + exists_ok: bool = False, + token: Union[bool, str, None] = None, + ) -> Collection: + """Add an item to a collection on the Hub. + + Args: + collection_slug (`str`): + Slug of the collection to update. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + item_id (`str`): + ID of the item to add to the collection. It can be the ID of a repo on the Hub (e.g. `"facebook/bart-large-mnli"`) + or a paper id (e.g. `"2307.09288"`). + item_type (`str`): + Type of the item to add. Can be one of `"model"`, `"dataset"`, `"space"` or `"paper"`. + note (`str`, *optional*): + A note to attach to the item in the collection. The maximum size for a note is 500 characters. + exists_ok (`bool`, *optional*): + If `True`, do not raise an error if item already exists. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: [`Collection`] + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the item you try to add to the collection does not exist on the Hub. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 409 if the item you try to add to the collection is already in the collection (and exists_ok=False) + + Example: + + ```py + >>> from huggingface_hub import add_collection_item + >>> collection = add_collection_item( + ... collection_slug="davanstrien/climate-64f99dc2a5067f6b65531bab", + ... item_id="pierre-loic/climate-news-articles", + ... item_type="dataset" + ... ) + >>> collection.items[-1].item_id + "pierre-loic/climate-news-articles" + # ^item got added to the collection on last position + + # Add item with a note + >>> add_collection_item( + ... collection_slug="davanstrien/climate-64f99dc2a5067f6b65531bab", + ... item_id="datasets/climate_fever", + ... item_type="dataset" + ... note="This dataset adopts the FEVER methodology that consists of 1,535 real-world claims regarding climate-change collected on the internet." + ... ) + (...) + ``` + """ + payload: Dict[str, Any] = {"item": {"id": item_id, "type": item_type}} + if note is not None: + payload["note"] = note + r = get_session().post( + f"{self.endpoint}/api/collections/{collection_slug}/items", + headers=self._build_hf_headers(token=token), + json=payload, + ) + try: + hf_raise_for_status(r) + except HTTPError as err: + if exists_ok and err.response.status_code == 409: + # Item already exists and `exists_ok=True` + return self.get_collection(collection_slug, token=token) + else: + raise + return Collection(**{**r.json(), "endpoint": self.endpoint}) + + def update_collection_item( + self, + collection_slug: str, + item_object_id: str, + *, + note: Optional[str] = None, + position: Optional[int] = None, + token: Union[bool, str, None] = None, + ) -> None: + """Update an item in a collection. + + Args: + collection_slug (`str`): + Slug of the collection to update. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + item_object_id (`str`): + ID of the item in the collection. This is not the id of the item on the Hub (repo_id or paper id). + It must be retrieved from a [`CollectionItem`] object. Example: `collection.items[0].item_object_id`. + note (`str`, *optional*): + A note to attach to the item in the collection. The maximum size for a note is 500 characters. + position (`int`, *optional*): + New position of the item in the collection. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Example: + + ```py + >>> from huggingface_hub import get_collection, update_collection_item + + # Get collection first + >>> collection = get_collection("TheBloke/recent-models-64f9a55bb3115b4f513ec026") + + # Update item based on its ID (add note + update position) + >>> update_collection_item( + ... collection_slug="TheBloke/recent-models-64f9a55bb3115b4f513ec026", + ... item_object_id=collection.items[-1].item_object_id, + ... note="Newly updated model!" + ... position=0, + ... ) + ``` + """ + payload = {"position": position, "note": note} + r = get_session().patch( + f"{self.endpoint}/api/collections/{collection_slug}/items/{item_object_id}", + headers=self._build_hf_headers(token=token), + # Only send not-none values to the API + json={key: value for key, value in payload.items() if value is not None}, + ) + hf_raise_for_status(r) + + def delete_collection_item( + self, + collection_slug: str, + item_object_id: str, + *, + missing_ok: bool = False, + token: Union[bool, str, None] = None, + ) -> None: + """Delete an item from a collection. + + Args: + collection_slug (`str`): + Slug of the collection to update. Example: `"TheBloke/recent-models-64f9a55bb3115b4f513ec026"`. + item_object_id (`str`): + ID of the item in the collection. This is not the id of the item on the Hub (repo_id or paper id). + It must be retrieved from a [`CollectionItem`] object. Example: `collection.items[0].item_object_id`. + missing_ok (`bool`, *optional*): + If `True`, do not raise an error if item doesn't exists. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Example: + + ```py + >>> from huggingface_hub import get_collection, delete_collection_item + + # Get collection first + >>> collection = get_collection("TheBloke/recent-models-64f9a55bb3115b4f513ec026") + + # Delete item based on its ID + >>> delete_collection_item( + ... collection_slug="TheBloke/recent-models-64f9a55bb3115b4f513ec026", + ... item_object_id=collection.items[-1].item_object_id, + ... ) + ``` + """ + r = get_session().delete( + f"{self.endpoint}/api/collections/{collection_slug}/items/{item_object_id}", + headers=self._build_hf_headers(token=token), + ) + try: + hf_raise_for_status(r) + except HTTPError as err: + if missing_ok and err.response.status_code == 404: + # Item already deleted and `missing_ok=True` + return + else: + raise + + ########################## + # Manage access requests # + ########################## + + @validate_hf_hub_args + def list_pending_access_requests( + self, repo_id: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> List[AccessRequest]: + """ + Get pending access requests for a given gated repo. + + A pending request means the user has requested access to the repo but the request has not been processed yet. + If the approval mode is automatic, this list should be empty. Pending requests can be accepted or rejected + using [`accept_access_request`] and [`reject_access_request`]. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to get access requests for. + repo_type (`str`, *optional*): + The type of the repo to get access requests for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[AccessRequest]`: A list of [`AccessRequest`] objects. Each time contains a `username`, `email`, + `status` and `timestamp` attribute. If the gated repo has a custom form, the `fields` attribute will + be populated with user's answers. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + + Example: + ```py + >>> from huggingface_hub import list_pending_access_requests, accept_access_request + + # List pending requests + >>> requests = list_pending_access_requests("meta-llama/Llama-2-7b") + >>> len(requests) + 411 + >>> requests[0] + [ + AccessRequest( + username='clem', + fullname='Clem 🤗', + email='***', + timestamp=datetime.datetime(2023, 11, 23, 18, 4, 53, 828000, tzinfo=datetime.timezone.utc), + status='pending', + fields=None, + ), + ... + ] + + # Accept Clem's request + >>> accept_access_request("meta-llama/Llama-2-7b", "clem") + ``` + """ + return self._list_access_requests(repo_id, "pending", repo_type=repo_type, token=token) + + @validate_hf_hub_args + def list_accepted_access_requests( + self, repo_id: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> List[AccessRequest]: + """ + Get accepted access requests for a given gated repo. + + An accepted request means the user has requested access to the repo and the request has been accepted. The user + can download any file of the repo. If the approval mode is automatic, this list should contains by default all + requests. Accepted requests can be cancelled or rejected at any time using [`cancel_access_request`] and + [`reject_access_request`]. A cancelled request will go back to the pending list while a rejected request will + go to the rejected list. In both cases, the user will lose access to the repo. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to get access requests for. + repo_type (`str`, *optional*): + The type of the repo to get access requests for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[AccessRequest]`: A list of [`AccessRequest`] objects. Each time contains a `username`, `email`, + `status` and `timestamp` attribute. If the gated repo has a custom form, the `fields` attribute will + be populated with user's answers. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + + Example: + ```py + >>> from huggingface_hub import list_accepted_access_requests + + >>> requests = list_accepted_access_requests("meta-llama/Llama-2-7b") + >>> len(requests) + 411 + >>> requests[0] + [ + AccessRequest( + username='clem', + fullname='Clem 🤗', + email='***', + timestamp=datetime.datetime(2023, 11, 23, 18, 4, 53, 828000, tzinfo=datetime.timezone.utc), + status='accepted', + fields=None, + ), + ... + ] + ``` + """ + return self._list_access_requests(repo_id, "accepted", repo_type=repo_type, token=token) + + @validate_hf_hub_args + def list_rejected_access_requests( + self, repo_id: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> List[AccessRequest]: + """ + Get rejected access requests for a given gated repo. + + A rejected request means the user has requested access to the repo and the request has been explicitly rejected + by a repo owner (either you or another user from your organization). The user cannot download any file of the + repo. Rejected requests can be accepted or cancelled at any time using [`accept_access_request`] and + [`cancel_access_request`]. A cancelled request will go back to the pending list while an accepted request will + go to the accepted list. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to get access requests for. + repo_type (`str`, *optional*): + The type of the repo to get access requests for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[AccessRequest]`: A list of [`AccessRequest`] objects. Each time contains a `username`, `email`, + `status` and `timestamp` attribute. If the gated repo has a custom form, the `fields` attribute will + be populated with user's answers. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + + Example: + ```py + >>> from huggingface_hub import list_rejected_access_requests + + >>> requests = list_rejected_access_requests("meta-llama/Llama-2-7b") + >>> len(requests) + 411 + >>> requests[0] + [ + AccessRequest( + username='clem', + fullname='Clem 🤗', + email='***', + timestamp=datetime.datetime(2023, 11, 23, 18, 4, 53, 828000, tzinfo=datetime.timezone.utc), + status='rejected', + fields=None, + ), + ... + ] + ``` + """ + return self._list_access_requests(repo_id, "rejected", repo_type=repo_type, token=token) + + def _list_access_requests( + self, + repo_id: str, + status: Literal["accepted", "rejected", "pending"], + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> List[AccessRequest]: + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + + response = get_session().get( + f"{constants.ENDPOINT}/api/{repo_type}s/{repo_id}/user-access-request/{status}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + return [ + AccessRequest( + username=request["user"]["user"], + fullname=request["user"]["fullname"], + email=request["user"].get("email"), + status=request["status"], + timestamp=parse_datetime(request["timestamp"]), + fields=request.get("fields"), # only if custom fields in form + ) + for request in response.json() + ] + + @validate_hf_hub_args + def cancel_access_request( + self, repo_id: str, user: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> None: + """ + Cancel an access request from a user for a given gated repo. + + A cancelled request will go back to the pending list and the user will lose access to the repo. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to cancel access request for. + user (`str`): + The username of the user which access request should be cancelled. + repo_type (`str`, *optional*): + The type of the repo to cancel access request for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user does not exist on the Hub. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request cannot be found. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request is already in the pending list. + """ + self._handle_access_request(repo_id, user, "pending", repo_type=repo_type, token=token) + + @validate_hf_hub_args + def accept_access_request( + self, repo_id: str, user: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> None: + """ + Accept an access request from a user for a given gated repo. + + Once the request is accepted, the user will be able to download any file of the repo and access the community + tab. If the approval mode is automatic, you don't have to accept requests manually. An accepted request can be + cancelled or rejected at any time using [`cancel_access_request`] and [`reject_access_request`]. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to accept access request for. + user (`str`): + The username of the user which access request should be accepted. + repo_type (`str`, *optional*): + The type of the repo to accept access request for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user does not exist on the Hub. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request cannot be found. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request is already in the accepted list. + """ + self._handle_access_request(repo_id, user, "accepted", repo_type=repo_type, token=token) + + @validate_hf_hub_args + def reject_access_request( + self, + repo_id: str, + user: str, + *, + repo_type: Optional[str] = None, + rejection_reason: Optional[str], + token: Union[bool, str, None] = None, + ) -> None: + """ + Reject an access request from a user for a given gated repo. + + A rejected request will go to the rejected list. The user cannot download any file of the repo. Rejected + requests can be accepted or cancelled at any time using [`accept_access_request`] and [`cancel_access_request`]. + A cancelled request will go back to the pending list while an accepted request will go to the accepted list. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to reject access request for. + user (`str`): + The username of the user which access request should be rejected. + repo_type (`str`, *optional*): + The type of the repo to reject access request for. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + rejection_reason (`str`, *optional*): + Optional rejection reason that will be visible to the user (max 200 characters). + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user does not exist on the Hub. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request cannot be found. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user access request is already in the rejected list. + """ + self._handle_access_request( + repo_id, user, "rejected", repo_type=repo_type, rejection_reason=rejection_reason, token=token + ) + + @validate_hf_hub_args + def _handle_access_request( + self, + repo_id: str, + user: str, + status: Literal["accepted", "rejected", "pending"], + repo_type: Optional[str] = None, + rejection_reason: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> None: + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + + payload = {"user": user, "status": status} + + if rejection_reason is not None: + if status != "rejected": + raise ValueError("`rejection_reason` can only be passed when rejecting an access request.") + payload["rejectionReason"] = rejection_reason + + response = get_session().post( + f"{constants.ENDPOINT}/api/{repo_type}s/{repo_id}/user-access-request/handle", + headers=self._build_hf_headers(token=token), + json=payload, + ) + hf_raise_for_status(response) + + @validate_hf_hub_args + def grant_access( + self, repo_id: str, user: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> None: + """ + Grant access to a user for a given gated repo. + + Granting access don't require for the user to send an access request by themselves. The user is automatically + added to the accepted list meaning they can download the files You can revoke the granted access at any time + using [`cancel_access_request`] or [`reject_access_request`]. + + For more info about gated repos, see https://huggingface.co/docs/hub/models-gated. + + Args: + repo_id (`str`): + The id of the repo to grant access to. + user (`str`): + The username of the user to grant access. + repo_type (`str`, *optional*): + The type of the repo to grant access to. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the repo is not gated. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 400 if the user already has access to the repo. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 403 if you only have read-only access to the repo. This can be the case if you don't have `write` + or `admin` role in the organization the repo belongs to or if you passed a `read` token. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 if the user does not exist on the Hub. + """ + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + + response = get_session().post( + f"{constants.ENDPOINT}/api/{repo_type}s/{repo_id}/user-access-request/grant", + headers=self._build_hf_headers(token=token), + json={"user": user}, + ) + hf_raise_for_status(response) + return response.json() + + ################### + # Manage webhooks # + ################### + + @validate_hf_hub_args + def get_webhook(self, webhook_id: str, *, token: Union[bool, str, None] = None) -> WebhookInfo: + """Get a webhook by its id. + + Args: + webhook_id (`str`): + The unique identifier of the webhook to get. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`WebhookInfo`]: + Info about the webhook. + + Example: + ```python + >>> from huggingface_hub import get_webhook + >>> webhook = get_webhook("654bbbc16f2ec14d77f109cc") + >>> print(webhook) + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + secret="my-secret", + domains=["repo", "discussion"], + disabled=False, + ) + ``` + """ + response = get_session().get( + f"{constants.ENDPOINT}/api/settings/webhooks/{webhook_id}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhook_data = response.json()["webhook"] + + watched_items = [WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook_data["watched"]] + + webhook = WebhookInfo( + id=webhook_data["id"], + url=webhook_data["url"], + watched=watched_items, + domains=webhook_data["domains"], + secret=webhook_data.get("secret"), + disabled=webhook_data["disabled"], + ) + + return webhook + + @validate_hf_hub_args + def list_webhooks(self, *, token: Union[bool, str, None] = None) -> List[WebhookInfo]: + """List all configured webhooks. + + Args: + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `List[WebhookInfo]`: + List of webhook info objects. + + Example: + ```python + >>> from huggingface_hub import list_webhooks + >>> webhooks = list_webhooks() + >>> len(webhooks) + 2 + >>> webhooks[0] + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + secret="my-secret", + domains=["repo", "discussion"], + disabled=False, + ) + ``` + """ + response = get_session().get( + f"{constants.ENDPOINT}/api/settings/webhooks", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhooks_data = response.json() + + return [ + WebhookInfo( + id=webhook["id"], + url=webhook["url"], + watched=[WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook["watched"]], + domains=webhook["domains"], + secret=webhook.get("secret"), + disabled=webhook["disabled"], + ) + for webhook in webhooks_data + ] + + @validate_hf_hub_args + def create_webhook( + self, + *, + url: str, + watched: List[Union[Dict, WebhookWatchedItem]], + domains: Optional[List[constants.WEBHOOK_DOMAIN_T]] = None, + secret: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> WebhookInfo: + """Create a new webhook. + + Args: + url (`str`): + URL to send the payload to. + watched (`List[WebhookWatchedItem]`): + List of [`WebhookWatchedItem`] to be watched by the webhook. It can be users, orgs, models, datasets or spaces. + Watched items can also be provided as plain dictionaries. + domains (`List[Literal["repo", "discussion"]]`, optional): + List of domains to watch. It can be "repo", "discussion" or both. + secret (`str`, optional): + A secret to sign the payload with. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`WebhookInfo`]: + Info about the newly created webhook. + + Example: + ```python + >>> from huggingface_hub import create_webhook + >>> payload = create_webhook( + ... watched=[{"type": "user", "name": "julien-c"}, {"type": "org", "name": "HuggingFaceH4"}], + ... url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + ... domains=["repo", "discussion"], + ... secret="my-secret", + ... ) + >>> print(payload) + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + domains=["repo", "discussion"], + secret="my-secret", + disabled=False, + ) + ``` + """ + watched_dicts = [asdict(item) if isinstance(item, WebhookWatchedItem) else item for item in watched] + + response = get_session().post( + f"{constants.ENDPOINT}/api/settings/webhooks", + json={"watched": watched_dicts, "url": url, "domains": domains, "secret": secret}, + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhook_data = response.json()["webhook"] + watched_items = [WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook_data["watched"]] + + webhook = WebhookInfo( + id=webhook_data["id"], + url=webhook_data["url"], + watched=watched_items, + domains=webhook_data["domains"], + secret=webhook_data.get("secret"), + disabled=webhook_data["disabled"], + ) + + return webhook + + @validate_hf_hub_args + def update_webhook( + self, + webhook_id: str, + *, + url: Optional[str] = None, + watched: Optional[List[Union[Dict, WebhookWatchedItem]]] = None, + domains: Optional[List[constants.WEBHOOK_DOMAIN_T]] = None, + secret: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> WebhookInfo: + """Update an existing webhook. + + Args: + webhook_id (`str`): + The unique identifier of the webhook to be updated. + url (`str`, optional): + The URL to which the payload will be sent. + watched (`List[WebhookWatchedItem]`, optional): + List of items to watch. It can be users, orgs, models, datasets, or spaces. + Refer to [`WebhookWatchedItem`] for more details. Watched items can also be provided as plain dictionaries. + domains (`List[Literal["repo", "discussion"]]`, optional): + The domains to watch. This can include "repo", "discussion", or both. + secret (`str`, optional): + A secret to sign the payload with, providing an additional layer of security. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`WebhookInfo`]: + Info about the updated webhook. + + Example: + ```python + >>> from huggingface_hub import update_webhook + >>> updated_payload = update_webhook( + ... webhook_id="654bbbc16f2ec14d77f109cc", + ... url="https://new.webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + ... watched=[{"type": "user", "name": "julien-c"}, {"type": "org", "name": "HuggingFaceH4"}], + ... domains=["repo"], + ... secret="my-secret", + ... ) + >>> print(updated_payload) + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + url="https://new.webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + domains=["repo"], + secret="my-secret", + disabled=False, + ``` + """ + if watched is None: + watched = [] + watched_dicts = [asdict(item) if isinstance(item, WebhookWatchedItem) else item for item in watched] + + response = get_session().post( + f"{constants.ENDPOINT}/api/settings/webhooks/{webhook_id}", + json={"watched": watched_dicts, "url": url, "domains": domains, "secret": secret}, + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhook_data = response.json()["webhook"] + + watched_items = [WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook_data["watched"]] + + webhook = WebhookInfo( + id=webhook_data["id"], + url=webhook_data["url"], + watched=watched_items, + domains=webhook_data["domains"], + secret=webhook_data.get("secret"), + disabled=webhook_data["disabled"], + ) + + return webhook + + @validate_hf_hub_args + def enable_webhook(self, webhook_id: str, *, token: Union[bool, str, None] = None) -> WebhookInfo: + """Enable a webhook (makes it "active"). + + Args: + webhook_id (`str`): + The unique identifier of the webhook to enable. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`WebhookInfo`]: + Info about the enabled webhook. + + Example: + ```python + >>> from huggingface_hub import enable_webhook + >>> enabled_webhook = enable_webhook("654bbbc16f2ec14d77f109cc") + >>> enabled_webhook + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + domains=["repo", "discussion"], + secret="my-secret", + disabled=False, + ) + ``` + """ + response = get_session().post( + f"{constants.ENDPOINT}/api/settings/webhooks/{webhook_id}/enable", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhook_data = response.json()["webhook"] + + watched_items = [WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook_data["watched"]] + + webhook = WebhookInfo( + id=webhook_data["id"], + url=webhook_data["url"], + watched=watched_items, + domains=webhook_data["domains"], + secret=webhook_data.get("secret"), + disabled=webhook_data["disabled"], + ) + + return webhook + + @validate_hf_hub_args + def disable_webhook(self, webhook_id: str, *, token: Union[bool, str, None] = None) -> WebhookInfo: + """Disable a webhook (makes it "disabled"). + + Args: + webhook_id (`str`): + The unique identifier of the webhook to disable. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + [`WebhookInfo`]: + Info about the disabled webhook. + + Example: + ```python + >>> from huggingface_hub import disable_webhook + >>> disabled_webhook = disable_webhook("654bbbc16f2ec14d77f109cc") + >>> disabled_webhook + WebhookInfo( + id="654bbbc16f2ec14d77f109cc", + url="https://webhook.site/a2176e82-5720-43ee-9e06-f91cb4c91548", + watched=[WebhookWatchedItem(type="user", name="julien-c"), WebhookWatchedItem(type="org", name="HuggingFaceH4")], + domains=["repo", "discussion"], + secret="my-secret", + disabled=True, + ) + ``` + """ + response = get_session().post( + f"{constants.ENDPOINT}/api/settings/webhooks/{webhook_id}/disable", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + webhook_data = response.json()["webhook"] + + watched_items = [WebhookWatchedItem(type=item["type"], name=item["name"]) for item in webhook_data["watched"]] + + webhook = WebhookInfo( + id=webhook_data["id"], + url=webhook_data["url"], + watched=watched_items, + domains=webhook_data["domains"], + secret=webhook_data.get("secret"), + disabled=webhook_data["disabled"], + ) + + return webhook + + @validate_hf_hub_args + def delete_webhook(self, webhook_id: str, *, token: Union[bool, str, None] = None) -> None: + """Delete a webhook. + + Args: + webhook_id (`str`): + The unique identifier of the webhook to delete. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved token, which is the recommended + method for authentication (see https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `None` + + Example: + ```python + >>> from huggingface_hub import delete_webhook + >>> delete_webhook("654bbbc16f2ec14d77f109cc") + ``` + """ + response = get_session().delete( + f"{constants.ENDPOINT}/api/settings/webhooks/{webhook_id}", + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(response) + + ############# + # Internals # + ############# + + def _build_hf_headers( + self, + token: Union[bool, str, None] = None, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, + ) -> Dict[str, str]: + """ + Alias for [`build_hf_headers`] that uses the token from [`HfApi`] client + when `token` is not provided. + """ + if token is None: + # Cannot do `token = token or self.token` as token can be `False`. + token = self.token + return build_hf_headers( + token=token, + library_name=library_name or self.library_name, + library_version=library_version or self.library_version, + user_agent=user_agent or self.user_agent, + headers=self.headers, + ) + + def _prepare_folder_deletions( + self, + repo_id: str, + repo_type: Optional[str], + revision: Optional[str], + path_in_repo: str, + delete_patterns: Optional[Union[List[str], str]], + token: Union[bool, str, None] = None, + ) -> List[CommitOperationDelete]: + """Generate the list of Delete operations for a commit to delete files from a repo. + + List remote files and match them against the `delete_patterns` constraints. Returns a list of [`CommitOperationDelete`] + with the matching items. + + Note: `.gitattributes` file is essential to make a repo work properly on the Hub. This file will always be + kept even if it matches the `delete_patterns` constraints. + """ + if delete_patterns is None: + # If no delete patterns, no need to list and filter remote files + return [] + + # List remote files + filenames = self.list_repo_files(repo_id=repo_id, revision=revision, repo_type=repo_type, token=token) + + # Compute relative path in repo + if path_in_repo and path_in_repo not in (".", "./"): + path_in_repo = path_in_repo.strip("/") + "/" # harmonize + relpath_to_abspath = { + file[len(path_in_repo) :]: file for file in filenames if file.startswith(path_in_repo) + } + else: + relpath_to_abspath = {file: file for file in filenames} + + # Apply filter on relative paths and return + return [ + CommitOperationDelete(path_in_repo=relpath_to_abspath[relpath], is_folder=False) + for relpath in filter_repo_objects(relpath_to_abspath.keys(), allow_patterns=delete_patterns) + if relpath_to_abspath[relpath] != ".gitattributes" + ] + + def _prepare_upload_folder_additions( + self, + folder_path: Union[str, Path], + path_in_repo: str, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + repo_type: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> List[CommitOperationAdd]: + """Generate the list of Add operations for a commit to upload a folder. + + Files not matching the `allow_patterns` (allowlist) and `ignore_patterns` (denylist) + constraints are discarded. + """ + + folder_path = Path(folder_path).expanduser().resolve() + if not folder_path.is_dir(): + raise ValueError(f"Provided path: '{folder_path}' is not a directory") + + # List files from folder + relpath_to_abspath = { + path.relative_to(folder_path).as_posix(): path + for path in sorted(folder_path.glob("**/*")) # sorted to be deterministic + if path.is_file() + } + + # Filter files + # Patterns are applied on the path relative to `folder_path`. `path_in_repo` is prefixed after the filtering. + filtered_repo_objects = list( + filter_repo_objects( + relpath_to_abspath.keys(), allow_patterns=allow_patterns, ignore_patterns=ignore_patterns + ) + ) + + prefix = f"{path_in_repo.strip('/')}/" if path_in_repo else "" + + # If updating a README.md file, make sure the metadata format is valid + # It's better to fail early than to fail after all the files have been hashed. + if "README.md" in filtered_repo_objects: + self._validate_yaml( + content=relpath_to_abspath["README.md"].read_text(encoding="utf8"), + repo_type=repo_type, + token=token, + ) + if len(filtered_repo_objects) > 30: + log = logger.warning if len(filtered_repo_objects) > 200 else logger.info + log( + "It seems you are trying to upload a large folder at once. This might take some time and then fail if " + "the folder is too large. For such cases, it is recommended to upload in smaller batches or to use " + "`HfApi().upload_large_folder(...)`/`huggingface-cli upload-large-folder` instead. For more details, " + "check out https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#upload-a-large-folder." + ) + + logger.info(f"Start hashing {len(filtered_repo_objects)} files.") + operations = [ + CommitOperationAdd( + path_or_fileobj=relpath_to_abspath[relpath], # absolute path on disk + path_in_repo=prefix + relpath, # "absolute" path in repo + ) + for relpath in filtered_repo_objects + ] + logger.info(f"Finished hashing {len(filtered_repo_objects)} files.") + return operations + + def _validate_yaml(self, content: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None): + """ + Validate YAML from `README.md`, used before file hashing and upload. + + Args: + content (`str`): + Content of `README.md` to validate. + repo_type (`str`, *optional*): + The type of the repo to grant access to. Must be one of `model`, `dataset` or `space`. + Defaults to `model`. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Raises: + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if YAML is invalid + """ + repo_type = repo_type if repo_type is not None else constants.REPO_TYPE_MODEL + headers = self._build_hf_headers(token=token) + + response = get_session().post( + f"{self.endpoint}/api/validate-yaml", + json={"content": content, "repoType": repo_type}, + headers=headers, + ) + # Handle warnings (example: empty metadata) + response_content = response.json() + message = "\n".join([f"- {warning.get('message')}" for warning in response_content.get("warnings", [])]) + if message: + warnings.warn(f"Warnings while validating metadata in README.md:\n{message}") + + # Raise on errors + try: + hf_raise_for_status(response) + except BadRequestError as e: + errors = response_content.get("errors", []) + message = "\n".join([f"- {error.get('message')}" for error in errors]) + raise ValueError(f"Invalid metadata in README.md.\n{message}") from e + + def get_user_overview(self, username: str, token: Union[bool, str, None] = None) -> User: + """ + Get an overview of a user on the Hub. + + Args: + username (`str`): + Username of the user to get an overview of. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `User`: A [`User`] object with the user's overview. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 If the user does not exist on the Hub. + """ + r = get_session().get( + f"{constants.ENDPOINT}/api/users/{username}/overview", headers=self._build_hf_headers(token=token) + ) + hf_raise_for_status(r) + return User(**r.json()) + + def list_organization_members(self, organization: str, token: Union[bool, str, None] = None) -> Iterable[User]: + """ + List of members of an organization on the Hub. + + Args: + organization (`str`): + Name of the organization to get the members of. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[User]`: A list of [`User`] objects with the members of the organization. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 If the organization does not exist on the Hub. + + """ + for member in paginate( + path=f"{constants.ENDPOINT}/api/organizations/{organization}/members", + params={}, + headers=self._build_hf_headers(token=token), + ): + yield User(**member) + + def list_user_followers(self, username: str, token: Union[bool, str, None] = None) -> Iterable[User]: + """ + Get the list of followers of a user on the Hub. + + Args: + username (`str`): + Username of the user to get the followers of. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[User]`: A list of [`User`] objects with the followers of the user. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 If the user does not exist on the Hub. + + """ + for follower in paginate( + path=f"{constants.ENDPOINT}/api/users/{username}/followers", + params={}, + headers=self._build_hf_headers(token=token), + ): + yield User(**follower) + + def list_user_following(self, username: str, token: Union[bool, str, None] = None) -> Iterable[User]: + """ + Get the list of users followed by a user on the Hub. + + Args: + username (`str`): + Username of the user to get the users followed by. + token (Union[bool, str, None], optional): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[User]`: A list of [`User`] objects with the users followed by the user. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 If the user does not exist on the Hub. + + """ + for followed_user in paginate( + path=f"{constants.ENDPOINT}/api/users/{username}/following", + params={}, + headers=self._build_hf_headers(token=token), + ): + yield User(**followed_user) + + def list_papers( + self, + *, + query: Optional[str] = None, + token: Union[bool, str, None] = None, + ) -> Iterable[PaperInfo]: + """ + List daily papers on the Hugging Face Hub given a search query. + + Args: + query (`str`, *optional*): + A search query string to find papers. + If provided, returns papers that match the query. + token (Union[bool, str, None], *optional*): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + + Returns: + `Iterable[PaperInfo]`: an iterable of [`huggingface_hub.hf_api.PaperInfo`] objects. + + Example: + + ```python + >>> from huggingface_hub import HfApi + + >>> api = HfApi() + + # List all papers with "attention" in their title + >>> api.list_papers(query="attention") + ``` + """ + path = f"{self.endpoint}/api/papers/search" + params = {} + if query: + params["q"] = query + r = get_session().get( + path, + params=params, + headers=self._build_hf_headers(token=token), + ) + hf_raise_for_status(r) + for paper in r.json(): + yield PaperInfo(**paper) + + def paper_info(self, id: str) -> PaperInfo: + """ + Get information for a paper on the Hub. + + Args: + id (`str`, **optional**): + ArXiv id of the paper. + + Returns: + `PaperInfo`: A `PaperInfo` object. + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError): + HTTP 404 If the paper does not exist on the Hub. + """ + path = f"{self.endpoint}/api/papers/{id}" + r = get_session().get(path) + hf_raise_for_status(r) + return PaperInfo(**r.json()) + + def auth_check( + self, repo_id: str, *, repo_type: Optional[str] = None, token: Union[bool, str, None] = None + ) -> None: + """ + Check if the provided user token has access to a specific repository on the Hugging Face Hub. + + This method verifies whether the user, authenticated via the provided token, has access to the specified + repository. If the repository is not found or if the user lacks the required permissions to access it, + the method raises an appropriate exception. + + Args: + repo_id (`str`): + The repository to check for access. Format should be `"user/repo_name"`. + Example: `"user/my-cool-model"`. + + repo_type (`str`, *optional*): + The type of the repository. Should be one of `"model"`, `"dataset"`, or `"space"`. + If not specified, the default is `"model"`. + + token `(Union[bool, str, None]`, *optional*): + A valid user access token. If not provided, the locally saved token will be used, which is the + recommended authentication method. Set to `False` to disable authentication. + Refer to: https://huggingface.co/docs/huggingface_hub/quick-start#authentication. + + Raises: + [`~utils.RepositoryNotFoundError`]: + Raised if the repository does not exist, is private, or the user does not have access. This can + occur if the `repo_id` or `repo_type` is incorrect or if the repository is private but the user + is not authenticated. + + [`~utils.GatedRepoError`]: + Raised if the repository exists but is gated and the user is not authorized to access it. + + Example: + Check if the user has access to a repository: + + ```python + >>> from huggingface_hub import auth_check + >>> from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError + + try: + auth_check("user/my-cool-model") + except GatedRepoError: + # Handle gated repository error + print("You do not have permission to access this gated repository.") + except RepositoryNotFoundError: + # Handle repository not found error + print("The repository was not found or you do not have access.") + ``` + + In this example: + - If the user has access, the method completes successfully. + - If the repository is gated or does not exist, appropriate exceptions are raised, allowing the user + to handle them accordingly. + """ + headers = self._build_hf_headers(token=token) + if repo_type is None: + repo_type = constants.REPO_TYPE_MODEL + if repo_type not in constants.REPO_TYPES: + raise ValueError(f"Invalid repo type, must be one of {constants.REPO_TYPES}") + path = f"{self.endpoint}/api/{repo_type}s/{repo_id}/auth-check" + r = get_session().get(path, headers=headers) + hf_raise_for_status(r) + + +def _parse_revision_from_pr_url(pr_url: str) -> str: + """Safely parse revision number from a PR url. + + Example: + ```py + >>> _parse_revision_from_pr_url("https://huggingface.co/bigscience/bloom/discussions/2") + "refs/pr/2" + ``` + """ + re_match = re.match(_REGEX_DISCUSSION_URL, pr_url) + if re_match is None: + raise RuntimeError(f"Unexpected response from the hub, expected a Pull Request URL but got: '{pr_url}'") + return f"refs/pr/{re_match[1]}" + + +api = HfApi() + +whoami = api.whoami +auth_check = api.auth_check +get_token_permission = api.get_token_permission + +list_models = api.list_models +model_info = api.model_info + +list_datasets = api.list_datasets +dataset_info = api.dataset_info + +list_spaces = api.list_spaces +space_info = api.space_info + +list_papers = api.list_papers +paper_info = api.paper_info + +repo_exists = api.repo_exists +revision_exists = api.revision_exists +file_exists = api.file_exists +repo_info = api.repo_info +list_repo_files = api.list_repo_files +list_repo_refs = api.list_repo_refs +list_repo_commits = api.list_repo_commits +list_repo_tree = api.list_repo_tree +get_paths_info = api.get_paths_info + +get_model_tags = api.get_model_tags +get_dataset_tags = api.get_dataset_tags + +create_commit = api.create_commit +create_repo = api.create_repo +delete_repo = api.delete_repo +update_repo_visibility = api.update_repo_visibility +update_repo_settings = api.update_repo_settings +move_repo = api.move_repo +upload_file = api.upload_file +upload_folder = api.upload_folder +delete_file = api.delete_file +delete_folder = api.delete_folder +delete_files = api.delete_files +upload_large_folder = api.upload_large_folder +preupload_lfs_files = api.preupload_lfs_files +create_branch = api.create_branch +delete_branch = api.delete_branch +create_tag = api.create_tag +delete_tag = api.delete_tag +get_full_repo_name = api.get_full_repo_name + +# Danger-zone API +super_squash_history = api.super_squash_history +list_lfs_files = api.list_lfs_files +permanently_delete_lfs_files = api.permanently_delete_lfs_files + +# Safetensors helpers +get_safetensors_metadata = api.get_safetensors_metadata +parse_safetensors_file_metadata = api.parse_safetensors_file_metadata + +# Background jobs +run_as_future = api.run_as_future + +# Activity API +list_liked_repos = api.list_liked_repos +list_repo_likers = api.list_repo_likers +unlike = api.unlike + +# Community API +get_discussion_details = api.get_discussion_details +get_repo_discussions = api.get_repo_discussions +create_discussion = api.create_discussion +create_pull_request = api.create_pull_request +change_discussion_status = api.change_discussion_status +comment_discussion = api.comment_discussion +edit_discussion_comment = api.edit_discussion_comment +rename_discussion = api.rename_discussion +merge_pull_request = api.merge_pull_request + +# Space API +add_space_secret = api.add_space_secret +delete_space_secret = api.delete_space_secret +get_space_variables = api.get_space_variables +add_space_variable = api.add_space_variable +delete_space_variable = api.delete_space_variable +get_space_runtime = api.get_space_runtime +request_space_hardware = api.request_space_hardware +set_space_sleep_time = api.set_space_sleep_time +pause_space = api.pause_space +restart_space = api.restart_space +duplicate_space = api.duplicate_space +request_space_storage = api.request_space_storage +delete_space_storage = api.delete_space_storage + +# Inference Endpoint API +list_inference_endpoints = api.list_inference_endpoints +create_inference_endpoint = api.create_inference_endpoint +get_inference_endpoint = api.get_inference_endpoint +update_inference_endpoint = api.update_inference_endpoint +delete_inference_endpoint = api.delete_inference_endpoint +pause_inference_endpoint = api.pause_inference_endpoint +resume_inference_endpoint = api.resume_inference_endpoint +scale_to_zero_inference_endpoint = api.scale_to_zero_inference_endpoint +create_inference_endpoint_from_catalog = api.create_inference_endpoint_from_catalog +list_inference_catalog = api.list_inference_catalog + +# Collections API +get_collection = api.get_collection +list_collections = api.list_collections +create_collection = api.create_collection +update_collection_metadata = api.update_collection_metadata +delete_collection = api.delete_collection +add_collection_item = api.add_collection_item +update_collection_item = api.update_collection_item +delete_collection_item = api.delete_collection_item +delete_collection_item = api.delete_collection_item + +# Access requests API +list_pending_access_requests = api.list_pending_access_requests +list_accepted_access_requests = api.list_accepted_access_requests +list_rejected_access_requests = api.list_rejected_access_requests +cancel_access_request = api.cancel_access_request +accept_access_request = api.accept_access_request +reject_access_request = api.reject_access_request +grant_access = api.grant_access + +# Webhooks API +create_webhook = api.create_webhook +disable_webhook = api.disable_webhook +delete_webhook = api.delete_webhook +enable_webhook = api.enable_webhook +get_webhook = api.get_webhook +list_webhooks = api.list_webhooks +update_webhook = api.update_webhook + + +# User API +get_user_overview = api.get_user_overview +list_organization_members = api.list_organization_members +list_user_followers = api.list_user_followers +list_user_following = api.list_user_following diff --git a/lib/python3.12/site-packages/huggingface_hub/hf_file_system.py b/lib/python3.12/site-packages/huggingface_hub/hf_file_system.py new file mode 100644 index 0000000000000000000000000000000000000000..4f438edc05041fb86af5cecd7081a8c1329ce1b2 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/hf_file_system.py @@ -0,0 +1,1142 @@ +import os +import re +import tempfile +from collections import deque +from dataclasses import dataclass, field +from datetime import datetime +from itertools import chain +from pathlib import Path +from typing import Any, Dict, Iterator, List, NoReturn, Optional, Tuple, Union +from urllib.parse import quote, unquote + +import fsspec +from fsspec.callbacks import _DEFAULT_CALLBACK, NoOpCallback, TqdmCallback +from fsspec.utils import isfilelike +from requests import Response + +from . import constants +from ._commit_api import CommitOperationCopy, CommitOperationDelete +from .errors import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError +from .file_download import hf_hub_url, http_get +from .hf_api import HfApi, LastCommitInfo, RepoFile +from .utils import HFValidationError, hf_raise_for_status, http_backoff + + +# Regex used to match special revisions with "/" in them (see #1710) +SPECIAL_REFS_REVISION_REGEX = re.compile( + r""" + (^refs\/convert\/\w+) # `refs/convert/parquet` revisions + | + (^refs\/pr\/\d+) # PR revisions + """, + re.VERBOSE, +) + + +@dataclass +class HfFileSystemResolvedPath: + """Data structure containing information about a resolved Hugging Face file system path.""" + + repo_type: str + repo_id: str + revision: str + path_in_repo: str + # The part placed after '@' in the initial path. It can be a quoted or unquoted refs revision. + # Used to reconstruct the unresolved path to return to the user. + _raw_revision: Optional[str] = field(default=None, repr=False) + + def unresolve(self) -> str: + repo_path = constants.REPO_TYPES_URL_PREFIXES.get(self.repo_type, "") + self.repo_id + if self._raw_revision: + return f"{repo_path}@{self._raw_revision}/{self.path_in_repo}".rstrip("/") + elif self.revision != constants.DEFAULT_REVISION: + return f"{repo_path}@{safe_revision(self.revision)}/{self.path_in_repo}".rstrip("/") + else: + return f"{repo_path}/{self.path_in_repo}".rstrip("/") + + +class HfFileSystem(fsspec.AbstractFileSystem): + """ + Access a remote Hugging Face Hub repository as if were a local file system. + + + + [`HfFileSystem`] provides fsspec compatibility, which is useful for libraries that require it (e.g., reading + Hugging Face datasets directly with `pandas`). However, it introduces additional overhead due to this compatibility + layer. For better performance and reliability, it's recommended to use `HfApi` methods when possible. + + + + Args: + token (`str` or `bool`, *optional*): + A valid user access token (string). Defaults to the locally saved + token, which is the recommended method for authentication (see + https://huggingface.co/docs/huggingface_hub/quick-start#authentication). + To disable authentication, pass `False`. + endpoint (`str`, *optional*): + Endpoint of the Hub. Defaults to . + Usage: + + ```python + >>> from huggingface_hub import HfFileSystem + + >>> fs = HfFileSystem() + + >>> # List files + >>> fs.glob("my-username/my-model/*.bin") + ['my-username/my-model/pytorch_model.bin'] + >>> fs.ls("datasets/my-username/my-dataset", detail=False) + ['datasets/my-username/my-dataset/.gitattributes', 'datasets/my-username/my-dataset/README.md', 'datasets/my-username/my-dataset/data.json'] + + >>> # Read/write files + >>> with fs.open("my-username/my-model/pytorch_model.bin") as f: + ... data = f.read() + >>> with fs.open("my-username/my-model/pytorch_model.bin", "wb") as f: + ... f.write(data) + ``` + """ + + root_marker = "" + protocol = "hf" + + def __init__( + self, + *args, + endpoint: Optional[str] = None, + token: Union[bool, str, None] = None, + **storage_options, + ): + super().__init__(*args, **storage_options) + self.endpoint = endpoint or constants.ENDPOINT + self.token = token + self._api = HfApi(endpoint=endpoint, token=token) + # Maps (repo_type, repo_id, revision) to a 2-tuple with: + # * the 1st element indicating whether the repositoy and the revision exist + # * the 2nd element being the exception raised if the repository or revision doesn't exist + self._repo_and_revision_exists_cache: Dict[ + Tuple[str, str, Optional[str]], Tuple[bool, Optional[Exception]] + ] = {} + + def _repo_and_revision_exist( + self, repo_type: str, repo_id: str, revision: Optional[str] + ) -> Tuple[bool, Optional[Exception]]: + if (repo_type, repo_id, revision) not in self._repo_and_revision_exists_cache: + try: + self._api.repo_info( + repo_id, revision=revision, repo_type=repo_type, timeout=constants.HF_HUB_ETAG_TIMEOUT + ) + except (RepositoryNotFoundError, HFValidationError) as e: + self._repo_and_revision_exists_cache[(repo_type, repo_id, revision)] = False, e + self._repo_and_revision_exists_cache[(repo_type, repo_id, None)] = False, e + except RevisionNotFoundError as e: + self._repo_and_revision_exists_cache[(repo_type, repo_id, revision)] = False, e + self._repo_and_revision_exists_cache[(repo_type, repo_id, None)] = True, None + else: + self._repo_and_revision_exists_cache[(repo_type, repo_id, revision)] = True, None + self._repo_and_revision_exists_cache[(repo_type, repo_id, None)] = True, None + return self._repo_and_revision_exists_cache[(repo_type, repo_id, revision)] + + def resolve_path(self, path: str, revision: Optional[str] = None) -> HfFileSystemResolvedPath: + """ + Resolve a Hugging Face file system path into its components. + + Args: + path (`str`): + Path to resolve. + revision (`str`, *optional*): + The revision of the repo to resolve. Defaults to the revision specified in the path. + + Returns: + [`HfFileSystemResolvedPath`]: Resolved path information containing `repo_type`, `repo_id`, `revision` and `path_in_repo`. + + Raises: + `ValueError`: + If path contains conflicting revision information. + `NotImplementedError`: + If trying to list repositories. + """ + + def _align_revision_in_path_with_revision( + revision_in_path: Optional[str], revision: Optional[str] + ) -> Optional[str]: + if revision is not None: + if revision_in_path is not None and revision_in_path != revision: + raise ValueError( + f'Revision specified in path ("{revision_in_path}") and in `revision` argument ("{revision}")' + " are not the same." + ) + else: + revision = revision_in_path + return revision + + path = self._strip_protocol(path) + if not path: + # can't list repositories at root + raise NotImplementedError("Access to repositories lists is not implemented.") + elif path.split("/")[0] + "/" in constants.REPO_TYPES_URL_PREFIXES.values(): + if "/" not in path: + # can't list repositories at the repository type level + raise NotImplementedError("Access to repositories lists is not implemented.") + repo_type, path = path.split("/", 1) + repo_type = constants.REPO_TYPES_MAPPING[repo_type] + else: + repo_type = constants.REPO_TYPE_MODEL + if path.count("/") > 0: + if "@" in path: + repo_id, revision_in_path = path.split("@", 1) + if "/" in revision_in_path: + match = SPECIAL_REFS_REVISION_REGEX.search(revision_in_path) + if match is not None and revision in (None, match.group()): + # Handle `refs/convert/parquet` and PR revisions separately + path_in_repo = SPECIAL_REFS_REVISION_REGEX.sub("", revision_in_path).lstrip("/") + revision_in_path = match.group() + else: + revision_in_path, path_in_repo = revision_in_path.split("/", 1) + else: + path_in_repo = "" + revision = _align_revision_in_path_with_revision(unquote(revision_in_path), revision) + repo_and_revision_exist, err = self._repo_and_revision_exist(repo_type, repo_id, revision) + if not repo_and_revision_exist: + _raise_file_not_found(path, err) + else: + revision_in_path = None + repo_id_with_namespace = "/".join(path.split("/")[:2]) + path_in_repo_with_namespace = "/".join(path.split("/")[2:]) + repo_id_without_namespace = path.split("/")[0] + path_in_repo_without_namespace = "/".join(path.split("/")[1:]) + repo_id = repo_id_with_namespace + path_in_repo = path_in_repo_with_namespace + repo_and_revision_exist, err = self._repo_and_revision_exist(repo_type, repo_id, revision) + if not repo_and_revision_exist: + if isinstance(err, (RepositoryNotFoundError, HFValidationError)): + repo_id = repo_id_without_namespace + path_in_repo = path_in_repo_without_namespace + repo_and_revision_exist, _ = self._repo_and_revision_exist(repo_type, repo_id, revision) + if not repo_and_revision_exist: + _raise_file_not_found(path, err) + else: + _raise_file_not_found(path, err) + else: + repo_id = path + path_in_repo = "" + if "@" in path: + repo_id, revision_in_path = path.split("@", 1) + revision = _align_revision_in_path_with_revision(unquote(revision_in_path), revision) + else: + revision_in_path = None + repo_and_revision_exist, _ = self._repo_and_revision_exist(repo_type, repo_id, revision) + if not repo_and_revision_exist: + raise NotImplementedError("Access to repositories lists is not implemented.") + + revision = revision if revision is not None else constants.DEFAULT_REVISION + return HfFileSystemResolvedPath(repo_type, repo_id, revision, path_in_repo, _raw_revision=revision_in_path) + + def invalidate_cache(self, path: Optional[str] = None) -> None: + """ + Clear the cache for a given path. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.invalidate_cache). + + Args: + path (`str`, *optional*): + Path to clear from cache. If not provided, clear the entire cache. + + """ + if not path: + self.dircache.clear() + self._repo_and_revision_exists_cache.clear() + else: + resolved_path = self.resolve_path(path) + path = resolved_path.unresolve() + while path: + self.dircache.pop(path, None) + path = self._parent(path) + + # Only clear repo cache if path is to repo root + if not resolved_path.path_in_repo: + self._repo_and_revision_exists_cache.pop((resolved_path.repo_type, resolved_path.repo_id, None), None) + self._repo_and_revision_exists_cache.pop( + (resolved_path.repo_type, resolved_path.repo_id, resolved_path.revision), None + ) + + def _open( + self, + path: str, + mode: str = "rb", + revision: Optional[str] = None, + block_size: Optional[int] = None, + **kwargs, + ) -> "HfFileSystemFile": + if "a" in mode: + raise NotImplementedError("Appending to remote files is not yet supported.") + if block_size == 0: + return HfFileSystemStreamFile(self, path, mode=mode, revision=revision, block_size=block_size, **kwargs) + else: + return HfFileSystemFile(self, path, mode=mode, revision=revision, block_size=block_size, **kwargs) + + def _rm(self, path: str, revision: Optional[str] = None, **kwargs) -> None: + resolved_path = self.resolve_path(path, revision=revision) + self._api.delete_file( + path_in_repo=resolved_path.path_in_repo, + repo_id=resolved_path.repo_id, + token=self.token, + repo_type=resolved_path.repo_type, + revision=resolved_path.revision, + commit_message=kwargs.get("commit_message"), + commit_description=kwargs.get("commit_description"), + ) + self.invalidate_cache(path=resolved_path.unresolve()) + + def rm( + self, + path: str, + recursive: bool = False, + maxdepth: Optional[int] = None, + revision: Optional[str] = None, + **kwargs, + ) -> None: + """ + Delete files from a repository. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.rm). + + + + Note: When possible, use `HfApi.delete_file()` for better performance. + + + + Args: + path (`str`): + Path to delete. + recursive (`bool`, *optional*): + If True, delete directory and all its contents. Defaults to False. + maxdepth (`int`, *optional*): + Maximum number of subdirectories to visit when deleting recursively. + revision (`str`, *optional*): + The git revision to delete from. + + """ + resolved_path = self.resolve_path(path, revision=revision) + paths = self.expand_path(path, recursive=recursive, maxdepth=maxdepth, revision=revision) + paths_in_repo = [self.resolve_path(path).path_in_repo for path in paths if not self.isdir(path)] + operations = [CommitOperationDelete(path_in_repo=path_in_repo) for path_in_repo in paths_in_repo] + commit_message = f"Delete {path} " + commit_message += "recursively " if recursive else "" + commit_message += f"up to depth {maxdepth} " if maxdepth is not None else "" + # TODO: use `commit_description` to list all the deleted paths? + self._api.create_commit( + repo_id=resolved_path.repo_id, + repo_type=resolved_path.repo_type, + token=self.token, + operations=operations, + revision=resolved_path.revision, + commit_message=kwargs.get("commit_message", commit_message), + commit_description=kwargs.get("commit_description"), + ) + self.invalidate_cache(path=resolved_path.unresolve()) + + def ls( + self, path: str, detail: bool = True, refresh: bool = False, revision: Optional[str] = None, **kwargs + ) -> List[Union[str, Dict[str, Any]]]: + """ + List the contents of a directory. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.ls). + + + + Note: When possible, use `HfApi.list_repo_tree()` for better performance. + + + + Args: + path (`str`): + Path to the directory. + detail (`bool`, *optional*): + If True, returns a list of dictionaries containing file information. If False, + returns a list of file paths. Defaults to True. + refresh (`bool`, *optional*): + If True, bypass the cache and fetch the latest data. Defaults to False. + revision (`str`, *optional*): + The git revision to list from. + + Returns: + `List[Union[str, Dict[str, Any]]]`: List of file paths (if detail=False) or list of file information + dictionaries (if detail=True). + """ + resolved_path = self.resolve_path(path, revision=revision) + path = resolved_path.unresolve() + kwargs = {"expand_info": detail, **kwargs} + try: + out = self._ls_tree(path, refresh=refresh, revision=revision, **kwargs) + except EntryNotFoundError: + # Path could be a file + if not resolved_path.path_in_repo: + _raise_file_not_found(path, None) + out = self._ls_tree(self._parent(path), refresh=refresh, revision=revision, **kwargs) + out = [o for o in out if o["name"] == path] + if len(out) == 0: + _raise_file_not_found(path, None) + return out if detail else [o["name"] for o in out] + + def _ls_tree( + self, + path: str, + recursive: bool = False, + refresh: bool = False, + revision: Optional[str] = None, + expand_info: bool = True, + ): + resolved_path = self.resolve_path(path, revision=revision) + path = resolved_path.unresolve() + root_path = HfFileSystemResolvedPath( + resolved_path.repo_type, + resolved_path.repo_id, + resolved_path.revision, + path_in_repo="", + _raw_revision=resolved_path._raw_revision, + ).unresolve() + + out = [] + if path in self.dircache and not refresh: + cached_path_infos = self.dircache[path] + out.extend(cached_path_infos) + dirs_not_in_dircache = [] + if recursive: + # Use BFS to traverse the cache and build the "recursive "output + # (The Hub uses a so-called "tree first" strategy for the tree endpoint but we sort the output to follow the spec so the result is (eventually) the same) + dirs_to_visit = deque( + [path_info for path_info in cached_path_infos if path_info["type"] == "directory"] + ) + while dirs_to_visit: + dir_info = dirs_to_visit.popleft() + if dir_info["name"] not in self.dircache: + dirs_not_in_dircache.append(dir_info["name"]) + else: + cached_path_infos = self.dircache[dir_info["name"]] + out.extend(cached_path_infos) + dirs_to_visit.extend( + [path_info for path_info in cached_path_infos if path_info["type"] == "directory"] + ) + + dirs_not_expanded = [] + if expand_info: + # Check if there are directories with non-expanded entries + dirs_not_expanded = [self._parent(o["name"]) for o in out if o["last_commit"] is None] + + if (recursive and dirs_not_in_dircache) or (expand_info and dirs_not_expanded): + # If the dircache is incomplete, find the common path of the missing and non-expanded entries + # and extend the output with the result of `_ls_tree(common_path, recursive=True)` + common_prefix = os.path.commonprefix(dirs_not_in_dircache + dirs_not_expanded) + # Get the parent directory if the common prefix itself is not a directory + common_path = ( + common_prefix.rstrip("/") + if common_prefix.endswith("/") + or common_prefix == root_path + or common_prefix in chain(dirs_not_in_dircache, dirs_not_expanded) + else self._parent(common_prefix) + ) + out = [o for o in out if not o["name"].startswith(common_path + "/")] + for cached_path in self.dircache: + if cached_path.startswith(common_path + "/"): + self.dircache.pop(cached_path, None) + self.dircache.pop(common_path, None) + out.extend( + self._ls_tree( + common_path, + recursive=recursive, + refresh=True, + revision=revision, + expand_info=expand_info, + ) + ) + else: + tree = self._api.list_repo_tree( + resolved_path.repo_id, + resolved_path.path_in_repo, + recursive=recursive, + expand=expand_info, + revision=resolved_path.revision, + repo_type=resolved_path.repo_type, + ) + for path_info in tree: + if isinstance(path_info, RepoFile): + cache_path_info = { + "name": root_path + "/" + path_info.path, + "size": path_info.size, + "type": "file", + "blob_id": path_info.blob_id, + "lfs": path_info.lfs, + "last_commit": path_info.last_commit, + "security": path_info.security, + } + else: + cache_path_info = { + "name": root_path + "/" + path_info.path, + "size": 0, + "type": "directory", + "tree_id": path_info.tree_id, + "last_commit": path_info.last_commit, + } + parent_path = self._parent(cache_path_info["name"]) + self.dircache.setdefault(parent_path, []).append(cache_path_info) + out.append(cache_path_info) + return out + + def walk(self, path: str, *args, **kwargs) -> Iterator[Tuple[str, List[str], List[str]]]: + """ + Return all files below the given path. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.walk). + + Args: + path (`str`): + Root path to list files from. + + Returns: + `Iterator[Tuple[str, List[str], List[str]]]`: An iterator of (path, list of directory names, list of file names) tuples. + """ + # Set expand_info=False by default to get a x10 speed boost + kwargs = {"expand_info": kwargs.get("detail", False), **kwargs} + path = self.resolve_path(path, revision=kwargs.get("revision")).unresolve() + yield from super().walk(path, *args, **kwargs) + + def glob(self, path: str, **kwargs) -> List[str]: + """ + Find files by glob-matching. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.glob). + + Args: + path (`str`): + Path pattern to match. + + Returns: + `List[str]`: List of paths matching the pattern. + """ + # Set expand_info=False by default to get a x10 speed boost + kwargs = {"expand_info": kwargs.get("detail", False), **kwargs} + path = self.resolve_path(path, revision=kwargs.get("revision")).unresolve() + return super().glob(path, **kwargs) + + def find( + self, + path: str, + maxdepth: Optional[int] = None, + withdirs: bool = False, + detail: bool = False, + refresh: bool = False, + revision: Optional[str] = None, + **kwargs, + ) -> Union[List[str], Dict[str, Dict[str, Any]]]: + """ + List all files below path. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.find). + + Args: + path (`str`): + Root path to list files from. + maxdepth (`int`, *optional*): + Maximum depth to descend into subdirectories. + withdirs (`bool`, *optional*): + Include directory paths in the output. Defaults to False. + detail (`bool`, *optional*): + If True, returns a dict mapping paths to file information. Defaults to False. + refresh (`bool`, *optional*): + If True, bypass the cache and fetch the latest data. Defaults to False. + revision (`str`, *optional*): + The git revision to list from. + + Returns: + `Union[List[str], Dict[str, Dict[str, Any]]]`: List of paths or dict of file information. + """ + if maxdepth: + return super().find( + path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, refresh=refresh, revision=revision, **kwargs + ) + resolved_path = self.resolve_path(path, revision=revision) + path = resolved_path.unresolve() + kwargs = {"expand_info": detail, **kwargs} + try: + out = self._ls_tree(path, recursive=True, refresh=refresh, revision=resolved_path.revision, **kwargs) + except EntryNotFoundError: + # Path could be a file + if self.info(path, revision=revision, **kwargs)["type"] == "file": + out = {path: {}} + else: + out = {} + else: + if not withdirs: + out = [o for o in out if o["type"] != "directory"] + else: + # If `withdirs=True`, include the directory itself to be consistent with the spec + path_info = self.info(path, revision=resolved_path.revision, **kwargs) + out = [path_info] + out if path_info["type"] == "directory" else out + out = {o["name"]: o for o in out} + names = sorted(out) + if not detail: + return names + else: + return {name: out[name] for name in names} + + def cp_file(self, path1: str, path2: str, revision: Optional[str] = None, **kwargs) -> None: + """ + Copy a file within or between repositories. + + + + Note: When possible, use `HfApi.upload_file()` for better performance. + + + + Args: + path1 (`str`): + Source path to copy from. + path2 (`str`): + Destination path to copy to. + revision (`str`, *optional*): + The git revision to copy from. + + """ + resolved_path1 = self.resolve_path(path1, revision=revision) + resolved_path2 = self.resolve_path(path2, revision=revision) + + same_repo = ( + resolved_path1.repo_type == resolved_path2.repo_type and resolved_path1.repo_id == resolved_path2.repo_id + ) + + if same_repo: + commit_message = f"Copy {path1} to {path2}" + self._api.create_commit( + repo_id=resolved_path1.repo_id, + repo_type=resolved_path1.repo_type, + revision=resolved_path2.revision, + commit_message=kwargs.get("commit_message", commit_message), + commit_description=kwargs.get("commit_description", ""), + operations=[ + CommitOperationCopy( + src_path_in_repo=resolved_path1.path_in_repo, + path_in_repo=resolved_path2.path_in_repo, + src_revision=resolved_path1.revision, + ) + ], + ) + else: + with self.open(path1, "rb", revision=resolved_path1.revision) as f: + content = f.read() + commit_message = f"Copy {path1} to {path2}" + self._api.upload_file( + path_or_fileobj=content, + path_in_repo=resolved_path2.path_in_repo, + repo_id=resolved_path2.repo_id, + token=self.token, + repo_type=resolved_path2.repo_type, + revision=resolved_path2.revision, + commit_message=kwargs.get("commit_message", commit_message), + commit_description=kwargs.get("commit_description"), + ) + self.invalidate_cache(path=resolved_path1.unresolve()) + self.invalidate_cache(path=resolved_path2.unresolve()) + + def modified(self, path: str, **kwargs) -> datetime: + """ + Get the last modified time of a file. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.modified). + + Args: + path (`str`): + Path to the file. + + Returns: + `datetime`: Last commit date of the file. + """ + info = self.info(path, **kwargs) + return info["last_commit"]["date"] + + def info(self, path: str, refresh: bool = False, revision: Optional[str] = None, **kwargs) -> Dict[str, Any]: + """ + Get information about a file or directory. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.info). + + + + Note: When possible, use `HfApi.get_paths_info()` or `HfApi.repo_info()` for better performance. + + + + Args: + path (`str`): + Path to get info for. + refresh (`bool`, *optional*): + If True, bypass the cache and fetch the latest data. Defaults to False. + revision (`str`, *optional*): + The git revision to get info from. + + Returns: + `Dict[str, Any]`: Dictionary containing file information (type, size, commit info, etc.). + + """ + resolved_path = self.resolve_path(path, revision=revision) + path = resolved_path.unresolve() + expand_info = kwargs.get( + "expand_info", True + ) # don't expose it as a parameter in the public API to follow the spec + if not resolved_path.path_in_repo: + # Path is the root directory + out = { + "name": path, + "size": 0, + "type": "directory", + } + if expand_info: + last_commit = self._api.list_repo_commits( + resolved_path.repo_id, repo_type=resolved_path.repo_type, revision=resolved_path.revision + )[-1] + out = { + **out, + "tree_id": None, # TODO: tree_id of the root directory? + "last_commit": LastCommitInfo( + oid=last_commit.commit_id, title=last_commit.title, date=last_commit.created_at + ), + } + else: + out = None + parent_path = self._parent(path) + if not expand_info and parent_path not in self.dircache: + # Fill the cache with cheap call + self.ls(parent_path, expand_info=False) + if parent_path in self.dircache: + # Check if the path is in the cache + out1 = [o for o in self.dircache[parent_path] if o["name"] == path] + if not out1: + _raise_file_not_found(path, None) + out = out1[0] + if refresh or out is None or (expand_info and out and out["last_commit"] is None): + paths_info = self._api.get_paths_info( + resolved_path.repo_id, + resolved_path.path_in_repo, + expand=expand_info, + revision=resolved_path.revision, + repo_type=resolved_path.repo_type, + ) + if not paths_info: + _raise_file_not_found(path, None) + path_info = paths_info[0] + root_path = HfFileSystemResolvedPath( + resolved_path.repo_type, + resolved_path.repo_id, + resolved_path.revision, + path_in_repo="", + _raw_revision=resolved_path._raw_revision, + ).unresolve() + if isinstance(path_info, RepoFile): + out = { + "name": root_path + "/" + path_info.path, + "size": path_info.size, + "type": "file", + "blob_id": path_info.blob_id, + "lfs": path_info.lfs, + "last_commit": path_info.last_commit, + "security": path_info.security, + } + else: + out = { + "name": root_path + "/" + path_info.path, + "size": 0, + "type": "directory", + "tree_id": path_info.tree_id, + "last_commit": path_info.last_commit, + } + if not expand_info: + out = {k: out[k] for k in ["name", "size", "type"]} + assert out is not None + return out + + def exists(self, path, **kwargs): + """ + Check if a file exists. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.exists). + + + + Note: When possible, use `HfApi.file_exists()` for better performance. + + + + Args: + path (`str`): + Path to check. + + Returns: + `bool`: True if file exists, False otherwise. + """ + try: + if kwargs.get("refresh", False): + self.invalidate_cache(path) + + self.info(path, **{**kwargs, "expand_info": False}) + return True + except: # noqa: E722 + return False + + def isdir(self, path): + """ + Check if a path is a directory. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.isdir). + + Args: + path (`str`): + Path to check. + + Returns: + `bool`: True if path is a directory, False otherwise. + """ + try: + return self.info(path, expand_info=False)["type"] == "directory" + except OSError: + return False + + def isfile(self, path): + """ + Check if a path is a file. + + For more details, refer to [fsspec documentation](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.isfile). + + Args: + path (`str`): + Path to check. + + Returns: + `bool`: True if path is a file, False otherwise. + """ + try: + return self.info(path, expand_info=False)["type"] == "file" + except: # noqa: E722 + return False + + def url(self, path: str) -> str: + """ + Get the HTTP URL of the given path. + + Args: + path (`str`): + Path to get URL for. + + Returns: + `str`: HTTP URL to access the file or directory on the Hub. + """ + resolved_path = self.resolve_path(path) + url = hf_hub_url( + resolved_path.repo_id, + resolved_path.path_in_repo, + repo_type=resolved_path.repo_type, + revision=resolved_path.revision, + endpoint=self.endpoint, + ) + if self.isdir(path): + url = url.replace("/resolve/", "/tree/", 1) + return url + + def get_file(self, rpath, lpath, callback=_DEFAULT_CALLBACK, outfile=None, **kwargs) -> None: + """ + Copy single remote file to local. + + + + Note: When possible, use `HfApi.hf_hub_download()` for better performance. + + + + Args: + rpath (`str`): + Remote path to download from. + lpath (`str`): + Local path to download to. + callback (`Callback`, *optional*): + Optional callback to track download progress. Defaults to no callback. + outfile (`IO`, *optional*): + Optional file-like object to write to. If provided, `lpath` is ignored. + + """ + revision = kwargs.get("revision") + unhandled_kwargs = set(kwargs.keys()) - {"revision"} + if not isinstance(callback, (NoOpCallback, TqdmCallback)) or len(unhandled_kwargs) > 0: + # for now, let's not handle custom callbacks + # and let's not handle custom kwargs + return super().get_file(rpath, lpath, callback=callback, outfile=outfile, **kwargs) + + # Taken from https://github.com/fsspec/filesystem_spec/blob/47b445ae4c284a82dd15e0287b1ffc410e8fc470/fsspec/spec.py#L883 + if isfilelike(lpath): + outfile = lpath + elif self.isdir(rpath): + os.makedirs(lpath, exist_ok=True) + return None + + if isinstance(lpath, (str, Path)): # otherwise, let's assume it's a file-like object + os.makedirs(os.path.dirname(lpath), exist_ok=True) + + # Open file if not already open + close_file = False + if outfile is None: + outfile = open(lpath, "wb") + close_file = True + initial_pos = outfile.tell() + + # Custom implementation of `get_file` to use `http_get`. + resolve_remote_path = self.resolve_path(rpath, revision=revision) + expected_size = self.info(rpath, revision=revision)["size"] + callback.set_size(expected_size) + try: + http_get( + url=hf_hub_url( + repo_id=resolve_remote_path.repo_id, + revision=resolve_remote_path.revision, + filename=resolve_remote_path.path_in_repo, + repo_type=resolve_remote_path.repo_type, + endpoint=self.endpoint, + ), + temp_file=outfile, + displayed_filename=rpath, + expected_size=expected_size, + resume_size=0, + headers=self._api._build_hf_headers(), + _tqdm_bar=callback.tqdm if isinstance(callback, TqdmCallback) else None, + ) + outfile.seek(initial_pos) + finally: + # Close file only if we opened it ourselves + if close_file: + outfile.close() + + @property + def transaction(self): + """A context within which files are committed together upon exit + + Requires the file class to implement `.commit()` and `.discard()` + for the normal and exception cases. + """ + # Taken from https://github.com/fsspec/filesystem_spec/blob/3fbb6fee33b46cccb015607630843dea049d3243/fsspec/spec.py#L231 + # See https://github.com/huggingface/huggingface_hub/issues/1733 + raise NotImplementedError("Transactional commits are not supported.") + + def start_transaction(self): + """Begin write transaction for deferring files, non-context version""" + # Taken from https://github.com/fsspec/filesystem_spec/blob/3fbb6fee33b46cccb015607630843dea049d3243/fsspec/spec.py#L241 + # See https://github.com/huggingface/huggingface_hub/issues/1733 + raise NotImplementedError("Transactional commits are not supported.") + + +class HfFileSystemFile(fsspec.spec.AbstractBufferedFile): + def __init__(self, fs: HfFileSystem, path: str, revision: Optional[str] = None, **kwargs): + try: + self.resolved_path = fs.resolve_path(path, revision=revision) + except FileNotFoundError as e: + if "w" in kwargs.get("mode", ""): + raise FileNotFoundError( + f"{e}.\nMake sure the repository and revision exist before writing data." + ) from e + raise + # avoid an unnecessary .info() call with expensive expand_info=True to instantiate .details + if kwargs.get("mode", "rb") == "rb": + self.details = fs.info(self.resolved_path.unresolve(), expand_info=False) + super().__init__(fs, self.resolved_path.unresolve(), **kwargs) + self.fs: HfFileSystem + + def __del__(self): + if not hasattr(self, "resolved_path"): + # Means that the constructor failed. Nothing to do. + return + return super().__del__() + + def _fetch_range(self, start: int, end: int) -> bytes: + headers = { + "range": f"bytes={start}-{end - 1}", + **self.fs._api._build_hf_headers(), + } + url = hf_hub_url( + repo_id=self.resolved_path.repo_id, + revision=self.resolved_path.revision, + filename=self.resolved_path.path_in_repo, + repo_type=self.resolved_path.repo_type, + endpoint=self.fs.endpoint, + ) + r = http_backoff( + "GET", + url, + headers=headers, + retry_on_status_codes=(500, 502, 503, 504), + timeout=constants.HF_HUB_DOWNLOAD_TIMEOUT, + ) + hf_raise_for_status(r) + return r.content + + def _initiate_upload(self) -> None: + self.temp_file = tempfile.NamedTemporaryFile(prefix="hffs-", delete=False) + + def _upload_chunk(self, final: bool = False) -> None: + self.buffer.seek(0) + block = self.buffer.read() + self.temp_file.write(block) + if final: + self.temp_file.close() + self.fs._api.upload_file( + path_or_fileobj=self.temp_file.name, + path_in_repo=self.resolved_path.path_in_repo, + repo_id=self.resolved_path.repo_id, + token=self.fs.token, + repo_type=self.resolved_path.repo_type, + revision=self.resolved_path.revision, + commit_message=self.kwargs.get("commit_message"), + commit_description=self.kwargs.get("commit_description"), + ) + os.remove(self.temp_file.name) + self.fs.invalidate_cache( + path=self.resolved_path.unresolve(), + ) + + def read(self, length=-1): + """Read remote file. + + If `length` is not provided or is -1, the entire file is downloaded and read. On POSIX systems and if + `hf_transfer` is not enabled, the file is loaded in memory directly. Otherwise, the file is downloaded to a + temporary file and read from there. + """ + if self.mode == "rb" and (length is None or length == -1) and self.loc == 0: + with self.fs.open(self.path, "rb", block_size=0) as f: # block_size=0 enables fast streaming + out = f.read() + self.loc += len(out) + return out + return super().read(length) + + def url(self) -> str: + return self.fs.url(self.path) + + +class HfFileSystemStreamFile(fsspec.spec.AbstractBufferedFile): + def __init__( + self, + fs: HfFileSystem, + path: str, + mode: str = "rb", + revision: Optional[str] = None, + block_size: int = 0, + cache_type: str = "none", + **kwargs, + ): + if block_size != 0: + raise ValueError(f"HfFileSystemStreamFile only supports block_size=0 but got {block_size}") + if cache_type != "none": + raise ValueError(f"HfFileSystemStreamFile only supports cache_type='none' but got {cache_type}") + if "w" in mode: + raise ValueError(f"HfFileSystemStreamFile only supports reading but got mode='{mode}'") + try: + self.resolved_path = fs.resolve_path(path, revision=revision) + except FileNotFoundError as e: + if "w" in kwargs.get("mode", ""): + raise FileNotFoundError( + f"{e}.\nMake sure the repository and revision exist before writing data." + ) from e + # avoid an unnecessary .info() call to instantiate .details + self.details = {"name": self.resolved_path.unresolve(), "size": None} + super().__init__( + fs, self.resolved_path.unresolve(), mode=mode, block_size=block_size, cache_type=cache_type, **kwargs + ) + self.response: Optional[Response] = None + self.fs: HfFileSystem + + def seek(self, loc: int, whence: int = 0): + if loc == 0 and whence == 1: + return + if loc == self.loc and whence == 0: + return + raise ValueError("Cannot seek streaming HF file") + + def read(self, length: int = -1): + read_args = (length,) if length >= 0 else () + if self.response is None: + url = hf_hub_url( + repo_id=self.resolved_path.repo_id, + revision=self.resolved_path.revision, + filename=self.resolved_path.path_in_repo, + repo_type=self.resolved_path.repo_type, + endpoint=self.fs.endpoint, + ) + self.response = http_backoff( + "GET", + url, + headers=self.fs._api._build_hf_headers(), + retry_on_status_codes=(500, 502, 503, 504), + stream=True, + timeout=constants.HF_HUB_DOWNLOAD_TIMEOUT, + ) + hf_raise_for_status(self.response) + try: + out = self.response.raw.read(*read_args) + except Exception: + self.response.close() + + # Retry by recreating the connection + url = hf_hub_url( + repo_id=self.resolved_path.repo_id, + revision=self.resolved_path.revision, + filename=self.resolved_path.path_in_repo, + repo_type=self.resolved_path.repo_type, + endpoint=self.fs.endpoint, + ) + self.response = http_backoff( + "GET", + url, + headers={"Range": "bytes=%d-" % self.loc, **self.fs._api._build_hf_headers()}, + retry_on_status_codes=(500, 502, 503, 504), + stream=True, + timeout=constants.HF_HUB_DOWNLOAD_TIMEOUT, + ) + hf_raise_for_status(self.response) + try: + out = self.response.raw.read(*read_args) + except Exception: + self.response.close() + raise + self.loc += len(out) + return out + + def url(self) -> str: + return self.fs.url(self.path) + + def __del__(self): + if not hasattr(self, "resolved_path"): + # Means that the constructor failed. Nothing to do. + return + return super().__del__() + + def __reduce__(self): + return reopen, (self.fs, self.path, self.mode, self.blocksize, self.cache.name) + + +def safe_revision(revision: str) -> str: + return revision if SPECIAL_REFS_REVISION_REGEX.match(revision) else safe_quote(revision) + + +def safe_quote(s: str) -> str: + return quote(s, safe="") + + +def _raise_file_not_found(path: str, err: Optional[Exception]) -> NoReturn: + msg = path + if isinstance(err, RepositoryNotFoundError): + msg = f"{path} (repository not found)" + elif isinstance(err, RevisionNotFoundError): + msg = f"{path} (revision not found)" + elif isinstance(err, HFValidationError): + msg = f"{path} (invalid repository id)" + raise FileNotFoundError(msg) from err + + +def reopen(fs: HfFileSystem, path: str, mode: str, block_size: int, cache_type: str): + return fs.open(path, mode=mode, block_size=block_size, cache_type=cache_type) diff --git a/lib/python3.12/site-packages/huggingface_hub/hub_mixin.py b/lib/python3.12/site-packages/huggingface_hub/hub_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..7cd8f490d9d4a5b6b73737ec9050959c11f7b391 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/hub_mixin.py @@ -0,0 +1,851 @@ +import inspect +import json +import os +from dataclasses import Field, asdict, dataclass, is_dataclass +from pathlib import Path +from typing import Any, Callable, ClassVar, Dict, List, Optional, Protocol, Tuple, Type, TypeVar, Union + +import packaging.version + +from . import constants +from .errors import EntryNotFoundError, HfHubHTTPError +from .file_download import hf_hub_download +from .hf_api import HfApi +from .repocard import ModelCard, ModelCardData +from .utils import ( + SoftTemporaryDirectory, + is_jsonable, + is_safetensors_available, + is_simple_optional_type, + is_torch_available, + logging, + unwrap_simple_optional_type, + validate_hf_hub_args, +) + + +if is_torch_available(): + import torch # type: ignore + +if is_safetensors_available(): + import safetensors + from safetensors.torch import load_model as load_model_as_safetensor + from safetensors.torch import save_model as save_model_as_safetensor + + +logger = logging.get_logger(__name__) + + +# Type alias for dataclass instances, copied from https://github.com/python/typeshed/blob/9f28171658b9ca6c32a7cb93fbb99fc92b17858b/stdlib/_typeshed/__init__.pyi#L349 +class DataclassInstance(Protocol): + __dataclass_fields__: ClassVar[Dict[str, Field]] + + +# Generic variable that is either ModelHubMixin or a subclass thereof +T = TypeVar("T", bound="ModelHubMixin") +# Generic variable to represent an args type +ARGS_T = TypeVar("ARGS_T") +ENCODER_T = Callable[[ARGS_T], Any] +DECODER_T = Callable[[Any], ARGS_T] +CODER_T = Tuple[ENCODER_T, DECODER_T] + + +DEFAULT_MODEL_CARD = """ +--- +# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1 +# Doc / guide: https://huggingface.co/docs/hub/model-cards +{{ card_data }} +--- + +This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration: +- Code: {{ repo_url | default("[More Information Needed]", true) }} +- Paper: {{ paper_url | default("[More Information Needed]", true) }} +- Docs: {{ docs_url | default("[More Information Needed]", true) }} +""" + + +@dataclass +class MixinInfo: + model_card_template: str + model_card_data: ModelCardData + docs_url: Optional[str] = None + paper_url: Optional[str] = None + repo_url: Optional[str] = None + + +class ModelHubMixin: + """ + A generic mixin to integrate ANY machine learning framework with the Hub. + + To integrate your framework, your model class must inherit from this class. Custom logic for saving/loading models + have to be overwritten in [`_from_pretrained`] and [`_save_pretrained`]. [`PyTorchModelHubMixin`] is a good example + of mixin integration with the Hub. Check out our [integration guide](../guides/integrations) for more instructions. + + When inheriting from [`ModelHubMixin`], you can define class-level attributes. These attributes are not passed to + `__init__` but to the class definition itself. This is useful to define metadata about the library integrating + [`ModelHubMixin`]. + + For more details on how to integrate the mixin with your library, checkout the [integration guide](../guides/integrations). + + Args: + repo_url (`str`, *optional*): + URL of the library repository. Used to generate model card. + paper_url (`str`, *optional*): + URL of the library paper. Used to generate model card. + docs_url (`str`, *optional*): + URL of the library documentation. Used to generate model card. + model_card_template (`str`, *optional*): + Template of the model card. Used to generate model card. Defaults to a generic template. + language (`str` or `List[str]`, *optional*): + Language supported by the library. Used to generate model card. + library_name (`str`, *optional*): + Name of the library integrating ModelHubMixin. Used to generate model card. + license (`str`, *optional*): + License of the library integrating ModelHubMixin. Used to generate model card. + E.g: "apache-2.0" + license_name (`str`, *optional*): + Name of the library integrating ModelHubMixin. Used to generate model card. + Only used if `license` is set to `other`. + E.g: "coqui-public-model-license". + license_link (`str`, *optional*): + URL to the license of the library integrating ModelHubMixin. Used to generate model card. + Only used if `license` is set to `other` and `license_name` is set. + E.g: "https://coqui.ai/cpml". + pipeline_tag (`str`, *optional*): + Tag of the pipeline. Used to generate model card. E.g. "text-classification". + tags (`List[str]`, *optional*): + Tags to be added to the model card. Used to generate model card. E.g. ["computer-vision"] + coders (`Dict[Type, Tuple[Callable, Callable]]`, *optional*): + Dictionary of custom types and their encoders/decoders. Used to encode/decode arguments that are not + jsonable by default. E.g dataclasses, argparse.Namespace, OmegaConf, etc. + + Example: + + ```python + >>> from huggingface_hub import ModelHubMixin + + # Inherit from ModelHubMixin + >>> class MyCustomModel( + ... ModelHubMixin, + ... library_name="my-library", + ... tags=["computer-vision"], + ... repo_url="https://github.com/huggingface/my-cool-library", + ... paper_url="https://arxiv.org/abs/2304.12244", + ... docs_url="https://huggingface.co/docs/my-cool-library", + ... # ^ optional metadata to generate model card + ... ): + ... def __init__(self, size: int = 512, device: str = "cpu"): + ... # define how to initialize your model + ... super().__init__() + ... ... + ... + ... def _save_pretrained(self, save_directory: Path) -> None: + ... # define how to serialize your model + ... ... + ... + ... @classmethod + ... def from_pretrained( + ... cls: Type[T], + ... pretrained_model_name_or_path: Union[str, Path], + ... *, + ... force_download: bool = False, + ... resume_download: Optional[bool] = None, + ... proxies: Optional[Dict] = None, + ... token: Optional[Union[str, bool]] = None, + ... cache_dir: Optional[Union[str, Path]] = None, + ... local_files_only: bool = False, + ... revision: Optional[str] = None, + ... **model_kwargs, + ... ) -> T: + ... # define how to deserialize your model + ... ... + + >>> model = MyCustomModel(size=256, device="gpu") + + # Save model weights to local directory + >>> model.save_pretrained("my-awesome-model") + + # Push model weights to the Hub + >>> model.push_to_hub("my-awesome-model") + + # Download and initialize weights from the Hub + >>> reloaded_model = MyCustomModel.from_pretrained("username/my-awesome-model") + >>> reloaded_model.size + 256 + + # Model card has been correctly populated + >>> from huggingface_hub import ModelCard + >>> card = ModelCard.load("username/my-awesome-model") + >>> card.data.tags + ["x-custom-tag", "pytorch_model_hub_mixin", "model_hub_mixin"] + >>> card.data.library_name + "my-library" + ``` + """ + + _hub_mixin_config: Optional[Union[dict, DataclassInstance]] = None + # ^ optional config attribute automatically set in `from_pretrained` + _hub_mixin_info: MixinInfo + # ^ information about the library integrating ModelHubMixin (used to generate model card) + _hub_mixin_inject_config: bool # whether `_from_pretrained` expects `config` or not + _hub_mixin_init_parameters: Dict[str, inspect.Parameter] # __init__ parameters + _hub_mixin_jsonable_default_values: Dict[str, Any] # default values for __init__ parameters + _hub_mixin_jsonable_custom_types: Tuple[Type, ...] # custom types that can be encoded/decoded + _hub_mixin_coders: Dict[Type, CODER_T] # encoders/decoders for custom types + # ^ internal values to handle config + + def __init_subclass__( + cls, + *, + # Generic info for model card + repo_url: Optional[str] = None, + paper_url: Optional[str] = None, + docs_url: Optional[str] = None, + # Model card template + model_card_template: str = DEFAULT_MODEL_CARD, + # Model card metadata + language: Optional[List[str]] = None, + library_name: Optional[str] = None, + license: Optional[str] = None, + license_name: Optional[str] = None, + license_link: Optional[str] = None, + pipeline_tag: Optional[str] = None, + tags: Optional[List[str]] = None, + # How to encode/decode arguments with custom type into a JSON config? + coders: Optional[ + Dict[Type, CODER_T] + # Key is a type. + # Value is a tuple (encoder, decoder). + # Example: {MyCustomType: (lambda x: x.value, lambda data: MyCustomType(data))} + ] = None, + ) -> None: + """Inspect __init__ signature only once when subclassing + handle modelcard.""" + super().__init_subclass__() + + # Will be reused when creating modelcard + tags = tags or [] + tags.append("model_hub_mixin") + + # Initialize MixinInfo if not existent + info = MixinInfo(model_card_template=model_card_template, model_card_data=ModelCardData()) + + # If parent class has a MixinInfo, inherit from it as a copy + if hasattr(cls, "_hub_mixin_info"): + # Inherit model card template from parent class if not explicitly set + if model_card_template == DEFAULT_MODEL_CARD: + info.model_card_template = cls._hub_mixin_info.model_card_template + + # Inherit from parent model card data + info.model_card_data = ModelCardData(**cls._hub_mixin_info.model_card_data.to_dict()) + + # Inherit other info + info.docs_url = cls._hub_mixin_info.docs_url + info.paper_url = cls._hub_mixin_info.paper_url + info.repo_url = cls._hub_mixin_info.repo_url + cls._hub_mixin_info = info + + # Update MixinInfo with metadata + if model_card_template is not None and model_card_template != DEFAULT_MODEL_CARD: + info.model_card_template = model_card_template + if repo_url is not None: + info.repo_url = repo_url + if paper_url is not None: + info.paper_url = paper_url + if docs_url is not None: + info.docs_url = docs_url + if language is not None: + info.model_card_data.language = language + if library_name is not None: + info.model_card_data.library_name = library_name + if license is not None: + info.model_card_data.license = license + if license_name is not None: + info.model_card_data.license_name = license_name + if license_link is not None: + info.model_card_data.license_link = license_link + if pipeline_tag is not None: + info.model_card_data.pipeline_tag = pipeline_tag + if tags is not None: + if info.model_card_data.tags is not None: + info.model_card_data.tags.extend(tags) + else: + info.model_card_data.tags = tags + + info.model_card_data.tags = sorted(set(info.model_card_data.tags)) + + # Handle encoders/decoders for args + cls._hub_mixin_coders = coders or {} + cls._hub_mixin_jsonable_custom_types = tuple(cls._hub_mixin_coders.keys()) + + # Inspect __init__ signature to handle config + cls._hub_mixin_init_parameters = dict(inspect.signature(cls.__init__).parameters) + cls._hub_mixin_jsonable_default_values = { + param.name: cls._encode_arg(param.default) + for param in cls._hub_mixin_init_parameters.values() + if param.default is not inspect.Parameter.empty and cls._is_jsonable(param.default) + } + cls._hub_mixin_inject_config = "config" in inspect.signature(cls._from_pretrained).parameters + + def __new__(cls: Type[T], *args, **kwargs) -> T: + """Create a new instance of the class and handle config. + + 3 cases: + - If `self._hub_mixin_config` is already set, do nothing. + - If `config` is passed as a dataclass, set it as `self._hub_mixin_config`. + - Otherwise, build `self._hub_mixin_config` from default values and passed values. + """ + instance = super().__new__(cls) + + # If `config` is already set, return early + if instance._hub_mixin_config is not None: + return instance + + # Infer passed values + passed_values = { + **{ + key: value + for key, value in zip( + # [1:] to skip `self` parameter + list(cls._hub_mixin_init_parameters)[1:], + args, + ) + }, + **kwargs, + } + + # If config passed as dataclass => set it and return early + if is_dataclass(passed_values.get("config")): + instance._hub_mixin_config = passed_values["config"] + return instance + + # Otherwise, build config from default + passed values + init_config = { + # default values + **cls._hub_mixin_jsonable_default_values, + # passed values + **{ + key: cls._encode_arg(value) # Encode custom types as jsonable value + for key, value in passed_values.items() + if instance._is_jsonable(value) # Only if jsonable or we have a custom encoder + }, + } + passed_config = init_config.pop("config", {}) + + # Populate `init_config` with provided config + if isinstance(passed_config, dict): + init_config.update(passed_config) + + # Set `config` attribute and return + if init_config != {}: + instance._hub_mixin_config = init_config + return instance + + @classmethod + def _is_jsonable(cls, value: Any) -> bool: + """Check if a value is JSON serializable.""" + if is_dataclass(value): + return True + if isinstance(value, cls._hub_mixin_jsonable_custom_types): + return True + return is_jsonable(value) + + @classmethod + def _encode_arg(cls, arg: Any) -> Any: + """Encode an argument into a JSON serializable format.""" + if is_dataclass(arg): + return asdict(arg) + for type_, (encoder, _) in cls._hub_mixin_coders.items(): + if isinstance(arg, type_): + if arg is None: + return None + return encoder(arg) + return arg + + @classmethod + def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> Optional[ARGS_T]: + """Decode a JSON serializable value into an argument.""" + if is_simple_optional_type(expected_type): + if value is None: + return None + expected_type = unwrap_simple_optional_type(expected_type) + # Dataclass => handle it + if is_dataclass(expected_type): + return _load_dataclass(expected_type, value) # type: ignore[return-value] + # Otherwise => check custom decoders + for type_, (_, decoder) in cls._hub_mixin_coders.items(): + if inspect.isclass(expected_type) and issubclass(expected_type, type_): + return decoder(value) + # Otherwise => don't decode + return value + + def save_pretrained( + self, + save_directory: Union[str, Path], + *, + config: Optional[Union[dict, DataclassInstance]] = None, + repo_id: Optional[str] = None, + push_to_hub: bool = False, + model_card_kwargs: Optional[Dict[str, Any]] = None, + **push_to_hub_kwargs, + ) -> Optional[str]: + """ + Save weights in local directory. + + Args: + save_directory (`str` or `Path`): + Path to directory in which the model weights and configuration will be saved. + config (`dict` or `DataclassInstance`, *optional*): + Model configuration specified as a key/value dictionary or a dataclass instance. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Huggingface Hub after saving it. + repo_id (`str`, *optional*): + ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if + not provided. + model_card_kwargs (`Dict[str, Any]`, *optional*): + Additional arguments passed to the model card template to customize the model card. + push_to_hub_kwargs: + Additional key word arguments passed along to the [`~ModelHubMixin.push_to_hub`] method. + Returns: + `str` or `None`: url of the commit on the Hub if `push_to_hub=True`, `None` otherwise. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + + # Remove config.json if already exists. After `_save_pretrained` we don't want to overwrite config.json + # as it might have been saved by the custom `_save_pretrained` already. However we do want to overwrite + # an existing config.json if it was not saved by `_save_pretrained`. + config_path = save_directory / constants.CONFIG_NAME + config_path.unlink(missing_ok=True) + + # save model weights/files (framework-specific) + self._save_pretrained(save_directory) + + # save config (if provided and if not serialized yet in `_save_pretrained`) + if config is None: + config = self._hub_mixin_config + if config is not None: + if is_dataclass(config): + config = asdict(config) # type: ignore[arg-type] + if not config_path.exists(): + config_str = json.dumps(config, sort_keys=True, indent=2) + config_path.write_text(config_str) + + # save model card + model_card_path = save_directory / "README.md" + model_card_kwargs = model_card_kwargs if model_card_kwargs is not None else {} + if not model_card_path.exists(): # do not overwrite if already exists + self.generate_model_card(**model_card_kwargs).save(save_directory / "README.md") + + # push to the Hub if required + if push_to_hub: + kwargs = push_to_hub_kwargs.copy() # soft-copy to avoid mutating input + if config is not None: # kwarg for `push_to_hub` + kwargs["config"] = config + if repo_id is None: + repo_id = save_directory.name # Defaults to `save_directory` name + return self.push_to_hub(repo_id=repo_id, model_card_kwargs=model_card_kwargs, **kwargs) + return None + + def _save_pretrained(self, save_directory: Path) -> None: + """ + Overwrite this method in subclass to define how to save your model. + Check out our [integration guide](../guides/integrations) for instructions. + + Args: + save_directory (`str` or `Path`): + Path to directory in which the model weights and configuration will be saved. + """ + raise NotImplementedError + + @classmethod + @validate_hf_hub_args + def from_pretrained( + cls: Type[T], + pretrained_model_name_or_path: Union[str, Path], + *, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[Dict] = None, + token: Optional[Union[str, bool]] = None, + cache_dir: Optional[Union[str, Path]] = None, + local_files_only: bool = False, + revision: Optional[str] = None, + **model_kwargs, + ) -> T: + """ + Download a model from the Huggingface Hub and instantiate it. + + Args: + pretrained_model_name_or_path (`str`, `Path`): + - Either the `model_id` (string) of a model hosted on the Hub, e.g. `bigscience/bloom`. + - Or a path to a `directory` containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `../path/to/my_model_directory/`. + revision (`str`, *optional*): + Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. + Defaults to the latest commit on `main` branch. + force_download (`bool`, *optional*, defaults to `False`): + Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding + the existing cache. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. By default, it will use the token + cached when running `huggingface-cli login`. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the local cached file if it exists. + model_kwargs (`Dict`, *optional*): + Additional kwargs to pass to the model during initialization. + """ + model_id = str(pretrained_model_name_or_path) + config_file: Optional[str] = None + if os.path.isdir(model_id): + if constants.CONFIG_NAME in os.listdir(model_id): + config_file = os.path.join(model_id, constants.CONFIG_NAME) + else: + logger.warning(f"{constants.CONFIG_NAME} not found in {Path(model_id).resolve()}") + else: + try: + config_file = hf_hub_download( + repo_id=model_id, + filename=constants.CONFIG_NAME, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + except HfHubHTTPError as e: + logger.info(f"{constants.CONFIG_NAME} not found on the HuggingFace Hub: {str(e)}") + + # Read config + config = None + if config_file is not None: + with open(config_file, "r", encoding="utf-8") as f: + config = json.load(f) + + # Decode custom types in config + for key, value in config.items(): + if key in cls._hub_mixin_init_parameters: + expected_type = cls._hub_mixin_init_parameters[key].annotation + if expected_type is not inspect.Parameter.empty: + config[key] = cls._decode_arg(expected_type, value) + + # Populate model_kwargs from config + for param in cls._hub_mixin_init_parameters.values(): + if param.name not in model_kwargs and param.name in config: + model_kwargs[param.name] = config[param.name] + + # Check if `config` argument was passed at init + if "config" in cls._hub_mixin_init_parameters and "config" not in model_kwargs: + # Decode `config` argument if it was passed + config_annotation = cls._hub_mixin_init_parameters["config"].annotation + config = cls._decode_arg(config_annotation, config) + + # Forward config to model initialization + model_kwargs["config"] = config + + # Inject config if `**kwargs` are expected + if is_dataclass(cls): + for key in cls.__dataclass_fields__: + if key not in model_kwargs and key in config: + model_kwargs[key] = config[key] + elif any(param.kind == inspect.Parameter.VAR_KEYWORD for param in cls._hub_mixin_init_parameters.values()): + for key, value in config.items(): + if key not in model_kwargs: + model_kwargs[key] = value + + # Finally, also inject if `_from_pretrained` expects it + if cls._hub_mixin_inject_config and "config" not in model_kwargs: + model_kwargs["config"] = config + + instance = cls._from_pretrained( + model_id=str(model_id), + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + **model_kwargs, + ) + + # Implicitly set the config as instance attribute if not already set by the class + # This way `config` will be available when calling `save_pretrained` or `push_to_hub`. + if config is not None and (getattr(instance, "_hub_mixin_config", None) in (None, {})): + instance._hub_mixin_config = config + + return instance + + @classmethod + def _from_pretrained( + cls: Type[T], + *, + model_id: str, + revision: Optional[str], + cache_dir: Optional[Union[str, Path]], + force_download: bool, + proxies: Optional[Dict], + resume_download: Optional[bool], + local_files_only: bool, + token: Optional[Union[str, bool]], + **model_kwargs, + ) -> T: + """Overwrite this method in subclass to define how to load your model from pretrained. + + Use [`hf_hub_download`] or [`snapshot_download`] to download files from the Hub before loading them. Most + args taken as input can be directly passed to those 2 methods. If needed, you can add more arguments to this + method using "model_kwargs". For example [`PyTorchModelHubMixin._from_pretrained`] takes as input a `map_location` + parameter to set on which device the model should be loaded. + + Check out our [integration guide](../guides/integrations) for more instructions. + + Args: + model_id (`str`): + ID of the model to load from the Huggingface Hub (e.g. `bigscience/bloom`). + revision (`str`, *optional*): + Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the + latest commit on `main` branch. + force_download (`bool`, *optional*, defaults to `False`): + Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding + the existing cache. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint (e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`). + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. By default, it will use the token + cached when running `huggingface-cli login`. + cache_dir (`str`, `Path`, *optional*): + Path to the folder where cached files are stored. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, avoid downloading the file and return the path to the local cached file if it exists. + model_kwargs: + Additional keyword arguments passed along to the [`~ModelHubMixin._from_pretrained`] method. + """ + raise NotImplementedError + + @validate_hf_hub_args + def push_to_hub( + self, + repo_id: str, + *, + config: Optional[Union[dict, DataclassInstance]] = None, + commit_message: str = "Push model using huggingface_hub.", + private: Optional[bool] = None, + token: Optional[str] = None, + branch: Optional[str] = None, + create_pr: Optional[bool] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + model_card_kwargs: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Upload model checkpoint to the Hub. + + Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use + `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more + details. + + Args: + repo_id (`str`): + ID of the repository to push to (example: `"username/my-model"`). + config (`dict` or `DataclassInstance`, *optional*): + Model configuration specified as a key/value dictionary or a dataclass instance. + commit_message (`str`, *optional*): + Message to commit while pushing. + private (`bool`, *optional*): + Whether the repository created should be private. + If `None` (default), the repo will be public unless the organization's default is private. + token (`str`, *optional*): + The token to use as HTTP bearer authorization for remote files. By default, it will use the token + cached when running `huggingface-cli login`. + branch (`str`, *optional*): + The git branch on which to push the model. This defaults to `"main"`. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are pushed. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not pushed. + delete_patterns (`List[str]` or `str`, *optional*): + If provided, remote files matching any of the patterns will be deleted from the repo. + model_card_kwargs (`Dict[str, Any]`, *optional*): + Additional arguments passed to the model card template to customize the model card. + + Returns: + The url of the commit of your model in the given repository. + """ + api = HfApi(token=token) + repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id + + # Push the files to the repo in a single commit + with SoftTemporaryDirectory() as tmp: + saved_path = Path(tmp) / repo_id + self.save_pretrained(saved_path, config=config, model_card_kwargs=model_card_kwargs) + return api.upload_folder( + repo_id=repo_id, + repo_type="model", + folder_path=saved_path, + commit_message=commit_message, + revision=branch, + create_pr=create_pr, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + delete_patterns=delete_patterns, + ) + + def generate_model_card(self, *args, **kwargs) -> ModelCard: + card = ModelCard.from_template( + card_data=self._hub_mixin_info.model_card_data, + template_str=self._hub_mixin_info.model_card_template, + repo_url=self._hub_mixin_info.repo_url, + paper_url=self._hub_mixin_info.paper_url, + docs_url=self._hub_mixin_info.docs_url, + **kwargs, + ) + return card + + +class PyTorchModelHubMixin(ModelHubMixin): + """ + Implementation of [`ModelHubMixin`] to provide model Hub upload/download capabilities to PyTorch models. The model + is set in evaluation mode by default using `model.eval()` (dropout modules are deactivated). To train the model, + you should first set it back in training mode with `model.train()`. + + See [`ModelHubMixin`] for more details on how to use the mixin. + + Example: + + ```python + >>> import torch + >>> import torch.nn as nn + >>> from huggingface_hub import PyTorchModelHubMixin + + >>> class MyModel( + ... nn.Module, + ... PyTorchModelHubMixin, + ... library_name="keras-nlp", + ... repo_url="https://github.com/keras-team/keras-nlp", + ... paper_url="https://arxiv.org/abs/2304.12244", + ... docs_url="https://keras.io/keras_nlp/", + ... # ^ optional metadata to generate model card + ... ): + ... def __init__(self, hidden_size: int = 512, vocab_size: int = 30000, output_size: int = 4): + ... super().__init__() + ... self.param = nn.Parameter(torch.rand(hidden_size, vocab_size)) + ... self.linear = nn.Linear(output_size, vocab_size) + + ... def forward(self, x): + ... return self.linear(x + self.param) + >>> model = MyModel(hidden_size=256) + + # Save model weights to local directory + >>> model.save_pretrained("my-awesome-model") + + # Push model weights to the Hub + >>> model.push_to_hub("my-awesome-model") + + # Download and initialize weights from the Hub + >>> model = MyModel.from_pretrained("username/my-awesome-model") + >>> model.hidden_size + 256 + ``` + """ + + def __init_subclass__(cls, *args, tags: Optional[List[str]] = None, **kwargs) -> None: + tags = tags or [] + tags.append("pytorch_model_hub_mixin") + kwargs["tags"] = tags + return super().__init_subclass__(*args, **kwargs) + + def _save_pretrained(self, save_directory: Path) -> None: + """Save weights from a Pytorch model to a local directory.""" + model_to_save = self.module if hasattr(self, "module") else self # type: ignore + save_model_as_safetensor(model_to_save, str(save_directory / constants.SAFETENSORS_SINGLE_FILE)) + + @classmethod + def _from_pretrained( + cls, + *, + model_id: str, + revision: Optional[str], + cache_dir: Optional[Union[str, Path]], + force_download: bool, + proxies: Optional[Dict], + resume_download: Optional[bool], + local_files_only: bool, + token: Union[str, bool, None], + map_location: str = "cpu", + strict: bool = False, + **model_kwargs, + ): + """Load Pytorch pretrained weights and return the loaded model.""" + model = cls(**model_kwargs) + if os.path.isdir(model_id): + print("Loading weights from local directory") + model_file = os.path.join(model_id, constants.SAFETENSORS_SINGLE_FILE) + return cls._load_as_safetensor(model, model_file, map_location, strict) + else: + try: + model_file = hf_hub_download( + repo_id=model_id, + filename=constants.SAFETENSORS_SINGLE_FILE, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + return cls._load_as_safetensor(model, model_file, map_location, strict) + except EntryNotFoundError: + model_file = hf_hub_download( + repo_id=model_id, + filename=constants.PYTORCH_WEIGHTS_NAME, + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + return cls._load_as_pickle(model, model_file, map_location, strict) + + @classmethod + def _load_as_pickle(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: + state_dict = torch.load(model_file, map_location=torch.device(map_location), weights_only=True) + model.load_state_dict(state_dict, strict=strict) # type: ignore + model.eval() # type: ignore + return model + + @classmethod + def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: + if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"): # type: ignore [attr-defined] + load_model_as_safetensor(model, model_file, strict=strict) # type: ignore [arg-type] + if map_location != "cpu": + logger.warning( + "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors." + " This means that the model is loaded on 'cpu' first and then copied to the device." + " This leads to a slower loading time." + " Please update safetensors to version 0.4.3 or above for improved performance." + ) + model.to(map_location) # type: ignore [attr-defined] + else: + safetensors.torch.load_model(model, model_file, strict=strict, device=map_location) # type: ignore [arg-type] + return model + + +def _load_dataclass(datacls: Type[DataclassInstance], data: dict) -> DataclassInstance: + """Load a dataclass instance from a dictionary. + + Fields not expected by the dataclass are ignored. + """ + return datacls(**{k: v for k, v in data.items() if k in datacls.__dataclass_fields__}) diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/__init__.py b/lib/python3.12/site-packages/huggingface_hub/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13da39d5da3f52d6663a56aa924eb63cbb4c5b08 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/_common.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5397c201c865af20e59cc5f8376d896e3858c7e4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/__pycache__/_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_client.py b/lib/python3.12/site-packages/huggingface_hub/inference/_client.py new file mode 100644 index 0000000000000000000000000000000000000000..793e29e3ac23d99a7cc4b3418b0b757a78bd064b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_client.py @@ -0,0 +1,3474 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# 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. +# +# Related resources: +# https://huggingface.co/tasks +# https://huggingface.co/docs/huggingface.js/inference/README +# https://github.com/huggingface/huggingface.js/tree/main/packages/inference/src +# https://github.com/huggingface/text-generation-inference/tree/main/clients/python +# https://github.com/huggingface/text-generation-inference/blob/main/clients/python/text_generation/client.py +# https://huggingface.slack.com/archives/C03E4DQ9LAJ/p1680169099087869 +# https://github.com/huggingface/unity-api#tasks +# +# Some TODO: +# - add all tasks +# +# NOTE: the philosophy of this client is "let's make it as easy as possible to use it, even if less optimized". Some +# examples of how it translates: +# - Timeout / Server unavailable is handled by the client in a single "timeout" parameter. +# - Files can be provided as bytes, file paths, or URLs and the client will try to "guess" the type. +# - Images are parsed as PIL.Image for easier manipulation. +# - Provides a "recommended model" for each task => suboptimal but user-wise quicker to get a first script running. +# - Only the main parameters are publicly exposed. Power users can always read the docs for more options. +import base64 +import logging +import re +import warnings +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Optional, Union, overload + +from requests import HTTPError + +from huggingface_hub import constants +from huggingface_hub.errors import BadRequestError, InferenceTimeoutError +from huggingface_hub.inference._common import ( + TASKS_EXPECTING_IMAGES, + ContentT, + ModelStatus, + RequestParameters, + _b64_encode, + _b64_to_image, + _bytes_to_dict, + _bytes_to_image, + _bytes_to_list, + _get_unsupported_text_generation_kwargs, + _import_numpy, + _open_as_binary, + _set_unsupported_text_generation_kwargs, + _stream_chat_completion_response, + _stream_text_generation_response, + raise_text_generation_error, +) +from huggingface_hub.inference._generated.types import ( + AudioClassificationOutputElement, + AudioClassificationOutputTransform, + AudioToAudioOutputElement, + AutomaticSpeechRecognitionOutput, + ChatCompletionInputGrammarType, + ChatCompletionInputStreamOptions, + ChatCompletionInputTool, + ChatCompletionInputToolChoiceClass, + ChatCompletionInputToolChoiceEnum, + ChatCompletionOutput, + ChatCompletionStreamOutput, + DocumentQuestionAnsweringOutputElement, + FillMaskOutputElement, + ImageClassificationOutputElement, + ImageClassificationOutputTransform, + ImageSegmentationOutputElement, + ImageSegmentationSubtask, + ImageToImageTargetSize, + ImageToTextOutput, + ObjectDetectionOutputElement, + Padding, + QuestionAnsweringOutputElement, + SummarizationOutput, + SummarizationTruncationStrategy, + TableQuestionAnsweringOutputElement, + TextClassificationOutputElement, + TextClassificationOutputTransform, + TextGenerationInputGrammarType, + TextGenerationOutput, + TextGenerationStreamOutput, + TextToSpeechEarlyStoppingEnum, + TokenClassificationAggregationStrategy, + TokenClassificationOutputElement, + TranslationOutput, + TranslationTruncationStrategy, + VisualQuestionAnsweringOutputElement, + ZeroShotClassificationOutputElement, + ZeroShotImageClassificationOutputElement, +) +from huggingface_hub.inference._providers import PROVIDER_T, get_provider_helper +from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status +from huggingface_hub.utils._auth import get_token +from huggingface_hub.utils._deprecation import _deprecate_method + + +if TYPE_CHECKING: + import numpy as np + from PIL.Image import Image + +logger = logging.getLogger(__name__) + + +MODEL_KWARGS_NOT_USED_REGEX = re.compile(r"The following `model_kwargs` are not used by the model: \[(.*?)\]") + + +class InferenceClient: + """ + Initialize a new Inference Client. + + [`InferenceClient`] aims to provide a unified experience to perform inference. The client can be used + seamlessly with either the (free) Inference API, self-hosted Inference Endpoints, or third-party Inference Providers. + + Args: + model (`str`, `optional`): + The model to run inference with. Can be a model id hosted on the Hugging Face Hub, e.g. `meta-llama/Meta-Llama-3-8B-Instruct` + or a URL to a deployed Inference Endpoint. Defaults to None, in which case a recommended model is + automatically selected for the task. + Note: for better compatibility with OpenAI's client, `model` has been aliased as `base_url`. Those 2 + arguments are mutually exclusive. If using `base_url` for chat completion, the `/chat/completions` suffix + path will be appended to the base URL (see the [TGI Messages API](https://huggingface.co/docs/text-generation-inference/en/messages_api) + documentation for details). When passing a URL as `model`, the client will not append any suffix path to it. + provider (`str`, *optional*): + Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"fireworks-ai"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"openai"`, `"replicate"`, "sambanova"` or `"together"`. + Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers. + If model is a URL or `base_url` is passed, then `provider` is not used. + token (`str`, *optional*): + Hugging Face token. Will default to the locally saved token if not provided. + Note: for better compatibility with OpenAI's client, `token` has been aliased as `api_key`. Those 2 + arguments are mutually exclusive and have the exact same behavior. + timeout (`float`, `optional`): + The maximum number of seconds to wait for a response from the server. Defaults to None, meaning it will loop until the server is available. + headers (`Dict[str, str]`, `optional`): + Additional headers to send to the server. By default only the authorization and user-agent headers are sent. + Values in this dictionary will override the default values. + bill_to (`str`, `optional`): + The billing account to use for the requests. By default the requests are billed on the user's account. + Requests can only be billed to an organization the user is a member of, and which has subscribed to Enterprise Hub. + cookies (`Dict[str, str]`, `optional`): + Additional cookies to send to the server. + proxies (`Any`, `optional`): + Proxies to use for the request. + base_url (`str`, `optional`): + Base URL to run inference. This is a duplicated argument from `model` to make [`InferenceClient`] + follow the same pattern as `openai.OpenAI` client. Cannot be used if `model` is set. Defaults to None. + api_key (`str`, `optional`): + Token to use for authentication. This is a duplicated argument from `token` to make [`InferenceClient`] + follow the same pattern as `openai.OpenAI` client. Cannot be used if `token` is set. Defaults to None. + """ + + def __init__( + self, + model: Optional[str] = None, + *, + provider: Union[Literal["auto"], PROVIDER_T, None] = None, + token: Optional[str] = None, + timeout: Optional[float] = None, + headers: Optional[Dict[str, str]] = None, + cookies: Optional[Dict[str, str]] = None, + proxies: Optional[Any] = None, + bill_to: Optional[str] = None, + # OpenAI compatibility + base_url: Optional[str] = None, + api_key: Optional[str] = None, + ) -> None: + if model is not None and base_url is not None: + raise ValueError( + "Received both `model` and `base_url` arguments. Please provide only one of them." + " `base_url` is an alias for `model` to make the API compatible with OpenAI's client." + " If using `base_url` for chat completion, the `/chat/completions` suffix path will be appended to the base url." + " When passing a URL as `model`, the client will not append any suffix path to it." + ) + if token is not None and api_key is not None: + raise ValueError( + "Received both `token` and `api_key` arguments. Please provide only one of them." + " `api_key` is an alias for `token` to make the API compatible with OpenAI's client." + " It has the exact same behavior as `token`." + ) + token = token if token is not None else api_key + if isinstance(token, bool): + # Legacy behavior: previously is was possible to pass `token=False` to disable authentication. This is not + # supported anymore as authentication is required. Better to explicitly raise here rather than risking + # sending the locally saved token without the user knowing about it. + if token is False: + raise ValueError( + "Cannot use `token=False` to disable authentication as authentication is required to run Inference." + ) + warnings.warn( + "Using `token=True` to automatically use the locally saved token is deprecated and will be removed in a future release. " + "Please use `token=None` instead (default).", + DeprecationWarning, + ) + token = get_token() + + self.model: Optional[str] = base_url or model + self.token: Optional[str] = token + + self.headers = {**headers} if headers is not None else {} + if bill_to is not None: + if ( + constants.HUGGINGFACE_HEADER_X_BILL_TO in self.headers + and self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO] != bill_to + ): + warnings.warn( + f"Overriding existing '{self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO]}' value in headers with '{bill_to}'.", + UserWarning, + ) + self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO] = bill_to + + if token is not None and not token.startswith("hf_"): + warnings.warn( + "You've provided an external provider's API key, so requests will be billed directly by the provider. " + "The `bill_to` parameter is only applicable for Hugging Face billing and will be ignored.", + UserWarning, + ) + + # Configure provider + self.provider = provider + + self.cookies = cookies + self.timeout = timeout + self.proxies = proxies + + def __repr__(self): + return f"" + + @overload + def _inner_post( # type: ignore[misc] + self, request_parameters: RequestParameters, *, stream: Literal[False] = ... + ) -> bytes: ... + + @overload + def _inner_post( # type: ignore[misc] + self, request_parameters: RequestParameters, *, stream: Literal[True] = ... + ) -> Iterable[bytes]: ... + + @overload + def _inner_post( + self, request_parameters: RequestParameters, *, stream: bool = False + ) -> Union[bytes, Iterable[bytes]]: ... + + def _inner_post( + self, request_parameters: RequestParameters, *, stream: bool = False + ) -> Union[bytes, Iterable[bytes]]: + """Make a request to the inference server.""" + # TODO: this should be handled in provider helpers directly + if request_parameters.task in TASKS_EXPECTING_IMAGES and "Accept" not in request_parameters.headers: + request_parameters.headers["Accept"] = "image/png" + + with _open_as_binary(request_parameters.data) as data_as_binary: + try: + response = get_session().post( + request_parameters.url, + json=request_parameters.json, + data=data_as_binary, + headers=request_parameters.headers, + cookies=self.cookies, + timeout=self.timeout, + stream=stream, + proxies=self.proxies, + ) + except TimeoutError as error: + # Convert any `TimeoutError` to a `InferenceTimeoutError` + raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error # type: ignore + + try: + hf_raise_for_status(response) + return response.iter_lines() if stream else response.content + except HTTPError as error: + if error.response.status_code == 422 and request_parameters.task != "unknown": + msg = str(error.args[0]) + if len(error.response.text) > 0: + msg += f"\n{error.response.text}\n" + error.args = (msg,) + error.args[1:] + raise + + def audio_classification( + self, + audio: ContentT, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + function_to_apply: Optional["AudioClassificationOutputTransform"] = None, + ) -> List[AudioClassificationOutputElement]: + """ + Perform audio classification on the provided audio content. + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The audio content to classify. It can be raw audio bytes, a local audio file, or a URL pointing to an + audio file. + model (`str`, *optional*): + The model to use for audio classification. Can be a model ID hosted on the Hugging Face Hub + or a URL to a deployed Inference Endpoint. If not provided, the default recommended model for + audio classification will be used. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + function_to_apply (`"AudioClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + + Returns: + `List[AudioClassificationOutputElement]`: List of [`AudioClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.audio_classification("audio.flac") + [ + AudioClassificationOutputElement(score=0.4976358711719513, label='hap'), + AudioClassificationOutputElement(score=0.3677836060523987, label='neu'), + ... + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="audio-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={"function_to_apply": function_to_apply, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return AudioClassificationOutputElement.parse_obj_as_list(response) + + def audio_to_audio( + self, + audio: ContentT, + *, + model: Optional[str] = None, + ) -> List[AudioToAudioOutputElement]: + """ + Performs multiple tasks related to audio-to-audio depending on the model (eg: speech enhancement, source separation). + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The audio content for the model. It can be raw audio bytes, a local audio file, or a URL pointing to an + audio file. + model (`str`, *optional*): + The model can be any model which takes an audio file and returns another audio file. Can be a model ID hosted on the Hugging Face Hub + or a URL to a deployed Inference Endpoint. If not provided, the default recommended model for + audio_to_audio will be used. + + Returns: + `List[AudioToAudioOutputElement]`: A list of [`AudioToAudioOutputElement`] items containing audios label, content-type, and audio content in blob. + + Raises: + `InferenceTimeoutError`: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> audio_output = client.audio_to_audio("audio.flac") + >>> for i, item in enumerate(audio_output): + >>> with open(f"output_{i}.flac", "wb") as f: + f.write(item.blob) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="audio-to-audio", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + audio_output = AudioToAudioOutputElement.parse_obj_as_list(response) + for item in audio_output: + item.blob = base64.b64decode(item.blob) + return audio_output + + def automatic_speech_recognition( + self, + audio: ContentT, + *, + model: Optional[str] = None, + extra_body: Optional[Dict] = None, + ) -> AutomaticSpeechRecognitionOutput: + """ + Perform automatic speech recognition (ASR or audio-to-text) on the given audio content. + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The content to transcribe. It can be raw audio bytes, local audio file, or a URL to an audio file. + model (`str`, *optional*): + The model to use for ASR. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for ASR will be used. + extra_body (`Dict`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + [`AutomaticSpeechRecognitionOutput`]: An item containing the transcribed text and optionally the timestamp chunks. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.automatic_speech_recognition("hello_world.flac").text + "hello world" + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="automatic-speech-recognition", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={**(extra_body or {})}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return AutomaticSpeechRecognitionOutput.parse_obj_as_instance(response) + + @overload + def chat_completion( # type: ignore + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: Literal[False] = False, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> ChatCompletionOutput: ... + + @overload + def chat_completion( # type: ignore + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: Literal[True] = True, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> Iterable[ChatCompletionStreamOutput]: ... + + @overload + def chat_completion( + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: bool = False, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> Union[ChatCompletionOutput, Iterable[ChatCompletionStreamOutput]]: ... + + def chat_completion( + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: bool = False, + # Parameters from ChatCompletionInput (handled manually) + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> Union[ChatCompletionOutput, Iterable[ChatCompletionStreamOutput]]: + """ + A method for completing conversations using a specified language model. + + + + The `client.chat_completion` method is aliased as `client.chat.completions.create` for compatibility with OpenAI's client. + Inputs and outputs are strictly the same and using either syntax will yield the same results. + Check out the [Inference guide](https://huggingface.co/docs/huggingface_hub/guides/inference#openai-compatibility) + for more details about OpenAI's compatibility. + + + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + messages (List of [`ChatCompletionInputMessage`]): + Conversation history consisting of roles and content pairs. + model (`str`, *optional*): + The model to use for chat-completion. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for chat-based text-generation will be used. + See https://huggingface.co/tasks/text-generation for more details. + If `model` is a model ID, it is passed to the server as the `model` parameter. If you want to define a + custom URL while setting `model` in the request payload, you must set `base_url` when initializing [`InferenceClient`]. + frequency_penalty (`float`, *optional*): + Penalizes new tokens based on their existing frequency + in the text so far. Range: [-2.0, 2.0]. Defaults to 0.0. + logit_bias (`List[float]`, *optional*): + Adjusts the likelihood of specific tokens appearing in the generated output. + logprobs (`bool`, *optional*): + Whether to return log probabilities of the output tokens or not. If true, returns the log + probabilities of each output token returned in the content of message. + max_tokens (`int`, *optional*): + Maximum number of tokens allowed in the response. Defaults to 100. + n (`int`, *optional*): + The number of completions to generate for each prompt. + presence_penalty (`float`, *optional*): + Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the + text so far, increasing the model's likelihood to talk about new topics. + response_format ([`ChatCompletionInputGrammarType`], *optional*): + Grammar constraints. Can be either a JSONSchema or a regex. + seed (Optional[`int`], *optional*): + Seed for reproducible control flow. Defaults to None. + stop (`List[str]`, *optional*): + Up to four strings which trigger the end of the response. + Defaults to None. + stream (`bool`, *optional*): + Enable realtime streaming of responses. Defaults to False. + stream_options ([`ChatCompletionInputStreamOptions`], *optional*): + Options for streaming completions. + temperature (`float`, *optional*): + Controls randomness of the generations. Lower values ensure + less random completions. Range: [0, 2]. Defaults to 1.0. + top_logprobs (`int`, *optional*): + An integer between 0 and 5 specifying the number of most likely tokens to return at each token + position, each with an associated log probability. logprobs must be set to true if this parameter is + used. + top_p (`float`, *optional*): + Fraction of the most likely next words to sample from. + Must be between 0 and 1. Defaults to 1.0. + tool_choice ([`ChatCompletionInputToolChoiceClass`] or [`ChatCompletionInputToolChoiceEnum`], *optional*): + The tool to use for the completion. Defaults to "auto". + tool_prompt (`str`, *optional*): + A prompt to be appended before the tools. + tools (List of [`ChatCompletionInputTool`], *optional*): + A list of tools the model may call. Currently, only functions are supported as a tool. Use this to + provide a list of functions the model may generate JSON inputs for. + extra_body (`Dict`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + [`ChatCompletionOutput`] or Iterable of [`ChatCompletionStreamOutput`]: + Generated text returned from the server: + - if `stream=False`, the generated text is returned as a [`ChatCompletionOutput`] (default). + - if `stream=True`, the generated text is returned token by token as a sequence of [`ChatCompletionStreamOutput`]. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + + ```py + >>> from huggingface_hub import InferenceClient + >>> messages = [{"role": "user", "content": "What is the capital of France?"}] + >>> client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct") + >>> client.chat_completion(messages, max_tokens=100) + ChatCompletionOutput( + choices=[ + ChatCompletionOutputComplete( + finish_reason='eos_token', + index=0, + message=ChatCompletionOutputMessage( + role='assistant', + content='The capital of France is Paris.', + name=None, + tool_calls=None + ), + logprobs=None + ) + ], + created=1719907176, + id='', + model='meta-llama/Meta-Llama-3-8B-Instruct', + object='text_completion', + system_fingerprint='2.0.4-sha-f426a33', + usage=ChatCompletionOutputUsage( + completion_tokens=8, + prompt_tokens=17, + total_tokens=25 + ) + ) + ``` + + Example using streaming: + ```py + >>> from huggingface_hub import InferenceClient + >>> messages = [{"role": "user", "content": "What is the capital of France?"}] + >>> client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct") + >>> for token in client.chat_completion(messages, max_tokens=10, stream=True): + ... print(token) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content='The', role='assistant'), index=0, finish_reason=None)], created=1710498504) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content=' capital', role='assistant'), index=0, finish_reason=None)], created=1710498504) + (...) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content=' may', role='assistant'), index=0, finish_reason=None)], created=1710498504) + ``` + + Example using OpenAI's syntax: + ```py + # instead of `from openai import OpenAI` + from huggingface_hub import InferenceClient + + # instead of `client = OpenAI(...)` + client = InferenceClient( + base_url=..., + api_key=..., + ) + + output = client.chat.completions.create( + model="meta-llama/Meta-Llama-3-8B-Instruct", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Count to 10"}, + ], + stream=True, + max_tokens=1024, + ) + + for chunk in output: + print(chunk.choices[0].delta.content) + ``` + + Example using a third-party provider directly with extra (provider-specific) parameters. Usage will be billed on your Together AI account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="together", # Use Together AI provider + ... api_key="", # Pass your Together API key directly + ... ) + >>> client.chat_completion( + ... model="meta-llama/Meta-Llama-3-8B-Instruct", + ... messages=[{"role": "user", "content": "What is the capital of France?"}], + ... extra_body={"safety_model": "Meta-Llama/Llama-Guard-7b"}, + ... ) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="sambanova", # Use Sambanova provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> client.chat_completion( + ... model="meta-llama/Meta-Llama-3-8B-Instruct", + ... messages=[{"role": "user", "content": "What is the capital of France?"}], + ... ) + ``` + + Example using Image + Text as input: + ```py + >>> from huggingface_hub import InferenceClient + + # provide a remote URL + >>> image_url ="https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" + # or a base64-encoded image + >>> image_path = "/path/to/image.jpeg" + >>> with open(image_path, "rb") as f: + ... base64_image = base64.b64encode(f.read()).decode("utf-8") + >>> image_url = f"data:image/jpeg;base64,{base64_image}" + + >>> client = InferenceClient("meta-llama/Llama-3.2-11B-Vision-Instruct") + >>> output = client.chat.completions.create( + ... messages=[ + ... { + ... "role": "user", + ... "content": [ + ... { + ... "type": "image_url", + ... "image_url": {"url": image_url}, + ... }, + ... { + ... "type": "text", + ... "text": "Describe this image in one sentence.", + ... }, + ... ], + ... }, + ... ], + ... ) + >>> output + The image depicts the iconic Statue of Liberty situated in New York Harbor, New York, on a clear day. + ``` + + Example using tools: + ```py + >>> client = InferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> messages = [ + ... { + ... "role": "system", + ... "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.", + ... }, + ... { + ... "role": "user", + ... "content": "What's the weather like the next 3 days in San Francisco, CA?", + ... }, + ... ] + >>> tools = [ + ... { + ... "type": "function", + ... "function": { + ... "name": "get_current_weather", + ... "description": "Get the current weather", + ... "parameters": { + ... "type": "object", + ... "properties": { + ... "location": { + ... "type": "string", + ... "description": "The city and state, e.g. San Francisco, CA", + ... }, + ... "format": { + ... "type": "string", + ... "enum": ["celsius", "fahrenheit"], + ... "description": "The temperature unit to use. Infer this from the users location.", + ... }, + ... }, + ... "required": ["location", "format"], + ... }, + ... }, + ... }, + ... { + ... "type": "function", + ... "function": { + ... "name": "get_n_day_weather_forecast", + ... "description": "Get an N-day weather forecast", + ... "parameters": { + ... "type": "object", + ... "properties": { + ... "location": { + ... "type": "string", + ... "description": "The city and state, e.g. San Francisco, CA", + ... }, + ... "format": { + ... "type": "string", + ... "enum": ["celsius", "fahrenheit"], + ... "description": "The temperature unit to use. Infer this from the users location.", + ... }, + ... "num_days": { + ... "type": "integer", + ... "description": "The number of days to forecast", + ... }, + ... }, + ... "required": ["location", "format", "num_days"], + ... }, + ... }, + ... }, + ... ] + + >>> response = client.chat_completion( + ... model="meta-llama/Meta-Llama-3-70B-Instruct", + ... messages=messages, + ... tools=tools, + ... tool_choice="auto", + ... max_tokens=500, + ... ) + >>> response.choices[0].message.tool_calls[0].function + ChatCompletionOutputFunctionDefinition( + arguments={ + 'location': 'San Francisco, CA', + 'format': 'fahrenheit', + 'num_days': 3 + }, + name='get_n_day_weather_forecast', + description=None + ) + ``` + + Example using response_format: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> messages = [ + ... { + ... "role": "user", + ... "content": "I saw a puppy a cat and a raccoon during my bike ride in the park. What did I saw and when?", + ... }, + ... ] + >>> response_format = { + ... "type": "json", + ... "value": { + ... "properties": { + ... "location": {"type": "string"}, + ... "activity": {"type": "string"}, + ... "animals_seen": {"type": "integer", "minimum": 1, "maximum": 5}, + ... "animals": {"type": "array", "items": {"type": "string"}}, + ... }, + ... "required": ["location", "activity", "animals_seen", "animals"], + ... }, + ... } + >>> response = client.chat_completion( + ... messages=messages, + ... response_format=response_format, + ... max_tokens=500, + ... ) + >>> response.choices[0].message.content + '{\n\n"activity": "bike ride",\n"animals": ["puppy", "cat", "raccoon"],\n"animals_seen": 3,\n"location": "park"}' + ``` + """ + # Since `chat_completion(..., model=xxx)` is also a payload parameter for the server, we need to handle 'model' differently. + # `self.model` takes precedence over 'model' argument for building URL. + # `model` takes precedence for payload value. + model_id_or_url = self.model or model + payload_model = model or self.model + + # Get the provider helper + provider_helper = get_provider_helper( + self.provider, + task="conversational", + model=model_id_or_url + if model_id_or_url is not None and model_id_or_url.startswith(("http://", "https://")) + else payload_model, + ) + + # Prepare the payload + parameters = { + "model": payload_model, + "frequency_penalty": frequency_penalty, + "logit_bias": logit_bias, + "logprobs": logprobs, + "max_tokens": max_tokens, + "n": n, + "presence_penalty": presence_penalty, + "response_format": response_format, + "seed": seed, + "stop": stop, + "temperature": temperature, + "tool_choice": tool_choice, + "tool_prompt": tool_prompt, + "tools": tools, + "top_logprobs": top_logprobs, + "top_p": top_p, + "stream": stream, + "stream_options": stream_options, + **(extra_body or {}), + } + request_parameters = provider_helper.prepare_request( + inputs=messages, + parameters=parameters, + headers=self.headers, + model=model_id_or_url, + api_key=self.token, + ) + data = self._inner_post(request_parameters, stream=stream) + + if stream: + return _stream_chat_completion_response(data) # type: ignore[arg-type] + + return ChatCompletionOutput.parse_obj_as_instance(data) # type: ignore[arg-type] + + def document_question_answering( + self, + image: ContentT, + question: str, + *, + model: Optional[str] = None, + doc_stride: Optional[int] = None, + handle_impossible_answer: Optional[bool] = None, + lang: Optional[str] = None, + max_answer_len: Optional[int] = None, + max_question_len: Optional[int] = None, + max_seq_len: Optional[int] = None, + top_k: Optional[int] = None, + word_boxes: Optional[List[Union[List[float], str]]] = None, + ) -> List[DocumentQuestionAnsweringOutputElement]: + """ + Answer questions on document images. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for the context. It can be raw bytes, an image file, or a URL to an online image. + question (`str`): + Question to be answered. + model (`str`, *optional*): + The model to use for the document question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended document question answering model will be used. + Defaults to None. + doc_stride (`int`, *optional*): + If the words in the document are too long to fit with the question for the model, it will be split in + several chunks with some overlap. This argument controls the size of that overlap. + handle_impossible_answer (`bool`, *optional*): + Whether to accept impossible as an answer + lang (`str`, *optional*): + Language to use while running OCR. Defaults to english. + max_answer_len (`int`, *optional*): + The maximum length of predicted answers (e.g., only answers with a shorter length are considered). + max_question_len (`int`, *optional*): + The maximum length of the question after tokenization. It will be truncated if needed. + max_seq_len (`int`, *optional*): + The maximum length of the total sentence (context + question) in tokens of each chunk passed to the + model. The context will be split in several chunks (using doc_stride as overlap) if needed. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Can return less than top_k + answers if there are not enough options available within the context. + word_boxes (`List[Union[List[float], str`, *optional*): + A list of words and bounding boxes (normalized 0->1000). If provided, the inference will skip the OCR + step and use the provided bounding boxes instead. + Returns: + `List[DocumentQuestionAnsweringOutputElement]`: a list of [`DocumentQuestionAnsweringOutputElement`] items containing the predicted label, associated probability, word ids, and page number. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.document_question_answering(image="https://huggingface.co/spaces/impira/docquery/resolve/2359223c1837a7587402bda0f2643382a6eefeab/invoice.png", question="What is the invoice number?") + [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16)] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="document-question-answering", model=model_id) + inputs: Dict[str, Any] = {"question": question, "image": _b64_encode(image)} + request_parameters = provider_helper.prepare_request( + inputs=inputs, + parameters={ + "doc_stride": doc_stride, + "handle_impossible_answer": handle_impossible_answer, + "lang": lang, + "max_answer_len": max_answer_len, + "max_question_len": max_question_len, + "max_seq_len": max_seq_len, + "top_k": top_k, + "word_boxes": word_boxes, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return DocumentQuestionAnsweringOutputElement.parse_obj_as_list(response) + + def feature_extraction( + self, + text: str, + *, + normalize: Optional[bool] = None, + prompt_name: Optional[str] = None, + truncate: Optional[bool] = None, + truncation_direction: Optional[Literal["Left", "Right"]] = None, + model: Optional[str] = None, + ) -> "np.ndarray": + """ + Generate embeddings for a given text. + + Args: + text (`str`): + The text to embed. + model (`str`, *optional*): + The model to use for the feature extraction task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended feature extraction model will be used. + Defaults to None. + normalize (`bool`, *optional*): + Whether to normalize the embeddings or not. + Only available on server powered by Text-Embedding-Inference. + prompt_name (`str`, *optional*): + The name of the prompt that should be used by for encoding. If not set, no prompt will be applied. + Must be a key in the `Sentence Transformers` configuration `prompts` dictionary. + For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",...}, + then the sentence "What is the capital of France?" will be encoded as "query: What is the capital of France?" + because the prompt text will be prepended before any text to encode. + truncate (`bool`, *optional*): + Whether to truncate the embeddings or not. + Only available on server powered by Text-Embedding-Inference. + truncation_direction (`Literal["Left", "Right"]`, *optional*): + Which side of the input should be truncated when `truncate=True` is passed. + + Returns: + `np.ndarray`: The embedding representing the input text as a float32 numpy array. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.feature_extraction("Hi, who are you?") + array([[ 2.424802 , 2.93384 , 1.1750331 , ..., 1.240499, -0.13776633, -0.7889173 ], + [-0.42943227, -0.6364878 , -1.693462 , ..., 0.41978157, -2.4336355 , 0.6162071 ], + ..., + [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="feature-extraction", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "normalize": normalize, + "prompt_name": prompt_name, + "truncate": truncate, + "truncation_direction": truncation_direction, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + np = _import_numpy() + return np.array(provider_helper.get_response(response), dtype="float32") + + def fill_mask( + self, + text: str, + *, + model: Optional[str] = None, + targets: Optional[List[str]] = None, + top_k: Optional[int] = None, + ) -> List[FillMaskOutputElement]: + """ + Fill in a hole with a missing word (token to be precise). + + Args: + text (`str`): + a string to be filled from, must contain the [MASK] token (check model card for exact name of the mask). + model (`str`, *optional*): + The model to use for the fill mask task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended fill mask model will be used. + targets (`List[str`, *optional*): + When passed, the model will limit the scores to the passed targets instead of looking up in the whole + vocabulary. If the provided targets are not in the model vocab, they will be tokenized and the first + resulting token will be used (with a warning, and that might be slower). + top_k (`int`, *optional*): + When passed, overrides the number of predictions to return. + Returns: + `List[FillMaskOutputElement]`: a list of [`FillMaskOutputElement`] items containing the predicted label, associated + probability, token reference, and completed text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.fill_mask("The goal of life is .") + [ + FillMaskOutputElement(score=0.06897063553333282, token=11098, token_str=' happiness', sequence='The goal of life is happiness.'), + FillMaskOutputElement(score=0.06554922461509705, token=45075, token_str=' immortality', sequence='The goal of life is immortality.') + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="fill-mask", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={"targets": targets, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return FillMaskOutputElement.parse_obj_as_list(response) + + def image_classification( + self, + image: ContentT, + *, + model: Optional[str] = None, + function_to_apply: Optional["ImageClassificationOutputTransform"] = None, + top_k: Optional[int] = None, + ) -> List[ImageClassificationOutputElement]: + """ + Perform image classification on the given image using the specified model. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to classify. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for image classification. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for image classification will be used. + function_to_apply (`"ImageClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + Returns: + `List[ImageClassificationOutputElement]`: a list of [`ImageClassificationOutputElement`] items containing the predicted label and associated probability. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.image_classification("https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg") + [ImageClassificationOutputElement(label='Blenheim spaniel', score=0.9779096841812134), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"function_to_apply": function_to_apply, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return ImageClassificationOutputElement.parse_obj_as_list(response) + + def image_segmentation( + self, + image: ContentT, + *, + model: Optional[str] = None, + mask_threshold: Optional[float] = None, + overlap_mask_area_threshold: Optional[float] = None, + subtask: Optional["ImageSegmentationSubtask"] = None, + threshold: Optional[float] = None, + ) -> List[ImageSegmentationOutputElement]: + """ + Perform image segmentation on the given image using the specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to segment. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for image segmentation. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for image segmentation will be used. + mask_threshold (`float`, *optional*): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*): + Mask overlap threshold to eliminate small, disconnected segments. + subtask (`"ImageSegmentationSubtask"`, *optional*): + Segmentation task to be performed, depending on model capabilities. + threshold (`float`, *optional*): + Probability threshold to filter out predicted masks. + Returns: + `List[ImageSegmentationOutputElement]`: A list of [`ImageSegmentationOutputElement`] items containing the segmented masks and associated attributes. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.image_segmentation("cat.jpg") + [ImageSegmentationOutputElement(score=0.989008, label='LABEL_184', mask=), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-segmentation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "mask_threshold": mask_threshold, + "overlap_mask_area_threshold": overlap_mask_area_threshold, + "subtask": subtask, + "threshold": threshold, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + output = ImageSegmentationOutputElement.parse_obj_as_list(response) + for item in output: + item.mask = _b64_to_image(item.mask) # type: ignore [assignment] + return output + + def image_to_image( + self, + image: ContentT, + prompt: Optional[str] = None, + *, + negative_prompt: Optional[str] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + model: Optional[str] = None, + target_size: Optional[ImageToImageTargetSize] = None, + **kwargs, + ) -> "Image": + """ + Perform image-to-image translation using a specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for translation. It can be raw bytes, an image file, or a URL to an online image. + prompt (`str`, *optional*): + The text prompt to guide the image generation. + negative_prompt (`str`, *optional*): + One prompt to guide what NOT to include in image generation. + num_inference_steps (`int`, *optional*): + For diffusion models. The number of denoising steps. More denoising steps usually lead to a higher + quality image at the expense of slower inference. + guidance_scale (`float`, *optional*): + For diffusion models. A higher guidance scale value encourages the model to generate images closely + linked to the text prompt at the expense of lower image quality. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + target_size (`ImageToImageTargetSize`, *optional*): + The size in pixel of the output image. + + Returns: + `Image`: The translated image. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> image = client.image_to_image("cat.jpg", prompt="turn the cat into a tiger") + >>> image.save("tiger.jpg") + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-to-image", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "prompt": prompt, + "negative_prompt": negative_prompt, + "target_size": target_size, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + **kwargs, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return _bytes_to_image(response) + + def image_to_text(self, image: ContentT, *, model: Optional[str] = None) -> ImageToTextOutput: + """ + Takes an input image and return text. + + Models can have very different outputs depending on your use case (image captioning, optical character recognition + (OCR), Pix2Struct, etc). Please have a look to the model card to learn more about a model's specificities. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image to caption. It can be raw bytes, an image file, or a URL to an online image.. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + [`ImageToTextOutput`]: The generated text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.image_to_text("cat.jpg") + 'a cat standing in a grassy field ' + >>> client.image_to_text("https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg") + 'a dog laying on the grass next to a flower pot ' + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-to-text", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + output = ImageToTextOutput.parse_obj(response) + return output[0] if isinstance(output, list) else output + + def object_detection( + self, image: ContentT, *, model: Optional[str] = None, threshold: Optional[float] = None + ) -> List[ObjectDetectionOutputElement]: + """ + Perform object detection on the given image using the specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to detect objects on. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for object detection. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for object detection (DETR) will be used. + threshold (`float`, *optional*): + The probability necessary to make a prediction. + Returns: + `List[ObjectDetectionOutputElement]`: A list of [`ObjectDetectionOutputElement`] items containing the bounding boxes and associated attributes. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + `ValueError`: + If the request output is not a List. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.object_detection("people.jpg") + [ObjectDetectionOutputElement(score=0.9486683011054993, label='person', box=ObjectDetectionBoundingBox(xmin=59, ymin=39, xmax=420, ymax=510)), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="object-detection", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"threshold": threshold}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return ObjectDetectionOutputElement.parse_obj_as_list(response) + + def question_answering( + self, + question: str, + context: str, + *, + model: Optional[str] = None, + align_to_words: Optional[bool] = None, + doc_stride: Optional[int] = None, + handle_impossible_answer: Optional[bool] = None, + max_answer_len: Optional[int] = None, + max_question_len: Optional[int] = None, + max_seq_len: Optional[int] = None, + top_k: Optional[int] = None, + ) -> Union[QuestionAnsweringOutputElement, List[QuestionAnsweringOutputElement]]: + """ + Retrieve the answer to a question from a given text. + + Args: + question (`str`): + Question to be answered. + context (`str`): + The context of the question. + model (`str`): + The model to use for the question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. + align_to_words (`bool`, *optional*): + Attempts to align the answer to real words. Improves quality on space separated languages. Might hurt + on non-space-separated languages (like Japanese or Chinese) + doc_stride (`int`, *optional*): + If the context is too long to fit with the question for the model, it will be split in several chunks + with some overlap. This argument controls the size of that overlap. + handle_impossible_answer (`bool`, *optional*): + Whether to accept impossible as an answer. + max_answer_len (`int`, *optional*): + The maximum length of predicted answers (e.g., only answers with a shorter length are considered). + max_question_len (`int`, *optional*): + The maximum length of the question after tokenization. It will be truncated if needed. + max_seq_len (`int`, *optional*): + The maximum length of the total sentence (context + question) in tokens of each chunk passed to the + model. The context will be split in several chunks (using docStride as overlap) if needed. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Note that we return less than + topk answers if there are not enough options available within the context. + + Returns: + Union[`QuestionAnsweringOutputElement`, List[`QuestionAnsweringOutputElement`]]: + When top_k is 1 or not provided, it returns a single `QuestionAnsweringOutputElement`. + When top_k is greater than 1, it returns a list of `QuestionAnsweringOutputElement`. + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.question_answering(question="What's my name?", context="My name is Clara and I live in Berkeley.") + QuestionAnsweringOutputElement(answer='Clara', end=16, score=0.9326565265655518, start=11) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={ + "align_to_words": align_to_words, + "doc_stride": doc_stride, + "handle_impossible_answer": handle_impossible_answer, + "max_answer_len": max_answer_len, + "max_question_len": max_question_len, + "max_seq_len": max_seq_len, + "top_k": top_k, + }, + extra_payload={"question": question, "context": context}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + # Parse the response as a single `QuestionAnsweringOutputElement` when top_k is 1 or not provided, or a list of `QuestionAnsweringOutputElement` to ensure backward compatibility. + output = QuestionAnsweringOutputElement.parse_obj(response) + return output + + def sentence_similarity( + self, sentence: str, other_sentences: List[str], *, model: Optional[str] = None + ) -> List[float]: + """ + Compute the semantic similarity between a sentence and a list of other sentences by comparing their embeddings. + + Args: + sentence (`str`): + The main sentence to compare to others. + other_sentences (`List[str]`): + The list of sentences to compare to. + model (`str`, *optional*): + The model to use for the sentence similarity task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended sentence similarity model will be used. + Defaults to None. + + Returns: + `List[float]`: The embedding representing the input text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.sentence_similarity( + ... "Machine learning is so easy.", + ... other_sentences=[ + ... "Deep learning is so straightforward.", + ... "This is so difficult, like rocket science.", + ... "I can't believe how much I struggled with this.", + ... ], + ... ) + [0.7785726189613342, 0.45876261591911316, 0.2906220555305481] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="sentence-similarity", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs={"source_sentence": sentence, "sentences": other_sentences}, + parameters={}, + extra_payload={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return _bytes_to_list(response) + + def summarization( + self, + text: str, + *, + model: Optional[str] = None, + clean_up_tokenization_spaces: Optional[bool] = None, + generate_parameters: Optional[Dict[str, Any]] = None, + truncation: Optional["SummarizationTruncationStrategy"] = None, + ) -> SummarizationOutput: + """ + Generate a summary of a given text using a specified model. + + Args: + text (`str`): + The input text to summarize. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for summarization will be used. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether to clean up the potential extra spaces in the text output. + generate_parameters (`Dict[str, Any]`, *optional*): + Additional parametrization of the text generation algorithm. + truncation (`"SummarizationTruncationStrategy"`, *optional*): + The truncation strategy to use. + Returns: + [`SummarizationOutput`]: The generated summary text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.summarization("The Eiffel tower...") + SummarizationOutput(generated_text="The Eiffel tower is one of the most famous landmarks in the world....") + ``` + """ + parameters = { + "clean_up_tokenization_spaces": clean_up_tokenization_spaces, + "generate_parameters": generate_parameters, + "truncation": truncation, + } + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="summarization", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters=parameters, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return SummarizationOutput.parse_obj_as_list(response)[0] + + def table_question_answering( + self, + table: Dict[str, Any], + query: str, + *, + model: Optional[str] = None, + padding: Optional["Padding"] = None, + sequential: Optional[bool] = None, + truncation: Optional[bool] = None, + ) -> TableQuestionAnsweringOutputElement: + """ + Retrieve the answer to a question from information given in a table. + + Args: + table (`str`): + A table of data represented as a dict of lists where entries are headers and the lists are all the + values, all lists must have the same size. + query (`str`): + The query in plain text that you want to ask the table. + model (`str`): + The model to use for the table-question-answering task. Can be a model ID hosted on the Hugging Face + Hub or a URL to a deployed Inference Endpoint. + padding (`"Padding"`, *optional*): + Activates and controls padding. + sequential (`bool`, *optional*): + Whether to do inference sequentially or as a batch. Batching is faster, but models like SQA require the + inference to be done sequentially to extract relations within sequences, given their conversational + nature. + truncation (`bool`, *optional*): + Activates and controls truncation. + + Returns: + [`TableQuestionAnsweringOutputElement`]: a table question answering output containing the answer, coordinates, cells and the aggregator used. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> query = "How many stars does the transformers repository have?" + >>> table = {"Repository": ["Transformers", "Datasets", "Tokenizers"], "Stars": ["36542", "4512", "3934"]} + >>> client.table_question_answering(table, query, model="google/tapas-base-finetuned-wtq") + TableQuestionAnsweringOutputElement(answer='36542', coordinates=[[0, 1]], cells=['36542'], aggregator='AVERAGE') + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="table-question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={"model": model, "padding": padding, "sequential": sequential, "truncation": truncation}, + extra_payload={"query": query, "table": table}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return TableQuestionAnsweringOutputElement.parse_obj_as_instance(response) + + def tabular_classification(self, table: Dict[str, Any], *, model: Optional[str] = None) -> List[str]: + """ + Classifying a target category (a group) based on a set of attributes. + + Args: + table (`Dict[str, Any]`): + Set of attributes to classify. + model (`str`, *optional*): + The model to use for the tabular classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended tabular classification model will be used. + Defaults to None. + + Returns: + `List`: a list of labels, one per row in the initial table. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> table = { + ... "fixed_acidity": ["7.4", "7.8", "10.3"], + ... "volatile_acidity": ["0.7", "0.88", "0.32"], + ... "citric_acid": ["0", "0", "0.45"], + ... "residual_sugar": ["1.9", "2.6", "6.4"], + ... "chlorides": ["0.076", "0.098", "0.073"], + ... "free_sulfur_dioxide": ["11", "25", "5"], + ... "total_sulfur_dioxide": ["34", "67", "13"], + ... "density": ["0.9978", "0.9968", "0.9976"], + ... "pH": ["3.51", "3.2", "3.23"], + ... "sulphates": ["0.56", "0.68", "0.82"], + ... "alcohol": ["9.4", "9.8", "12.6"], + ... } + >>> client.tabular_classification(table=table, model="julien-c/wine-quality") + ["5", "5", "5"] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="tabular-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + extra_payload={"table": table}, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return _bytes_to_list(response) + + def tabular_regression(self, table: Dict[str, Any], *, model: Optional[str] = None) -> List[float]: + """ + Predicting a numerical target value given a set of attributes/features in a table. + + Args: + table (`Dict[str, Any]`): + Set of attributes stored in a table. The attributes used to predict the target can be both numerical and categorical. + model (`str`, *optional*): + The model to use for the tabular regression task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended tabular regression model will be used. + Defaults to None. + + Returns: + `List`: a list of predicted numerical target values. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> table = { + ... "Height": ["11.52", "12.48", "12.3778"], + ... "Length1": ["23.2", "24", "23.9"], + ... "Length2": ["25.4", "26.3", "26.5"], + ... "Length3": ["30", "31.2", "31.1"], + ... "Species": ["Bream", "Bream", "Bream"], + ... "Width": ["4.02", "4.3056", "4.6961"], + ... } + >>> client.tabular_regression(table, model="scikit-learn/Fish-Weight") + [110, 120, 130] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="tabular-regression", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={}, + extra_payload={"table": table}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return _bytes_to_list(response) + + def text_classification( + self, + text: str, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + function_to_apply: Optional["TextClassificationOutputTransform"] = None, + ) -> List[TextClassificationOutputElement]: + """ + Perform text classification (e.g. sentiment-analysis) on the given text. + + Args: + text (`str`): + A string to be classified. + model (`str`, *optional*): + The model to use for the text classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended text classification model will be used. + Defaults to None. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + function_to_apply (`"TextClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + + Returns: + `List[TextClassificationOutputElement]`: a list of [`TextClassificationOutputElement`] items containing the predicted label and associated probability. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.text_classification("I like you") + [ + TextClassificationOutputElement(label='POSITIVE', score=0.9998695850372314), + TextClassificationOutputElement(label='NEGATIVE', score=0.0001304351753788069), + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "function_to_apply": function_to_apply, + "top_k": top_k, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return TextClassificationOutputElement.parse_obj_as_list(response)[0] # type: ignore [return-value] + + @overload + def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[False] = ..., + stream: Literal[False] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> str: ... + + @overload + def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: Literal[False] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> TextGenerationOutput: ... + + @overload + def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[False] = ..., + stream: Literal[True] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Iterable[str]: ... + + @overload + def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: Literal[True] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Iterable[TextGenerationStreamOutput]: ... + + @overload + def text_generation( + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: bool = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Union[TextGenerationOutput, Iterable[TextGenerationStreamOutput]]: ... + + def text_generation( + self, + prompt: str, + *, + details: bool = False, + stream: bool = False, + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Union[str, TextGenerationOutput, Iterable[str], Iterable[TextGenerationStreamOutput]]: + """ + Given a prompt, generate the following text. + + + + If you want to generate a response from chat messages, you should use the [`InferenceClient.chat_completion`] method. + It accepts a list of messages instead of a single text prompt and handles the chat templating for you. + + + + Args: + prompt (`str`): + Input text. + details (`bool`, *optional*): + By default, text_generation returns a string. Pass `details=True` if you want a detailed output (tokens, + probabilities, seed, finish reason, etc.). Only available for models running on with the + `text-generation-inference` backend. + stream (`bool`, *optional*): + By default, text_generation returns the full generated text. Pass `stream=True` if you want a stream of + tokens to be returned. Only available for models running on with the `text-generation-inference` + backend. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + adapter_id (`str`, *optional*): + Lora adapter id. + best_of (`int`, *optional*): + Generate best_of sequences and return the one if the highest token logprobs. + decoder_input_details (`bool`, *optional*): + Return the decoder input token logprobs and ids. You must set `details=True` as well for it to be taken + into account. Defaults to `False`. + do_sample (`bool`, *optional*): + Activate logits sampling + frequency_penalty (`float`, *optional*): + Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in + the text so far, decreasing the model's likelihood to repeat the same line verbatim. + grammar ([`TextGenerationInputGrammarType`], *optional*): + Grammar constraints. Can be either a JSONSchema or a regex. + max_new_tokens (`int`, *optional*): + Maximum number of generated tokens. Defaults to 100. + repetition_penalty (`float`, *optional*): + The parameter for repetition penalty. 1.0 means no penalty. See [this + paper](https://arxiv.org/pdf/1909.05858.pdf) for more details. + return_full_text (`bool`, *optional*): + Whether to prepend the prompt to the generated text + seed (`int`, *optional*): + Random sampling seed + stop (`List[str]`, *optional*): + Stop generating tokens if a member of `stop` is generated. + stop_sequences (`List[str]`, *optional*): + Deprecated argument. Use `stop` instead. + temperature (`float`, *optional*): + The value used to module the logits distribution. + top_n_tokens (`int`, *optional*): + Return information about the `top_n_tokens` most likely tokens at each generation step, instead of + just the sampled token. + top_k (`int`, *optional`): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + top_p (`float`, *optional`): + If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or + higher are kept for generation. + truncate (`int`, *optional`): + Truncate inputs tokens to the given size. + typical_p (`float`, *optional`): + Typical Decoding mass + See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information + watermark (`bool`, *optional`): + Watermarking with [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226) + + Returns: + `Union[str, TextGenerationOutput, Iterable[str], Iterable[TextGenerationStreamOutput]]`: + Generated text returned from the server: + - if `stream=False` and `details=False`, the generated text is returned as a `str` (default) + - if `stream=True` and `details=False`, the generated text is returned token by token as a `Iterable[str]` + - if `stream=False` and `details=True`, the generated text is returned with more details as a [`~huggingface_hub.TextGenerationOutput`] + - if `details=True` and `stream=True`, the generated text is returned token by token as a iterable of [`~huggingface_hub.TextGenerationStreamOutput`] + + Raises: + `ValidationError`: + If input values are not valid. No HTTP call is made to the server. + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + + # Case 1: generate text + >>> client.text_generation("The huggingface_hub library is ", max_new_tokens=12) + '100% open source and built to be easy to use.' + + # Case 2: iterate over the generated tokens. Useful for large generation. + >>> for token in client.text_generation("The huggingface_hub library is ", max_new_tokens=12, stream=True): + ... print(token) + 100 + % + open + source + and + built + to + be + easy + to + use + . + + # Case 3: get more details about the generation process. + >>> client.text_generation("The huggingface_hub library is ", max_new_tokens=12, details=True) + TextGenerationOutput( + generated_text='100% open source and built to be easy to use.', + details=TextGenerationDetails( + finish_reason='length', + generated_tokens=12, + seed=None, + prefill=[ + TextGenerationPrefillOutputToken(id=487, text='The', logprob=None), + TextGenerationPrefillOutputToken(id=53789, text=' hugging', logprob=-13.171875), + (...) + TextGenerationPrefillOutputToken(id=204, text=' ', logprob=-7.0390625) + ], + tokens=[ + TokenElement(id=1425, text='100', logprob=-1.0175781, special=False), + TokenElement(id=16, text='%', logprob=-0.0463562, special=False), + (...) + TokenElement(id=25, text='.', logprob=-0.5703125, special=False) + ], + best_of_sequences=None + ) + ) + + # Case 4: iterate over the generated tokens with more details. + # Last object is more complete, containing the full generated text and the finish reason. + >>> for details in client.text_generation("The huggingface_hub library is ", max_new_tokens=12, details=True, stream=True): + ... print(details) + ... + TextGenerationStreamOutput(token=TokenElement(id=1425, text='100', logprob=-1.0175781, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=16, text='%', logprob=-0.0463562, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=1314, text=' open', logprob=-1.3359375, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=3178, text=' source', logprob=-0.28100586, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=273, text=' and', logprob=-0.5961914, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=3426, text=' built', logprob=-1.9423828, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=271, text=' to', logprob=-1.4121094, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=314, text=' be', logprob=-1.5224609, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=1833, text=' easy', logprob=-2.1132812, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=271, text=' to', logprob=-0.08520508, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=745, text=' use', logprob=-0.39453125, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement( + id=25, + text='.', + logprob=-0.5703125, + special=False), + generated_text='100% open source and built to be easy to use.', + details=TextGenerationStreamOutputStreamDetails(finish_reason='length', generated_tokens=12, seed=None) + ) + + # Case 5: generate constrained output using grammar + >>> response = client.text_generation( + ... prompt="I saw a puppy a cat and a raccoon during my bike ride in the park", + ... model="HuggingFaceH4/zephyr-orpo-141b-A35b-v0.1", + ... max_new_tokens=100, + ... repetition_penalty=1.3, + ... grammar={ + ... "type": "json", + ... "value": { + ... "properties": { + ... "location": {"type": "string"}, + ... "activity": {"type": "string"}, + ... "animals_seen": {"type": "integer", "minimum": 1, "maximum": 5}, + ... "animals": {"type": "array", "items": {"type": "string"}}, + ... }, + ... "required": ["location", "activity", "animals_seen", "animals"], + ... }, + ... }, + ... ) + >>> json.loads(response) + { + "activity": "bike riding", + "animals": ["puppy", "cat", "raccoon"], + "animals_seen": 3, + "location": "park" + } + ``` + """ + if decoder_input_details and not details: + warnings.warn( + "`decoder_input_details=True` has been passed to the server but `details=False` is set meaning that" + " the output from the server will be truncated." + ) + decoder_input_details = False + + if stop_sequences is not None: + warnings.warn( + "`stop_sequences` is a deprecated argument for `text_generation` task" + " and will be removed in version '0.28.0'. Use `stop` instead.", + FutureWarning, + ) + if stop is None: + stop = stop_sequences # use deprecated arg if provided + + # Build payload + parameters = { + "adapter_id": adapter_id, + "best_of": best_of, + "decoder_input_details": decoder_input_details, + "details": details, + "do_sample": do_sample, + "frequency_penalty": frequency_penalty, + "grammar": grammar, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + "return_full_text": return_full_text, + "seed": seed, + "stop": stop if stop is not None else [], + "temperature": temperature, + "top_k": top_k, + "top_n_tokens": top_n_tokens, + "top_p": top_p, + "truncate": truncate, + "typical_p": typical_p, + "watermark": watermark, + } + + # Remove some parameters if not a TGI server + unsupported_kwargs = _get_unsupported_text_generation_kwargs(model) + if len(unsupported_kwargs) > 0: + # The server does not support some parameters + # => means it is not a TGI server + # => remove unsupported parameters and warn the user + + ignored_parameters = [] + for key in unsupported_kwargs: + if parameters.get(key): + ignored_parameters.append(key) + parameters.pop(key, None) + if len(ignored_parameters) > 0: + warnings.warn( + "API endpoint/model for text-generation is not served via TGI. Ignoring following parameters:" + f" {', '.join(ignored_parameters)}.", + UserWarning, + ) + if details: + warnings.warn( + "API endpoint/model for text-generation is not served via TGI. Parameter `details=True` will" + " be ignored meaning only the generated text will be returned.", + UserWarning, + ) + details = False + if stream: + raise ValueError( + "API endpoint/model for text-generation is not served via TGI. Cannot return output as a stream." + " Please pass `stream=False` as input." + ) + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-generation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters=parameters, + extra_payload={"stream": stream}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + + # Handle errors separately for more precise error messages + try: + bytes_output = self._inner_post(request_parameters, stream=stream) + except HTTPError as e: + match = MODEL_KWARGS_NOT_USED_REGEX.search(str(e)) + if isinstance(e, BadRequestError) and match: + unused_params = [kwarg.strip("' ") for kwarg in match.group(1).split(",")] + _set_unsupported_text_generation_kwargs(model, unused_params) + return self.text_generation( # type: ignore + prompt=prompt, + details=details, + stream=stream, + model=model_id, + adapter_id=adapter_id, + best_of=best_of, + decoder_input_details=decoder_input_details, + do_sample=do_sample, + frequency_penalty=frequency_penalty, + grammar=grammar, + max_new_tokens=max_new_tokens, + repetition_penalty=repetition_penalty, + return_full_text=return_full_text, + seed=seed, + stop=stop, + temperature=temperature, + top_k=top_k, + top_n_tokens=top_n_tokens, + top_p=top_p, + truncate=truncate, + typical_p=typical_p, + watermark=watermark, + ) + raise_text_generation_error(e) + + # Parse output + if stream: + return _stream_text_generation_response(bytes_output, details) # type: ignore + + data = _bytes_to_dict(bytes_output) # type: ignore[arg-type] + + # Data can be a single element (dict) or an iterable of dicts where we select the first element of. + if isinstance(data, list): + data = data[0] + response = provider_helper.get_response(data, request_parameters) + return TextGenerationOutput.parse_obj_as_instance(response) if details else response["generated_text"] + + def text_to_image( + self, + prompt: str, + *, + negative_prompt: Optional[str] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + model: Optional[str] = None, + scheduler: Optional[str] = None, + seed: Optional[int] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> "Image": + """ + Generate an image based on a given text using a specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + prompt (`str`): + The prompt to generate an image from. + negative_prompt (`str`, *optional*): + One prompt to guide what NOT to include in image generation. + height (`int`, *optional*): + The height in pixels of the output image + width (`int`, *optional*): + The width in pixels of the output image + num_inference_steps (`int`, *optional*): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*): + A higher guidance scale value encourages the model to generate images closely linked to the text + prompt, but values too high may cause saturation and other artifacts. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-image model will be used. + Defaults to None. + scheduler (`str`, *optional*): + Override the scheduler with a compatible one. + seed (`int`, *optional*): + Seed for the random number generator. + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + + Returns: + `Image`: The generated image. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + + >>> image = client.text_to_image("An astronaut riding a horse on the moon.") + >>> image.save("astronaut.png") + + >>> image = client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... negative_prompt="low resolution, blurry", + ... model="stabilityai/stable-diffusion-2-1", + ... ) + >>> image.save("better_astronaut.png") + ``` + Example using a third-party provider directly. Usage will be billed on your fal.ai account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="fal-ai", # Use fal.ai provider + ... api_key="fal-ai-api-key", # Pass your fal.ai API key + ... ) + >>> image = client.text_to_image( + ... "A majestic lion in a fantasy forest", + ... model="black-forest-labs/FLUX.1-schnell", + ... ) + >>> image.save("lion.png") + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> image = client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... model="black-forest-labs/FLUX.1-dev", + ... ) + >>> image.save("astronaut.png") + ``` + + Example using Replicate provider with extra parameters + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> image = client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... model="black-forest-labs/FLUX.1-schnell", + ... extra_body={"output_quality": 100}, + ... ) + >>> image.save("astronaut.png") + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-image", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters={ + "negative_prompt": negative_prompt, + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "scheduler": scheduler, + "seed": seed, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + response = provider_helper.get_response(response) + return _bytes_to_image(response) + + def text_to_video( + self, + prompt: str, + *, + model: Optional[str] = None, + guidance_scale: Optional[float] = None, + negative_prompt: Optional[List[str]] = None, + num_frames: Optional[float] = None, + num_inference_steps: Optional[int] = None, + seed: Optional[int] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> bytes: + """ + Generate a video based on a given text. + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + prompt (`str`): + The prompt to generate a video from. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-video model will be used. + Defaults to None. + guidance_scale (`float`, *optional*): + A higher guidance scale value encourages the model to generate videos closely linked to the text + prompt, but values too high may cause saturation and other artifacts. + negative_prompt (`List[str]`, *optional*): + One or several prompt to guide what NOT to include in video generation. + num_frames (`float`, *optional*): + The num_frames parameter determines how many video frames are generated. + num_inference_steps (`int`, *optional*): + The number of denoising steps. More denoising steps usually lead to a higher quality video at the + expense of slower inference. + seed (`int`, *optional*): + Seed for the random number generator. + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + + Returns: + `bytes`: The generated video. + + Example: + + Example using a third-party provider directly. Usage will be billed on your fal.ai account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="fal-ai", # Using fal.ai provider + ... api_key="fal-ai-api-key", # Pass your fal.ai API key + ... ) + >>> video = client.text_to_video( + ... "A majestic lion running in a fantasy forest", + ... model="tencent/HunyuanVideo", + ... ) + >>> with open("lion.mp4", "wb") as file: + ... file.write(video) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Using replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> video = client.text_to_video( + ... "A cat running in a park", + ... model="genmo/mochi-1-preview", + ... ) + >>> with open("cat.mp4", "wb") as file: + ... file.write(video) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-video", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters={ + "guidance_scale": guidance_scale, + "negative_prompt": negative_prompt, + "num_frames": num_frames, + "num_inference_steps": num_inference_steps, + "seed": seed, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + response = provider_helper.get_response(response, request_parameters) + return response + + def text_to_speech( + self, + text: str, + *, + model: Optional[str] = None, + do_sample: Optional[bool] = None, + early_stopping: Optional[Union[bool, "TextToSpeechEarlyStoppingEnum"]] = None, + epsilon_cutoff: Optional[float] = None, + eta_cutoff: Optional[float] = None, + max_length: Optional[int] = None, + max_new_tokens: Optional[int] = None, + min_length: Optional[int] = None, + min_new_tokens: Optional[int] = None, + num_beam_groups: Optional[int] = None, + num_beams: Optional[int] = None, + penalty_alpha: Optional[float] = None, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + use_cache: Optional[bool] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> bytes: + """ + Synthesize an audio of a voice pronouncing a given text. + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + text (`str`): + The text to synthesize. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-speech model will be used. + Defaults to None. + do_sample (`bool`, *optional*): + Whether to use sampling instead of greedy decoding when generating new tokens. + early_stopping (`Union[bool, "TextToSpeechEarlyStoppingEnum"]`, *optional*): + Controls the stopping condition for beam-based methods. + epsilon_cutoff (`float`, *optional*): + If set to float strictly between 0 and 1, only tokens with a conditional probability greater than + epsilon_cutoff will be sampled. In the paper, suggested values range from 3e-4 to 9e-4, depending on + the size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://hf.co/papers/2210.15191) for more details. + eta_cutoff (`float`, *optional*): + Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to float strictly + between 0 and 1, a token is only considered if it is greater than either eta_cutoff or sqrt(eta_cutoff) + * exp(-entropy(softmax(next_token_logits))). The latter term is intuitively the expected next token + probability, scaled by sqrt(eta_cutoff). In the paper, suggested values range from 3e-4 to 2e-3, + depending on the size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://hf.co/papers/2210.15191) for more details. + max_length (`int`, *optional*): + The maximum length (in tokens) of the generated text, including the input. + max_new_tokens (`int`, *optional*): + The maximum number of tokens to generate. Takes precedence over max_length. + min_length (`int`, *optional*): + The minimum length (in tokens) of the generated text, including the input. + min_new_tokens (`int`, *optional*): + The minimum number of tokens to generate. Takes precedence over min_length. + num_beam_groups (`int`, *optional*): + Number of groups to divide num_beams into in order to ensure diversity among different groups of beams. + See [this paper](https://hf.co/papers/1610.02424) for more details. + num_beams (`int`, *optional*): + Number of beams to use for beam search. + penalty_alpha (`float`, *optional*): + The value balances the model confidence and the degeneration penalty in contrastive search decoding. + temperature (`float`, *optional*): + The value used to modulate the next token probabilities. + top_k (`int`, *optional*): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + top_p (`float`, *optional*): + If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to + top_p or higher are kept for generation. + typical_p (`float`, *optional*): + Local typicality measures how similar the conditional probability of predicting a target token next is + to the expected conditional probability of predicting a random token next, given the partial text + already generated. If set to float < 1, the smallest set of the most locally typical tokens with + probabilities that add up to typical_p or higher are kept for generation. See [this + paper](https://hf.co/papers/2202.00666) for more details. + use_cache (`bool`, *optional*): + Whether the model should use the past last key/values attentions to speed up decoding + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + `bytes`: The generated audio. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from pathlib import Path + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + + >>> audio = client.text_to_speech("Hello world") + >>> Path("hello_world.flac").write_bytes(audio) + ``` + + Example using a third-party provider directly. Usage will be billed on your Replicate account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", + ... api_key="your-replicate-api-key", # Pass your Replicate API key directly + ... ) + >>> audio = client.text_to_speech( + ... text="Hello world", + ... model="OuteAI/OuteTTS-0.3-500M", + ... ) + >>> Path("hello_world.flac").write_bytes(audio) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", + ... api_key="hf_...", # Pass your HF token + ... ) + >>> audio =client.text_to_speech( + ... text="Hello world", + ... model="OuteAI/OuteTTS-0.3-500M", + ... ) + >>> Path("hello_world.flac").write_bytes(audio) + ``` + Example using Replicate provider with extra parameters + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> audio = client.text_to_speech( + ... "Hello, my name is Kororo, an awesome text-to-speech model.", + ... model="hexgrad/Kokoro-82M", + ... extra_body={"voice": "af_nicole"}, + ... ) + >>> Path("hello.flac").write_bytes(audio) + ``` + + Example music-gen using "YuE-s1-7B-anneal-en-cot" on fal.ai + ```py + >>> from huggingface_hub import InferenceClient + >>> lyrics = ''' + ... [verse] + ... In the town where I was born + ... Lived a man who sailed to sea + ... And he told us of his life + ... In the land of submarines + ... So we sailed on to the sun + ... 'Til we found a sea of green + ... And we lived beneath the waves + ... In our yellow submarine + + ... [chorus] + ... We all live in a yellow submarine + ... Yellow submarine, yellow submarine + ... We all live in a yellow submarine + ... Yellow submarine, yellow submarine + ... ''' + >>> genres = "pavarotti-style tenor voice" + >>> client = InferenceClient( + ... provider="fal-ai", + ... model="m-a-p/YuE-s1-7B-anneal-en-cot", + ... api_key=..., + ... ) + >>> audio = client.text_to_speech(lyrics, extra_body={"genres": genres}) + >>> with open("output.mp3", "wb") as f: + ... f.write(audio) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-speech", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "do_sample": do_sample, + "early_stopping": early_stopping, + "epsilon_cutoff": epsilon_cutoff, + "eta_cutoff": eta_cutoff, + "max_length": max_length, + "max_new_tokens": max_new_tokens, + "min_length": min_length, + "min_new_tokens": min_new_tokens, + "num_beam_groups": num_beam_groups, + "num_beams": num_beams, + "penalty_alpha": penalty_alpha, + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + "typical_p": typical_p, + "use_cache": use_cache, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + response = provider_helper.get_response(response) + return response + + def token_classification( + self, + text: str, + *, + model: Optional[str] = None, + aggregation_strategy: Optional["TokenClassificationAggregationStrategy"] = None, + ignore_labels: Optional[List[str]] = None, + stride: Optional[int] = None, + ) -> List[TokenClassificationOutputElement]: + """ + Perform token classification on the given text. + Usually used for sentence parsing, either grammatical, or Named Entity Recognition (NER) to understand keywords contained within text. + + Args: + text (`str`): + A string to be classified. + model (`str`, *optional*): + The model to use for the token classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended token classification model will be used. + Defaults to None. + aggregation_strategy (`"TokenClassificationAggregationStrategy"`, *optional*): + The strategy used to fuse tokens based on model predictions + ignore_labels (`List[str`, *optional*): + A list of labels to ignore + stride (`int`, *optional*): + The number of overlapping tokens between chunks when splitting the input text. + + Returns: + `List[TokenClassificationOutputElement]`: List of [`TokenClassificationOutputElement`] items containing the entity group, confidence score, word, start and end index. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.token_classification("My name is Sarah Jessica Parker but you can call me Jessica") + [ + TokenClassificationOutputElement( + entity_group='PER', + score=0.9971321225166321, + word='Sarah Jessica Parker', + start=11, + end=31, + ), + TokenClassificationOutputElement( + entity_group='PER', + score=0.9773476123809814, + word='Jessica', + start=52, + end=59, + ) + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="token-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "aggregation_strategy": aggregation_strategy, + "ignore_labels": ignore_labels, + "stride": stride, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return TokenClassificationOutputElement.parse_obj_as_list(response) + + def translation( + self, + text: str, + *, + model: Optional[str] = None, + src_lang: Optional[str] = None, + tgt_lang: Optional[str] = None, + clean_up_tokenization_spaces: Optional[bool] = None, + truncation: Optional["TranslationTruncationStrategy"] = None, + generate_parameters: Optional[Dict[str, Any]] = None, + ) -> TranslationOutput: + """ + Convert text from one language to another. + + Check out https://huggingface.co/tasks/translation for more information on how to choose the best model for + your specific use case. Source and target languages usually depend on the model. + However, it is possible to specify source and target languages for certain models. If you are working with one of these models, + you can use `src_lang` and `tgt_lang` arguments to pass the relevant information. + + Args: + text (`str`): + A string to be translated. + model (`str`, *optional*): + The model to use for the translation task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended translation model will be used. + Defaults to None. + src_lang (`str`, *optional*): + The source language of the text. Required for models that can translate from multiple languages. + tgt_lang (`str`, *optional*): + Target language to translate to. Required for models that can translate to multiple languages. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether to clean up the potential extra spaces in the text output. + truncation (`"TranslationTruncationStrategy"`, *optional*): + The truncation strategy to use. + generate_parameters (`Dict[str, Any]`, *optional*): + Additional parametrization of the text generation algorithm. + + Returns: + [`TranslationOutput`]: The generated translated text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + `ValueError`: + If only one of the `src_lang` and `tgt_lang` arguments are provided. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.translation("My name is Wolfgang and I live in Berlin") + 'Mein Name ist Wolfgang und ich lebe in Berlin.' + >>> client.translation("My name is Wolfgang and I live in Berlin", model="Helsinki-NLP/opus-mt-en-fr") + TranslationOutput(translation_text='Je m'appelle Wolfgang et je vis à Berlin.') + ``` + + Specifying languages: + ```py + >>> client.translation("My name is Sarah Jessica Parker but you can call me Jessica", model="facebook/mbart-large-50-many-to-many-mmt", src_lang="en_XX", tgt_lang="fr_XX") + "Mon nom est Sarah Jessica Parker mais vous pouvez m'appeler Jessica" + ``` + """ + # Throw error if only one of `src_lang` and `tgt_lang` was given + if src_lang is not None and tgt_lang is None: + raise ValueError("You cannot specify `src_lang` without specifying `tgt_lang`.") + + if src_lang is None and tgt_lang is not None: + raise ValueError("You cannot specify `tgt_lang` without specifying `src_lang`.") + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="translation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "src_lang": src_lang, + "tgt_lang": tgt_lang, + "clean_up_tokenization_spaces": clean_up_tokenization_spaces, + "truncation": truncation, + "generate_parameters": generate_parameters, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return TranslationOutput.parse_obj_as_list(response)[0] + + def visual_question_answering( + self, + image: ContentT, + question: str, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + ) -> List[VisualQuestionAnsweringOutputElement]: + """ + Answering open-ended questions based on an image. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for the context. It can be raw bytes, an image file, or a URL to an online image. + question (`str`): + Question to be answered. + model (`str`, *optional*): + The model to use for the visual question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended visual question answering model will be used. + Defaults to None. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Note that we return less than + topk answers if there are not enough options available within the context. + Returns: + `List[VisualQuestionAnsweringOutputElement]`: a list of [`VisualQuestionAnsweringOutputElement`] items containing the predicted label and associated probability. + + Raises: + `InferenceTimeoutError`: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.visual_question_answering( + ... image="https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg", + ... question="What is the animal doing?" + ... ) + [ + VisualQuestionAnsweringOutputElement(score=0.778609573841095, answer='laying down'), + VisualQuestionAnsweringOutputElement(score=0.6957435607910156, answer='sitting'), + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="visual-question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + extra_payload={"question": question, "image": _b64_encode(image)}, + ) + response = self._inner_post(request_parameters) + return VisualQuestionAnsweringOutputElement.parse_obj_as_list(response) + + def zero_shot_classification( + self, + text: str, + candidate_labels: List[str], + *, + multi_label: Optional[bool] = False, + hypothesis_template: Optional[str] = None, + model: Optional[str] = None, + ) -> List[ZeroShotClassificationOutputElement]: + """ + Provide as input a text and a set of candidate labels to classify the input text. + + Args: + text (`str`): + The input text to classify. + candidate_labels (`List[str]`): + The set of possible class labels to classify the text into. + labels (`List[str]`, *optional*): + (deprecated) List of strings. Each string is the verbalization of a possible label for the input text. + multi_label (`bool`, *optional*): + Whether multiple candidate labels can be true. If false, the scores are normalized such that the sum of + the label likelihoods for each sequence is 1. If true, the labels are considered independent and + probabilities are normalized for each candidate. + hypothesis_template (`str`, *optional*): + The sentence used in conjunction with `candidate_labels` to attempt the text classification by + replacing the placeholder with the candidate labels. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. If not provided, the default recommended zero-shot classification model will be used. + + + Returns: + `List[ZeroShotClassificationOutputElement]`: List of [`ZeroShotClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example with `multi_label=False`: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> text = ( + ... "A new model offers an explanation for how the Galilean satellites formed around the solar system's" + ... "largest world. Konstantin Batygin did not set out to solve one of the solar system's most puzzling" + ... " mysteries when he went for a run up a hill in Nice, France." + ... ) + >>> labels = ["space & cosmos", "scientific discovery", "microbiology", "robots", "archeology"] + >>> client.zero_shot_classification(text, labels) + [ + ZeroShotClassificationOutputElement(label='scientific discovery', score=0.7961668968200684), + ZeroShotClassificationOutputElement(label='space & cosmos', score=0.18570658564567566), + ZeroShotClassificationOutputElement(label='microbiology', score=0.00730885099619627), + ZeroShotClassificationOutputElement(label='archeology', score=0.006258360575884581), + ZeroShotClassificationOutputElement(label='robots', score=0.004559356719255447), + ] + >>> client.zero_shot_classification(text, labels, multi_label=True) + [ + ZeroShotClassificationOutputElement(label='scientific discovery', score=0.9829297661781311), + ZeroShotClassificationOutputElement(label='space & cosmos', score=0.755190908908844), + ZeroShotClassificationOutputElement(label='microbiology', score=0.0005462635890580714), + ZeroShotClassificationOutputElement(label='archeology', score=0.00047131875180639327), + ZeroShotClassificationOutputElement(label='robots', score=0.00030448526376858354), + ] + ``` + + Example with `multi_label=True` and a custom `hypothesis_template`: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.zero_shot_classification( + ... text="I really like our dinner and I'm very happy. I don't like the weather though.", + ... labels=["positive", "negative", "pessimistic", "optimistic"], + ... multi_label=True, + ... hypothesis_template="This text is {} towards the weather" + ... ) + [ + ZeroShotClassificationOutputElement(label='negative', score=0.9231801629066467), + ZeroShotClassificationOutputElement(label='pessimistic', score=0.8760990500450134), + ZeroShotClassificationOutputElement(label='optimistic', score=0.0008674879791215062), + ZeroShotClassificationOutputElement(label='positive', score=0.0005250611575320363) + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="zero-shot-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "candidate_labels": candidate_labels, + "multi_label": multi_label, + "hypothesis_template": hypothesis_template, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + output = _bytes_to_dict(response) + return [ + ZeroShotClassificationOutputElement.parse_obj_as_instance({"label": label, "score": score}) + for label, score in zip(output["labels"], output["scores"]) + ] + + def zero_shot_image_classification( + self, + image: ContentT, + candidate_labels: List[str], + *, + model: Optional[str] = None, + hypothesis_template: Optional[str] = None, + # deprecated argument + labels: List[str] = None, # type: ignore + ) -> List[ZeroShotImageClassificationOutputElement]: + """ + Provide input image and text labels to predict text labels for the image. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image to caption. It can be raw bytes, an image file, or a URL to an online image. + candidate_labels (`List[str]`): + The candidate labels for this image + labels (`List[str]`, *optional*): + (deprecated) List of string possible labels. There must be at least 2 labels. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. If not provided, the default recommended zero-shot image classification model will be used. + hypothesis_template (`str`, *optional*): + The sentence used in conjunction with `candidate_labels` to attempt the image classification by + replacing the placeholder with the candidate labels. + + Returns: + `List[ZeroShotImageClassificationOutputElement]`: List of [`ZeroShotImageClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `HTTPError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + + >>> client.zero_shot_image_classification( + ... "https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg", + ... labels=["dog", "cat", "horse"], + ... ) + [ZeroShotImageClassificationOutputElement(label='dog', score=0.956),...] + ``` + """ + # Raise ValueError if input is less than 2 labels + if len(candidate_labels) < 2: + raise ValueError("You must specify at least 2 classes to compare.") + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="zero-shot-image-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "candidate_labels": candidate_labels, + "hypothesis_template": hypothesis_template, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = self._inner_post(request_parameters) + return ZeroShotImageClassificationOutputElement.parse_obj_as_list(response) + + @_deprecate_method( + version="0.33.0", + message=( + "HF Inference API is getting revamped and will only support warm models in the future (no cold start allowed)." + " Use `HfApi.list_models(..., inference_provider='...')` to list warm models per provider." + ), + ) + def list_deployed_models( + self, frameworks: Union[None, str, Literal["all"], List[str]] = None + ) -> Dict[str, List[str]]: + """ + List models deployed on the HF Serverless Inference API service. + + This helper checks deployed models framework by framework. By default, it will check the 4 main frameworks that + are supported and account for 95% of the hosted models. However, if you want a complete list of models you can + specify `frameworks="all"` as input. Alternatively, if you know before-hand which framework you are interested + in, you can also restrict to search to this one (e.g. `frameworks="text-generation-inference"`). The more + frameworks are checked, the more time it will take. + + + + This endpoint method does not return a live list of all models available for the HF Inference API service. + It searches over a cached list of models that were recently available and the list may not be up to date. + If you want to know the live status of a specific model, use [`~InferenceClient.get_model_status`]. + + + + + + This endpoint method is mostly useful for discoverability. If you already know which model you want to use and want to + check its availability, you can directly use [`~InferenceClient.get_model_status`]. + + + + Args: + frameworks (`Literal["all"]` or `List[str]` or `str`, *optional*): + The frameworks to filter on. By default only a subset of the available frameworks are tested. If set to + "all", all available frameworks will be tested. It is also possible to provide a single framework or a + custom set of frameworks to check. + + Returns: + `Dict[str, List[str]]`: A dictionary mapping task names to a sorted list of model IDs. + + Example: + ```python + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + + # Discover zero-shot-classification models currently deployed + >>> models = client.list_deployed_models() + >>> models["zero-shot-classification"] + ['Narsil/deberta-large-mnli-zero-cls', 'facebook/bart-large-mnli', ...] + + # List from only 1 framework + >>> client.list_deployed_models("text-generation-inference") + {'text-generation': ['bigcode/starcoder', 'meta-llama/Llama-2-70b-chat-hf', ...], ...} + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Listing deployed models is not supported on '{self.provider}'.") + + # Resolve which frameworks to check + if frameworks is None: + frameworks = constants.MAIN_INFERENCE_API_FRAMEWORKS + elif frameworks == "all": + frameworks = constants.ALL_INFERENCE_API_FRAMEWORKS + elif isinstance(frameworks, str): + frameworks = [frameworks] + frameworks = list(set(frameworks)) + + # Fetch them iteratively + models_by_task: Dict[str, List[str]] = {} + + def _unpack_response(framework: str, items: List[Dict]) -> None: + for model in items: + if framework == "sentence-transformers": + # Model running with the `sentence-transformers` framework can work with both tasks even if not + # branded as such in the API response + models_by_task.setdefault("feature-extraction", []).append(model["model_id"]) + models_by_task.setdefault("sentence-similarity", []).append(model["model_id"]) + else: + models_by_task.setdefault(model["task"], []).append(model["model_id"]) + + for framework in frameworks: + response = get_session().get( + f"{constants.INFERENCE_ENDPOINT}/framework/{framework}", headers=build_hf_headers(token=self.token) + ) + hf_raise_for_status(response) + _unpack_response(framework, response.json()) + + # Sort alphabetically for discoverability and return + for task, models in models_by_task.items(): + models_by_task[task] = sorted(set(models), key=lambda x: x.lower()) + return models_by_task + + def get_endpoint_info(self, *, model: Optional[str] = None) -> Dict[str, Any]: + """ + Get information about the deployed endpoint. + + This endpoint is only available on endpoints powered by Text-Generation-Inference (TGI) or Text-Embedding-Inference (TEI). + Endpoints powered by `transformers` return an empty payload. + + Args: + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + `Dict[str, Any]`: Information about the endpoint. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> client.get_endpoint_info() + { + 'model_id': 'meta-llama/Meta-Llama-3-70B-Instruct', + 'model_sha': None, + 'model_dtype': 'torch.float16', + 'model_device_type': 'cuda', + 'model_pipeline_tag': None, + 'max_concurrent_requests': 128, + 'max_best_of': 2, + 'max_stop_sequences': 4, + 'max_input_length': 8191, + 'max_total_tokens': 8192, + 'waiting_served_ratio': 0.3, + 'max_batch_total_tokens': 1259392, + 'max_waiting_tokens': 20, + 'max_batch_size': None, + 'validation_workers': 32, + 'max_client_batch_size': 4, + 'version': '2.0.2', + 'sha': 'dccab72549635c7eb5ddb17f43f0b7cdff07c214', + 'docker_label': 'sha-dccab72' + } + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Getting endpoint info is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if model.startswith(("http://", "https://")): + url = model.rstrip("/") + "/info" + else: + url = f"{constants.INFERENCE_ENDPOINT}/models/{model}/info" + + response = get_session().get(url, headers=build_hf_headers(token=self.token)) + hf_raise_for_status(response) + return response.json() + + def health_check(self, model: Optional[str] = None) -> bool: + """ + Check the health of the deployed endpoint. + + Health check is only available with Inference Endpoints powered by Text-Generation-Inference (TGI) or Text-Embedding-Inference (TEI). + For Inference API, please use [`InferenceClient.get_model_status`] instead. + + Args: + model (`str`, *optional*): + URL of the Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + `bool`: True if everything is working fine. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient("https://jzgu0buei5.us-east-1.aws.endpoints.huggingface.cloud") + >>> client.health_check() + True + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Health check is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if not model.startswith(("http://", "https://")): + raise ValueError( + "Model must be an Inference Endpoint URL. For serverless Inference API, please use `InferenceClient.get_model_status`." + ) + url = model.rstrip("/") + "/health" + + response = get_session().get(url, headers=build_hf_headers(token=self.token)) + return response.status_code == 200 + + @_deprecate_method( + version="0.33.0", + message=( + "HF Inference API is getting revamped and will only support warm models in the future (no cold start allowed)." + " Use `HfApi.model_info` to get the model status both with HF Inference API and external providers." + ), + ) + def get_model_status(self, model: Optional[str] = None) -> ModelStatus: + """ + Get the status of a model hosted on the HF Inference API. + + + + This endpoint is mostly useful when you already know which model you want to use and want to check its + availability. If you want to discover already deployed models, you should rather use [`~InferenceClient.list_deployed_models`]. + + + + Args: + model (`str`, *optional*): + Identifier of the model for witch the status gonna be checked. If model is not provided, + the model associated with this instance of [`InferenceClient`] will be used. Only HF Inference API service can be checked so the + identifier cannot be a URL. + + + Returns: + [`ModelStatus`]: An instance of ModelStatus dataclass, containing information, + about the state of the model: load, state, compute type and framework. + + Example: + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient() + >>> client.get_model_status("meta-llama/Meta-Llama-3-8B-Instruct") + ModelStatus(loaded=True, state='Loaded', compute_type='gpu', framework='text-generation-inference') + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Getting model status is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if model.startswith("https://"): + raise NotImplementedError("Model status is only available for Inference API endpoints.") + url = f"{constants.INFERENCE_ENDPOINT}/status/{model}" + + response = get_session().get(url, headers=build_hf_headers(token=self.token)) + hf_raise_for_status(response) + response_data = response.json() + + if "error" in response_data: + raise ValueError(response_data["error"]) + + return ModelStatus( + loaded=response_data["loaded"], + state=response_data["state"], + compute_type=response_data["compute_type"], + framework=response_data["framework"], + ) + + @property + def chat(self) -> "ProxyClientChat": + return ProxyClientChat(self) + + +class _ProxyClient: + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + def __init__(self, client: InferenceClient): + self._client = client + + +class ProxyClientChat(_ProxyClient): + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + @property + def completions(self) -> "ProxyClientChatCompletions": + return ProxyClientChatCompletions(self._client) + + +class ProxyClientChatCompletions(_ProxyClient): + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + @property + def create(self): + return self._client.chat_completion diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_common.py b/lib/python3.12/site-packages/huggingface_hub/inference/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..574f726b67dfd11baa8db8914f37969306f2a925 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_common.py @@ -0,0 +1,422 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities used by both the sync and async inference clients.""" + +import base64 +import io +import json +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterable, + BinaryIO, + ContextManager, + Dict, + Generator, + Iterable, + List, + Literal, + NoReturn, + Optional, + Union, + overload, +) + +from requests import HTTPError + +from huggingface_hub.errors import ( + GenerationError, + IncompleteGenerationError, + OverloadedError, + TextGenerationError, + UnknownError, + ValidationError, +) + +from ..utils import get_session, is_aiohttp_available, is_numpy_available, is_pillow_available +from ._generated.types import ChatCompletionStreamOutput, TextGenerationStreamOutput + + +if TYPE_CHECKING: + from aiohttp import ClientResponse, ClientSession + from PIL.Image import Image + +# TYPES +UrlT = str +PathT = Union[str, Path] +BinaryT = Union[bytes, BinaryIO] +ContentT = Union[BinaryT, PathT, UrlT] + +# Use to set a Accept: image/png header +TASKS_EXPECTING_IMAGES = {"text-to-image", "image-to-image"} + +logger = logging.getLogger(__name__) + + +@dataclass +class RequestParameters: + url: str + task: str + model: Optional[str] + json: Optional[Union[str, Dict, List]] + data: Optional[ContentT] + headers: Dict[str, Any] + + +# Add dataclass for ModelStatus. We use this dataclass in get_model_status function. +@dataclass +class ModelStatus: + """ + This Dataclass represents the model status in the HF Inference API. + + Args: + loaded (`bool`): + If the model is currently loaded into HF's Inference API. Models + are loaded on-demand, leading to the user's first request taking longer. + If a model is loaded, you can be assured that it is in a healthy state. + state (`str`): + The current state of the model. This can be 'Loaded', 'Loadable', 'TooBig'. + If a model's state is 'Loadable', it's not too big and has a supported + backend. Loadable models are automatically loaded when the user first + requests inference on the endpoint. This means it is transparent for the + user to load a model, except that the first call takes longer to complete. + compute_type (`Dict`): + Information about the compute resource the model is using or will use, such as 'gpu' type and number of + replicas. + framework (`str`): + The name of the framework that the model was built with, such as 'transformers' + or 'text-generation-inference'. + """ + + loaded: bool + state: str + compute_type: Dict + framework: str + + +## IMPORT UTILS + + +def _import_aiohttp(): + # Make sure `aiohttp` is installed on the machine. + if not is_aiohttp_available(): + raise ImportError("Please install aiohttp to use `AsyncInferenceClient` (`pip install aiohttp`).") + import aiohttp + + return aiohttp + + +def _import_numpy(): + """Make sure `numpy` is installed on the machine.""" + if not is_numpy_available(): + raise ImportError("Please install numpy to use deal with embeddings (`pip install numpy`).") + import numpy + + return numpy + + +def _import_pil_image(): + """Make sure `PIL` is installed on the machine.""" + if not is_pillow_available(): + raise ImportError( + "Please install Pillow to use deal with images (`pip install Pillow`). If you don't want the image to be" + " post-processed, use `client.post(...)` and get the raw response from the server." + ) + from PIL import Image + + return Image + + +## ENCODING / DECODING UTILS + + +@overload +def _open_as_binary( + content: ContentT, +) -> ContextManager[BinaryT]: ... # means "if input is not None, output is not None" + + +@overload +def _open_as_binary( + content: Literal[None], +) -> ContextManager[Literal[None]]: ... # means "if input is None, output is None" + + +@contextmanager # type: ignore +def _open_as_binary(content: Optional[ContentT]) -> Generator[Optional[BinaryT], None, None]: + """Open `content` as a binary file, either from a URL, a local path, or raw bytes. + + Do nothing if `content` is None, + + TODO: handle a PIL.Image as input + TODO: handle base64 as input + """ + # If content is a string => must be either a URL or a path + if isinstance(content, str): + if content.startswith("https://") or content.startswith("http://"): + logger.debug(f"Downloading content from {content}") + yield get_session().get(content).content # TODO: retrieve as stream and pipe to post request ? + return + content = Path(content) + if not content.exists(): + raise FileNotFoundError( + f"File not found at {content}. If `data` is a string, it must either be a URL or a path to a local" + " file. To pass raw content, please encode it as bytes first." + ) + + # If content is a Path => open it + if isinstance(content, Path): + logger.debug(f"Opening content from {content}") + with content.open("rb") as f: + yield f + else: + # Otherwise: already a file-like object or None + yield content + + +def _b64_encode(content: ContentT) -> str: + """Encode a raw file (image, audio) into base64. Can be bytes, an opened file, a path or a URL.""" + with _open_as_binary(content) as data: + data_as_bytes = data if isinstance(data, bytes) else data.read() + return base64.b64encode(data_as_bytes).decode() + + +def _b64_to_image(encoded_image: str) -> "Image": + """Parse a base64-encoded string into a PIL Image.""" + Image = _import_pil_image() + return Image.open(io.BytesIO(base64.b64decode(encoded_image))) + + +def _bytes_to_list(content: bytes) -> List: + """Parse bytes from a Response object into a Python list. + + Expects the response body to be JSON-encoded data. + + NOTE: This is exactly the same implementation as `_bytes_to_dict` and will not complain if the returned data is a + dictionary. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect. + """ + return json.loads(content.decode()) + + +def _bytes_to_dict(content: bytes) -> Dict: + """Parse bytes from a Response object into a Python dictionary. + + Expects the response body to be JSON-encoded data. + + NOTE: This is exactly the same implementation as `_bytes_to_list` and will not complain if the returned data is a + list. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect. + """ + return json.loads(content.decode()) + + +def _bytes_to_image(content: bytes) -> "Image": + """Parse bytes from a Response object into a PIL Image. + + Expects the response body to be raw bytes. To deal with b64 encoded images, use `_b64_to_image` instead. + """ + Image = _import_pil_image() + return Image.open(io.BytesIO(content)) + + +def _as_dict(response: Union[bytes, Dict]) -> Dict: + return json.loads(response) if isinstance(response, bytes) else response + + +## PAYLOAD UTILS + + +## STREAMING UTILS + + +def _stream_text_generation_response( + bytes_output_as_lines: Iterable[bytes], details: bool +) -> Union[Iterable[str], Iterable[TextGenerationStreamOutput]]: + """Used in `InferenceClient.text_generation`.""" + # Parse ServerSentEvents + for byte_payload in bytes_output_as_lines: + try: + output = _format_text_generation_stream_output(byte_payload, details) + except StopIteration: + break + if output is not None: + yield output + + +async def _async_stream_text_generation_response( + bytes_output_as_lines: AsyncIterable[bytes], details: bool +) -> Union[AsyncIterable[str], AsyncIterable[TextGenerationStreamOutput]]: + """Used in `AsyncInferenceClient.text_generation`.""" + # Parse ServerSentEvents + async for byte_payload in bytes_output_as_lines: + try: + output = _format_text_generation_stream_output(byte_payload, details) + except StopIteration: + break + if output is not None: + yield output + + +def _format_text_generation_stream_output( + byte_payload: bytes, details: bool +) -> Optional[Union[str, TextGenerationStreamOutput]]: + if not byte_payload.startswith(b"data:"): + return None # empty line + + if byte_payload.strip() == b"data: [DONE]": + raise StopIteration("[DONE] signal received.") + + # Decode payload + payload = byte_payload.decode("utf-8") + json_payload = json.loads(payload.lstrip("data:").rstrip("/n")) + + # Either an error as being returned + if json_payload.get("error") is not None: + raise _parse_text_generation_error(json_payload["error"], json_payload.get("error_type")) + + # Or parse token payload + output = TextGenerationStreamOutput.parse_obj_as_instance(json_payload) + return output.token.text if not details else output + + +def _stream_chat_completion_response( + bytes_lines: Iterable[bytes], +) -> Iterable[ChatCompletionStreamOutput]: + """Used in `InferenceClient.chat_completion` if model is served with TGI.""" + for item in bytes_lines: + try: + output = _format_chat_completion_stream_output(item) + except StopIteration: + break + if output is not None: + yield output + + +async def _async_stream_chat_completion_response( + bytes_lines: AsyncIterable[bytes], +) -> AsyncIterable[ChatCompletionStreamOutput]: + """Used in `AsyncInferenceClient.chat_completion`.""" + async for item in bytes_lines: + try: + output = _format_chat_completion_stream_output(item) + except StopIteration: + break + if output is not None: + yield output + + +def _format_chat_completion_stream_output( + byte_payload: bytes, +) -> Optional[ChatCompletionStreamOutput]: + if not byte_payload.startswith(b"data:"): + return None # empty line + + if byte_payload.strip() == b"data: [DONE]": + raise StopIteration("[DONE] signal received.") + + # Decode payload + payload = byte_payload.decode("utf-8") + json_payload = json.loads(payload.lstrip("data:").rstrip("/n")) + + # Either an error as being returned + if json_payload.get("error") is not None: + raise _parse_text_generation_error(json_payload["error"], json_payload.get("error_type")) + + # Or parse token payload + return ChatCompletionStreamOutput.parse_obj_as_instance(json_payload) + + +async def _async_yield_from(client: "ClientSession", response: "ClientResponse") -> AsyncIterable[bytes]: + async for byte_payload in response.content: + yield byte_payload.strip() + await client.close() + + +# "TGI servers" are servers running with the `text-generation-inference` backend. +# This backend is the go-to solution to run large language models at scale. However, +# for some smaller models (e.g. "gpt2") the default `transformers` + `api-inference` +# solution is still in use. +# +# Both approaches have very similar APIs, but not exactly the same. What we do first in +# the `text_generation` method is to assume the model is served via TGI. If we realize +# it's not the case (i.e. we receive an HTTP 400 Bad Request), we fallback to the +# default API with a warning message. When that's the case, We remember the unsupported +# attributes for this model in the `_UNSUPPORTED_TEXT_GENERATION_KWARGS` global variable. +# +# In addition, TGI servers have a built-in API route for chat-completion, which is not +# available on the default API. We use this route to provide a more consistent behavior +# when available. +# +# For more details, see https://github.com/huggingface/text-generation-inference and +# https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task. + +_UNSUPPORTED_TEXT_GENERATION_KWARGS: Dict[Optional[str], List[str]] = {} + + +def _set_unsupported_text_generation_kwargs(model: Optional[str], unsupported_kwargs: List[str]) -> None: + _UNSUPPORTED_TEXT_GENERATION_KWARGS.setdefault(model, []).extend(unsupported_kwargs) + + +def _get_unsupported_text_generation_kwargs(model: Optional[str]) -> List[str]: + return _UNSUPPORTED_TEXT_GENERATION_KWARGS.get(model, []) + + +# TEXT GENERATION ERRORS +# ---------------------- +# Text-generation errors are parsed separately to handle as much as possible the errors returned by the text generation +# inference project (https://github.com/huggingface/text-generation-inference). +# ---------------------- + + +def raise_text_generation_error(http_error: HTTPError) -> NoReturn: + """ + Try to parse text-generation-inference error message and raise HTTPError in any case. + + Args: + error (`HTTPError`): + The HTTPError that have been raised. + """ + # Try to parse a Text Generation Inference error + + try: + # Hacky way to retrieve payload in case of aiohttp error + payload = getattr(http_error, "response_error_payload", None) or http_error.response.json() + error = payload.get("error") + error_type = payload.get("error_type") + except Exception: # no payload + raise http_error + + # If error_type => more information than `hf_raise_for_status` + if error_type is not None: + exception = _parse_text_generation_error(error, error_type) + raise exception from http_error + + # Otherwise, fallback to default error + raise http_error + + +def _parse_text_generation_error(error: Optional[str], error_type: Optional[str]) -> TextGenerationError: + if error_type == "generation": + return GenerationError(error) # type: ignore + if error_type == "incomplete_generation": + return IncompleteGenerationError(error) # type: ignore + if error_type == "overloaded": + return OverloadedError(error) # type: ignore + if error_type == "validation": + return ValidationError(error) # type: ignore + return UnknownError(error) # type: ignore diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__init__.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24b2ee5e7fef73216045b9f3309398fdfdd450e1 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/_async_client.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/_async_client.py new file mode 100644 index 0000000000000000000000000000000000000000..5c89f2f7824e7e8b26c021507480ed55d3f458d3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/_async_client.py @@ -0,0 +1,3586 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# 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. +# +# WARNING +# This entire file has been adapted from the sync-client code in `src/huggingface_hub/inference/_client.py`. +# Any change in InferenceClient will be automatically reflected in AsyncInferenceClient. +# To re-generate the code, run `make style` or `python ./utils/generate_async_inference_client.py --update`. +# WARNING +import asyncio +import base64 +import logging +import re +import warnings +from typing import TYPE_CHECKING, Any, AsyncIterable, Dict, List, Literal, Optional, Set, Union, overload + +from huggingface_hub import constants +from huggingface_hub.errors import InferenceTimeoutError +from huggingface_hub.inference._common import ( + TASKS_EXPECTING_IMAGES, + ContentT, + ModelStatus, + RequestParameters, + _async_stream_chat_completion_response, + _async_stream_text_generation_response, + _b64_encode, + _b64_to_image, + _bytes_to_dict, + _bytes_to_image, + _bytes_to_list, + _get_unsupported_text_generation_kwargs, + _import_numpy, + _open_as_binary, + _set_unsupported_text_generation_kwargs, + raise_text_generation_error, +) +from huggingface_hub.inference._generated.types import ( + AudioClassificationOutputElement, + AudioClassificationOutputTransform, + AudioToAudioOutputElement, + AutomaticSpeechRecognitionOutput, + ChatCompletionInputGrammarType, + ChatCompletionInputStreamOptions, + ChatCompletionInputTool, + ChatCompletionInputToolChoiceClass, + ChatCompletionInputToolChoiceEnum, + ChatCompletionOutput, + ChatCompletionStreamOutput, + DocumentQuestionAnsweringOutputElement, + FillMaskOutputElement, + ImageClassificationOutputElement, + ImageClassificationOutputTransform, + ImageSegmentationOutputElement, + ImageSegmentationSubtask, + ImageToImageTargetSize, + ImageToTextOutput, + ObjectDetectionOutputElement, + Padding, + QuestionAnsweringOutputElement, + SummarizationOutput, + SummarizationTruncationStrategy, + TableQuestionAnsweringOutputElement, + TextClassificationOutputElement, + TextClassificationOutputTransform, + TextGenerationInputGrammarType, + TextGenerationOutput, + TextGenerationStreamOutput, + TextToSpeechEarlyStoppingEnum, + TokenClassificationAggregationStrategy, + TokenClassificationOutputElement, + TranslationOutput, + TranslationTruncationStrategy, + VisualQuestionAnsweringOutputElement, + ZeroShotClassificationOutputElement, + ZeroShotImageClassificationOutputElement, +) +from huggingface_hub.inference._providers import PROVIDER_T, get_provider_helper +from huggingface_hub.utils import build_hf_headers, get_session, hf_raise_for_status +from huggingface_hub.utils._auth import get_token +from huggingface_hub.utils._deprecation import _deprecate_method + +from .._common import _async_yield_from, _import_aiohttp + + +if TYPE_CHECKING: + import numpy as np + from aiohttp import ClientResponse, ClientSession + from PIL.Image import Image + +logger = logging.getLogger(__name__) + + +MODEL_KWARGS_NOT_USED_REGEX = re.compile(r"The following `model_kwargs` are not used by the model: \[(.*?)\]") + + +class AsyncInferenceClient: + """ + Initialize a new Inference Client. + + [`InferenceClient`] aims to provide a unified experience to perform inference. The client can be used + seamlessly with either the (free) Inference API, self-hosted Inference Endpoints, or third-party Inference Providers. + + Args: + model (`str`, `optional`): + The model to run inference with. Can be a model id hosted on the Hugging Face Hub, e.g. `meta-llama/Meta-Llama-3-8B-Instruct` + or a URL to a deployed Inference Endpoint. Defaults to None, in which case a recommended model is + automatically selected for the task. + Note: for better compatibility with OpenAI's client, `model` has been aliased as `base_url`. Those 2 + arguments are mutually exclusive. If using `base_url` for chat completion, the `/chat/completions` suffix + path will be appended to the base URL (see the [TGI Messages API](https://huggingface.co/docs/text-generation-inference/en/messages_api) + documentation for details). When passing a URL as `model`, the client will not append any suffix path to it. + provider (`str`, *optional*): + Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"fireworks-ai"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"openai"`, `"replicate"`, "sambanova"` or `"together"`. + Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers. + If model is a URL or `base_url` is passed, then `provider` is not used. + token (`str`, *optional*): + Hugging Face token. Will default to the locally saved token if not provided. + Note: for better compatibility with OpenAI's client, `token` has been aliased as `api_key`. Those 2 + arguments are mutually exclusive and have the exact same behavior. + timeout (`float`, `optional`): + The maximum number of seconds to wait for a response from the server. Defaults to None, meaning it will loop until the server is available. + headers (`Dict[str, str]`, `optional`): + Additional headers to send to the server. By default only the authorization and user-agent headers are sent. + Values in this dictionary will override the default values. + bill_to (`str`, `optional`): + The billing account to use for the requests. By default the requests are billed on the user's account. + Requests can only be billed to an organization the user is a member of, and which has subscribed to Enterprise Hub. + cookies (`Dict[str, str]`, `optional`): + Additional cookies to send to the server. + trust_env ('bool', 'optional'): + Trust environment settings for proxy configuration if the parameter is `True` (`False` by default). + proxies (`Any`, `optional`): + Proxies to use for the request. + base_url (`str`, `optional`): + Base URL to run inference. This is a duplicated argument from `model` to make [`InferenceClient`] + follow the same pattern as `openai.OpenAI` client. Cannot be used if `model` is set. Defaults to None. + api_key (`str`, `optional`): + Token to use for authentication. This is a duplicated argument from `token` to make [`InferenceClient`] + follow the same pattern as `openai.OpenAI` client. Cannot be used if `token` is set. Defaults to None. + """ + + def __init__( + self, + model: Optional[str] = None, + *, + provider: Union[Literal["auto"], PROVIDER_T, None] = None, + token: Optional[str] = None, + timeout: Optional[float] = None, + headers: Optional[Dict[str, str]] = None, + cookies: Optional[Dict[str, str]] = None, + trust_env: bool = False, + proxies: Optional[Any] = None, + bill_to: Optional[str] = None, + # OpenAI compatibility + base_url: Optional[str] = None, + api_key: Optional[str] = None, + ) -> None: + if model is not None and base_url is not None: + raise ValueError( + "Received both `model` and `base_url` arguments. Please provide only one of them." + " `base_url` is an alias for `model` to make the API compatible with OpenAI's client." + " If using `base_url` for chat completion, the `/chat/completions` suffix path will be appended to the base url." + " When passing a URL as `model`, the client will not append any suffix path to it." + ) + if token is not None and api_key is not None: + raise ValueError( + "Received both `token` and `api_key` arguments. Please provide only one of them." + " `api_key` is an alias for `token` to make the API compatible with OpenAI's client." + " It has the exact same behavior as `token`." + ) + token = token if token is not None else api_key + if isinstance(token, bool): + # Legacy behavior: previously is was possible to pass `token=False` to disable authentication. This is not + # supported anymore as authentication is required. Better to explicitly raise here rather than risking + # sending the locally saved token without the user knowing about it. + if token is False: + raise ValueError( + "Cannot use `token=False` to disable authentication as authentication is required to run Inference." + ) + warnings.warn( + "Using `token=True` to automatically use the locally saved token is deprecated and will be removed in a future release. " + "Please use `token=None` instead (default).", + DeprecationWarning, + ) + token = get_token() + + self.model: Optional[str] = base_url or model + self.token: Optional[str] = token + + self.headers = {**headers} if headers is not None else {} + if bill_to is not None: + if ( + constants.HUGGINGFACE_HEADER_X_BILL_TO in self.headers + and self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO] != bill_to + ): + warnings.warn( + f"Overriding existing '{self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO]}' value in headers with '{bill_to}'.", + UserWarning, + ) + self.headers[constants.HUGGINGFACE_HEADER_X_BILL_TO] = bill_to + + if token is not None and not token.startswith("hf_"): + warnings.warn( + "You've provided an external provider's API key, so requests will be billed directly by the provider. " + "The `bill_to` parameter is only applicable for Hugging Face billing and will be ignored.", + UserWarning, + ) + + # Configure provider + self.provider = provider + + self.cookies = cookies + self.timeout = timeout + self.trust_env = trust_env + self.proxies = proxies + + # Keep track of the sessions to close them properly + self._sessions: Dict["ClientSession", Set["ClientResponse"]] = dict() + + def __repr__(self): + return f"" + + @overload + async def _inner_post( # type: ignore[misc] + self, request_parameters: RequestParameters, *, stream: Literal[False] = ... + ) -> bytes: ... + + @overload + async def _inner_post( # type: ignore[misc] + self, request_parameters: RequestParameters, *, stream: Literal[True] = ... + ) -> AsyncIterable[bytes]: ... + + @overload + async def _inner_post( + self, request_parameters: RequestParameters, *, stream: bool = False + ) -> Union[bytes, AsyncIterable[bytes]]: ... + + async def _inner_post( + self, request_parameters: RequestParameters, *, stream: bool = False + ) -> Union[bytes, AsyncIterable[bytes]]: + """Make a request to the inference server.""" + + aiohttp = _import_aiohttp() + + # TODO: this should be handled in provider helpers directly + if request_parameters.task in TASKS_EXPECTING_IMAGES and "Accept" not in request_parameters.headers: + request_parameters.headers["Accept"] = "image/png" + + with _open_as_binary(request_parameters.data) as data_as_binary: + # Do not use context manager as we don't want to close the connection immediately when returning + # a stream + session = self._get_client_session(headers=request_parameters.headers) + + try: + response = await session.post( + request_parameters.url, json=request_parameters.json, data=data_as_binary, proxy=self.proxies + ) + response_error_payload = None + if response.status != 200: + try: + response_error_payload = await response.json() # get payload before connection closed + except Exception: + pass + response.raise_for_status() + if stream: + return _async_yield_from(session, response) + else: + content = await response.read() + await session.close() + return content + except asyncio.TimeoutError as error: + await session.close() + # Convert any `TimeoutError` to a `InferenceTimeoutError` + raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error # type: ignore + except aiohttp.ClientResponseError as error: + error.response_error_payload = response_error_payload + await session.close() + raise error + except Exception: + await session.close() + raise + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.close() + + def __del__(self): + if len(self._sessions) > 0: + warnings.warn( + "Deleting 'AsyncInferenceClient' client but some sessions are still open. " + "This can happen if you've stopped streaming data from the server before the stream was complete. " + "To close the client properly, you must call `await client.close()` " + "or use an async context (e.g. `async with AsyncInferenceClient(): ...`." + ) + + async def close(self): + """Close all open sessions. + + By default, 'aiohttp.ClientSession' objects are closed automatically when a call is completed. However, if you + are streaming data from the server and you stop before the stream is complete, you must call this method to + close the session properly. + + Another possibility is to use an async context (e.g. `async with AsyncInferenceClient(): ...`). + """ + await asyncio.gather(*[session.close() for session in self._sessions.keys()]) + + async def audio_classification( + self, + audio: ContentT, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + function_to_apply: Optional["AudioClassificationOutputTransform"] = None, + ) -> List[AudioClassificationOutputElement]: + """ + Perform audio classification on the provided audio content. + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The audio content to classify. It can be raw audio bytes, a local audio file, or a URL pointing to an + audio file. + model (`str`, *optional*): + The model to use for audio classification. Can be a model ID hosted on the Hugging Face Hub + or a URL to a deployed Inference Endpoint. If not provided, the default recommended model for + audio classification will be used. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + function_to_apply (`"AudioClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + + Returns: + `List[AudioClassificationOutputElement]`: List of [`AudioClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.audio_classification("audio.flac") + [ + AudioClassificationOutputElement(score=0.4976358711719513, label='hap'), + AudioClassificationOutputElement(score=0.3677836060523987, label='neu'), + ... + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="audio-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={"function_to_apply": function_to_apply, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return AudioClassificationOutputElement.parse_obj_as_list(response) + + async def audio_to_audio( + self, + audio: ContentT, + *, + model: Optional[str] = None, + ) -> List[AudioToAudioOutputElement]: + """ + Performs multiple tasks related to audio-to-audio depending on the model (eg: speech enhancement, source separation). + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The audio content for the model. It can be raw audio bytes, a local audio file, or a URL pointing to an + audio file. + model (`str`, *optional*): + The model can be any model which takes an audio file and returns another audio file. Can be a model ID hosted on the Hugging Face Hub + or a URL to a deployed Inference Endpoint. If not provided, the default recommended model for + audio_to_audio will be used. + + Returns: + `List[AudioToAudioOutputElement]`: A list of [`AudioToAudioOutputElement`] items containing audios label, content-type, and audio content in blob. + + Raises: + `InferenceTimeoutError`: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> audio_output = await client.audio_to_audio("audio.flac") + >>> async for i, item in enumerate(audio_output): + >>> with open(f"output_{i}.flac", "wb") as f: + f.write(item.blob) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="audio-to-audio", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + audio_output = AudioToAudioOutputElement.parse_obj_as_list(response) + for item in audio_output: + item.blob = base64.b64decode(item.blob) + return audio_output + + async def automatic_speech_recognition( + self, + audio: ContentT, + *, + model: Optional[str] = None, + extra_body: Optional[Dict] = None, + ) -> AutomaticSpeechRecognitionOutput: + """ + Perform automatic speech recognition (ASR or audio-to-text) on the given audio content. + + Args: + audio (Union[str, Path, bytes, BinaryIO]): + The content to transcribe. It can be raw audio bytes, local audio file, or a URL to an audio file. + model (`str`, *optional*): + The model to use for ASR. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for ASR will be used. + extra_body (`Dict`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + [`AutomaticSpeechRecognitionOutput`]: An item containing the transcribed text and optionally the timestamp chunks. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.automatic_speech_recognition("hello_world.flac").text + "hello world" + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="automatic-speech-recognition", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=audio, + parameters={**(extra_body or {})}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return AutomaticSpeechRecognitionOutput.parse_obj_as_instance(response) + + @overload + async def chat_completion( # type: ignore + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: Literal[False] = False, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> ChatCompletionOutput: ... + + @overload + async def chat_completion( # type: ignore + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: Literal[True] = True, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> AsyncIterable[ChatCompletionStreamOutput]: ... + + @overload + async def chat_completion( + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: bool = False, + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> Union[ChatCompletionOutput, AsyncIterable[ChatCompletionStreamOutput]]: ... + + async def chat_completion( + self, + messages: List[Dict], + *, + model: Optional[str] = None, + stream: bool = False, + # Parameters from ChatCompletionInput (handled manually) + frequency_penalty: Optional[float] = None, + logit_bias: Optional[List[float]] = None, + logprobs: Optional[bool] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + presence_penalty: Optional[float] = None, + response_format: Optional[ChatCompletionInputGrammarType] = None, + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stream_options: Optional[ChatCompletionInputStreamOptions] = None, + temperature: Optional[float] = None, + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None, + tool_prompt: Optional[str] = None, + tools: Optional[List[ChatCompletionInputTool]] = None, + top_logprobs: Optional[int] = None, + top_p: Optional[float] = None, + extra_body: Optional[Dict] = None, + ) -> Union[ChatCompletionOutput, AsyncIterable[ChatCompletionStreamOutput]]: + """ + A method for completing conversations using a specified language model. + + + + The `client.chat_completion` method is aliased as `client.chat.completions.create` for compatibility with OpenAI's client. + Inputs and outputs are strictly the same and using either syntax will yield the same results. + Check out the [Inference guide](https://huggingface.co/docs/huggingface_hub/guides/inference#openai-compatibility) + for more details about OpenAI's compatibility. + + + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + messages (List of [`ChatCompletionInputMessage`]): + Conversation history consisting of roles and content pairs. + model (`str`, *optional*): + The model to use for chat-completion. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for chat-based text-generation will be used. + See https://huggingface.co/tasks/text-generation for more details. + If `model` is a model ID, it is passed to the server as the `model` parameter. If you want to define a + custom URL while setting `model` in the request payload, you must set `base_url` when initializing [`InferenceClient`]. + frequency_penalty (`float`, *optional*): + Penalizes new tokens based on their existing frequency + in the text so far. Range: [-2.0, 2.0]. Defaults to 0.0. + logit_bias (`List[float]`, *optional*): + Adjusts the likelihood of specific tokens appearing in the generated output. + logprobs (`bool`, *optional*): + Whether to return log probabilities of the output tokens or not. If true, returns the log + probabilities of each output token returned in the content of message. + max_tokens (`int`, *optional*): + Maximum number of tokens allowed in the response. Defaults to 100. + n (`int`, *optional*): + The number of completions to generate for each prompt. + presence_penalty (`float`, *optional*): + Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the + text so far, increasing the model's likelihood to talk about new topics. + response_format ([`ChatCompletionInputGrammarType`], *optional*): + Grammar constraints. Can be either a JSONSchema or a regex. + seed (Optional[`int`], *optional*): + Seed for reproducible control flow. Defaults to None. + stop (`List[str]`, *optional*): + Up to four strings which trigger the end of the response. + Defaults to None. + stream (`bool`, *optional*): + Enable realtime streaming of responses. Defaults to False. + stream_options ([`ChatCompletionInputStreamOptions`], *optional*): + Options for streaming completions. + temperature (`float`, *optional*): + Controls randomness of the generations. Lower values ensure + less random completions. Range: [0, 2]. Defaults to 1.0. + top_logprobs (`int`, *optional*): + An integer between 0 and 5 specifying the number of most likely tokens to return at each token + position, each with an associated log probability. logprobs must be set to true if this parameter is + used. + top_p (`float`, *optional*): + Fraction of the most likely next words to sample from. + Must be between 0 and 1. Defaults to 1.0. + tool_choice ([`ChatCompletionInputToolChoiceClass`] or [`ChatCompletionInputToolChoiceEnum`], *optional*): + The tool to use for the completion. Defaults to "auto". + tool_prompt (`str`, *optional*): + A prompt to be appended before the tools. + tools (List of [`ChatCompletionInputTool`], *optional*): + A list of tools the model may call. Currently, only functions are supported as a tool. Use this to + provide a list of functions the model may generate JSON inputs for. + extra_body (`Dict`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + [`ChatCompletionOutput`] or Iterable of [`ChatCompletionStreamOutput`]: + Generated text returned from the server: + - if `stream=False`, the generated text is returned as a [`ChatCompletionOutput`] (default). + - if `stream=True`, the generated text is returned token by token as a sequence of [`ChatCompletionStreamOutput`]. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> messages = [{"role": "user", "content": "What is the capital of France?"}] + >>> client = AsyncInferenceClient("meta-llama/Meta-Llama-3-8B-Instruct") + >>> await client.chat_completion(messages, max_tokens=100) + ChatCompletionOutput( + choices=[ + ChatCompletionOutputComplete( + finish_reason='eos_token', + index=0, + message=ChatCompletionOutputMessage( + role='assistant', + content='The capital of France is Paris.', + name=None, + tool_calls=None + ), + logprobs=None + ) + ], + created=1719907176, + id='', + model='meta-llama/Meta-Llama-3-8B-Instruct', + object='text_completion', + system_fingerprint='2.0.4-sha-f426a33', + usage=ChatCompletionOutputUsage( + completion_tokens=8, + prompt_tokens=17, + total_tokens=25 + ) + ) + ``` + + Example using streaming: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> messages = [{"role": "user", "content": "What is the capital of France?"}] + >>> client = AsyncInferenceClient("meta-llama/Meta-Llama-3-8B-Instruct") + >>> async for token in await client.chat_completion(messages, max_tokens=10, stream=True): + ... print(token) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content='The', role='assistant'), index=0, finish_reason=None)], created=1710498504) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content=' capital', role='assistant'), index=0, finish_reason=None)], created=1710498504) + (...) + ChatCompletionStreamOutput(choices=[ChatCompletionStreamOutputChoice(delta=ChatCompletionStreamOutputDelta(content=' may', role='assistant'), index=0, finish_reason=None)], created=1710498504) + ``` + + Example using OpenAI's syntax: + ```py + # Must be run in an async context + # instead of `from openai import OpenAI` + from huggingface_hub import AsyncInferenceClient + + # instead of `client = OpenAI(...)` + client = AsyncInferenceClient( + base_url=..., + api_key=..., + ) + + output = await client.chat.completions.create( + model="meta-llama/Meta-Llama-3-8B-Instruct", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Count to 10"}, + ], + stream=True, + max_tokens=1024, + ) + + for chunk in output: + print(chunk.choices[0].delta.content) + ``` + + Example using a third-party provider directly with extra (provider-specific) parameters. Usage will be billed on your Together AI account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="together", # Use Together AI provider + ... api_key="", # Pass your Together API key directly + ... ) + >>> client.chat_completion( + ... model="meta-llama/Meta-Llama-3-8B-Instruct", + ... messages=[{"role": "user", "content": "What is the capital of France?"}], + ... extra_body={"safety_model": "Meta-Llama/Llama-Guard-7b"}, + ... ) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="sambanova", # Use Sambanova provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> client.chat_completion( + ... model="meta-llama/Meta-Llama-3-8B-Instruct", + ... messages=[{"role": "user", "content": "What is the capital of France?"}], + ... ) + ``` + + Example using Image + Text as input: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + + # provide a remote URL + >>> image_url ="https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg" + # or a base64-encoded image + >>> image_path = "/path/to/image.jpeg" + >>> with open(image_path, "rb") as f: + ... base64_image = base64.b64encode(f.read()).decode("utf-8") + >>> image_url = f"data:image/jpeg;base64,{base64_image}" + + >>> client = AsyncInferenceClient("meta-llama/Llama-3.2-11B-Vision-Instruct") + >>> output = await client.chat.completions.create( + ... messages=[ + ... { + ... "role": "user", + ... "content": [ + ... { + ... "type": "image_url", + ... "image_url": {"url": image_url}, + ... }, + ... { + ... "type": "text", + ... "text": "Describe this image in one sentence.", + ... }, + ... ], + ... }, + ... ], + ... ) + >>> output + The image depicts the iconic Statue of Liberty situated in New York Harbor, New York, on a clear day. + ``` + + Example using tools: + ```py + # Must be run in an async context + >>> client = AsyncInferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> messages = [ + ... { + ... "role": "system", + ... "content": "Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.", + ... }, + ... { + ... "role": "user", + ... "content": "What's the weather like the next 3 days in San Francisco, CA?", + ... }, + ... ] + >>> tools = [ + ... { + ... "type": "function", + ... "function": { + ... "name": "get_current_weather", + ... "description": "Get the current weather", + ... "parameters": { + ... "type": "object", + ... "properties": { + ... "location": { + ... "type": "string", + ... "description": "The city and state, e.g. San Francisco, CA", + ... }, + ... "format": { + ... "type": "string", + ... "enum": ["celsius", "fahrenheit"], + ... "description": "The temperature unit to use. Infer this from the users location.", + ... }, + ... }, + ... "required": ["location", "format"], + ... }, + ... }, + ... }, + ... { + ... "type": "function", + ... "function": { + ... "name": "get_n_day_weather_forecast", + ... "description": "Get an N-day weather forecast", + ... "parameters": { + ... "type": "object", + ... "properties": { + ... "location": { + ... "type": "string", + ... "description": "The city and state, e.g. San Francisco, CA", + ... }, + ... "format": { + ... "type": "string", + ... "enum": ["celsius", "fahrenheit"], + ... "description": "The temperature unit to use. Infer this from the users location.", + ... }, + ... "num_days": { + ... "type": "integer", + ... "description": "The number of days to forecast", + ... }, + ... }, + ... "required": ["location", "format", "num_days"], + ... }, + ... }, + ... }, + ... ] + + >>> response = await client.chat_completion( + ... model="meta-llama/Meta-Llama-3-70B-Instruct", + ... messages=messages, + ... tools=tools, + ... tool_choice="auto", + ... max_tokens=500, + ... ) + >>> response.choices[0].message.tool_calls[0].function + ChatCompletionOutputFunctionDefinition( + arguments={ + 'location': 'San Francisco, CA', + 'format': 'fahrenheit', + 'num_days': 3 + }, + name='get_n_day_weather_forecast', + description=None + ) + ``` + + Example using response_format: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> messages = [ + ... { + ... "role": "user", + ... "content": "I saw a puppy a cat and a raccoon during my bike ride in the park. What did I saw and when?", + ... }, + ... ] + >>> response_format = { + ... "type": "json", + ... "value": { + ... "properties": { + ... "location": {"type": "string"}, + ... "activity": {"type": "string"}, + ... "animals_seen": {"type": "integer", "minimum": 1, "maximum": 5}, + ... "animals": {"type": "array", "items": {"type": "string"}}, + ... }, + ... "required": ["location", "activity", "animals_seen", "animals"], + ... }, + ... } + >>> response = await client.chat_completion( + ... messages=messages, + ... response_format=response_format, + ... max_tokens=500, + ... ) + >>> response.choices[0].message.content + '{\n\n"activity": "bike ride",\n"animals": ["puppy", "cat", "raccoon"],\n"animals_seen": 3,\n"location": "park"}' + ``` + """ + # Since `chat_completion(..., model=xxx)` is also a payload parameter for the server, we need to handle 'model' differently. + # `self.model` takes precedence over 'model' argument for building URL. + # `model` takes precedence for payload value. + model_id_or_url = self.model or model + payload_model = model or self.model + + # Get the provider helper + provider_helper = get_provider_helper( + self.provider, + task="conversational", + model=model_id_or_url + if model_id_or_url is not None and model_id_or_url.startswith(("http://", "https://")) + else payload_model, + ) + + # Prepare the payload + parameters = { + "model": payload_model, + "frequency_penalty": frequency_penalty, + "logit_bias": logit_bias, + "logprobs": logprobs, + "max_tokens": max_tokens, + "n": n, + "presence_penalty": presence_penalty, + "response_format": response_format, + "seed": seed, + "stop": stop, + "temperature": temperature, + "tool_choice": tool_choice, + "tool_prompt": tool_prompt, + "tools": tools, + "top_logprobs": top_logprobs, + "top_p": top_p, + "stream": stream, + "stream_options": stream_options, + **(extra_body or {}), + } + request_parameters = provider_helper.prepare_request( + inputs=messages, + parameters=parameters, + headers=self.headers, + model=model_id_or_url, + api_key=self.token, + ) + data = await self._inner_post(request_parameters, stream=stream) + + if stream: + return _async_stream_chat_completion_response(data) # type: ignore[arg-type] + + return ChatCompletionOutput.parse_obj_as_instance(data) # type: ignore[arg-type] + + async def document_question_answering( + self, + image: ContentT, + question: str, + *, + model: Optional[str] = None, + doc_stride: Optional[int] = None, + handle_impossible_answer: Optional[bool] = None, + lang: Optional[str] = None, + max_answer_len: Optional[int] = None, + max_question_len: Optional[int] = None, + max_seq_len: Optional[int] = None, + top_k: Optional[int] = None, + word_boxes: Optional[List[Union[List[float], str]]] = None, + ) -> List[DocumentQuestionAnsweringOutputElement]: + """ + Answer questions on document images. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for the context. It can be raw bytes, an image file, or a URL to an online image. + question (`str`): + Question to be answered. + model (`str`, *optional*): + The model to use for the document question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended document question answering model will be used. + Defaults to None. + doc_stride (`int`, *optional*): + If the words in the document are too long to fit with the question for the model, it will be split in + several chunks with some overlap. This argument controls the size of that overlap. + handle_impossible_answer (`bool`, *optional*): + Whether to accept impossible as an answer + lang (`str`, *optional*): + Language to use while running OCR. Defaults to english. + max_answer_len (`int`, *optional*): + The maximum length of predicted answers (e.g., only answers with a shorter length are considered). + max_question_len (`int`, *optional*): + The maximum length of the question after tokenization. It will be truncated if needed. + max_seq_len (`int`, *optional*): + The maximum length of the total sentence (context + question) in tokens of each chunk passed to the + model. The context will be split in several chunks (using doc_stride as overlap) if needed. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Can return less than top_k + answers if there are not enough options available within the context. + word_boxes (`List[Union[List[float], str`, *optional*): + A list of words and bounding boxes (normalized 0->1000). If provided, the inference will skip the OCR + step and use the provided bounding boxes instead. + Returns: + `List[DocumentQuestionAnsweringOutputElement]`: a list of [`DocumentQuestionAnsweringOutputElement`] items containing the predicted label, associated probability, word ids, and page number. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.document_question_answering(image="https://huggingface.co/spaces/impira/docquery/resolve/2359223c1837a7587402bda0f2643382a6eefeab/invoice.png", question="What is the invoice number?") + [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16)] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="document-question-answering", model=model_id) + inputs: Dict[str, Any] = {"question": question, "image": _b64_encode(image)} + request_parameters = provider_helper.prepare_request( + inputs=inputs, + parameters={ + "doc_stride": doc_stride, + "handle_impossible_answer": handle_impossible_answer, + "lang": lang, + "max_answer_len": max_answer_len, + "max_question_len": max_question_len, + "max_seq_len": max_seq_len, + "top_k": top_k, + "word_boxes": word_boxes, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return DocumentQuestionAnsweringOutputElement.parse_obj_as_list(response) + + async def feature_extraction( + self, + text: str, + *, + normalize: Optional[bool] = None, + prompt_name: Optional[str] = None, + truncate: Optional[bool] = None, + truncation_direction: Optional[Literal["Left", "Right"]] = None, + model: Optional[str] = None, + ) -> "np.ndarray": + """ + Generate embeddings for a given text. + + Args: + text (`str`): + The text to embed. + model (`str`, *optional*): + The model to use for the feature extraction task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended feature extraction model will be used. + Defaults to None. + normalize (`bool`, *optional*): + Whether to normalize the embeddings or not. + Only available on server powered by Text-Embedding-Inference. + prompt_name (`str`, *optional*): + The name of the prompt that should be used by for encoding. If not set, no prompt will be applied. + Must be a key in the `Sentence Transformers` configuration `prompts` dictionary. + For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",...}, + then the sentence "What is the capital of France?" will be encoded as "query: What is the capital of France?" + because the prompt text will be prepended before any text to encode. + truncate (`bool`, *optional*): + Whether to truncate the embeddings or not. + Only available on server powered by Text-Embedding-Inference. + truncation_direction (`Literal["Left", "Right"]`, *optional*): + Which side of the input should be truncated when `truncate=True` is passed. + + Returns: + `np.ndarray`: The embedding representing the input text as a float32 numpy array. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.feature_extraction("Hi, who are you?") + array([[ 2.424802 , 2.93384 , 1.1750331 , ..., 1.240499, -0.13776633, -0.7889173 ], + [-0.42943227, -0.6364878 , -1.693462 , ..., 0.41978157, -2.4336355 , 0.6162071 ], + ..., + [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="feature-extraction", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "normalize": normalize, + "prompt_name": prompt_name, + "truncate": truncate, + "truncation_direction": truncation_direction, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + np = _import_numpy() + return np.array(provider_helper.get_response(response), dtype="float32") + + async def fill_mask( + self, + text: str, + *, + model: Optional[str] = None, + targets: Optional[List[str]] = None, + top_k: Optional[int] = None, + ) -> List[FillMaskOutputElement]: + """ + Fill in a hole with a missing word (token to be precise). + + Args: + text (`str`): + a string to be filled from, must contain the [MASK] token (check model card for exact name of the mask). + model (`str`, *optional*): + The model to use for the fill mask task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended fill mask model will be used. + targets (`List[str`, *optional*): + When passed, the model will limit the scores to the passed targets instead of looking up in the whole + vocabulary. If the provided targets are not in the model vocab, they will be tokenized and the first + resulting token will be used (with a warning, and that might be slower). + top_k (`int`, *optional*): + When passed, overrides the number of predictions to return. + Returns: + `List[FillMaskOutputElement]`: a list of [`FillMaskOutputElement`] items containing the predicted label, associated + probability, token reference, and completed text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.fill_mask("The goal of life is .") + [ + FillMaskOutputElement(score=0.06897063553333282, token=11098, token_str=' happiness', sequence='The goal of life is happiness.'), + FillMaskOutputElement(score=0.06554922461509705, token=45075, token_str=' immortality', sequence='The goal of life is immortality.') + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="fill-mask", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={"targets": targets, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return FillMaskOutputElement.parse_obj_as_list(response) + + async def image_classification( + self, + image: ContentT, + *, + model: Optional[str] = None, + function_to_apply: Optional["ImageClassificationOutputTransform"] = None, + top_k: Optional[int] = None, + ) -> List[ImageClassificationOutputElement]: + """ + Perform image classification on the given image using the specified model. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to classify. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for image classification. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for image classification will be used. + function_to_apply (`"ImageClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + Returns: + `List[ImageClassificationOutputElement]`: a list of [`ImageClassificationOutputElement`] items containing the predicted label and associated probability. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.image_classification("https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg") + [ImageClassificationOutputElement(label='Blenheim spaniel', score=0.9779096841812134), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"function_to_apply": function_to_apply, "top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return ImageClassificationOutputElement.parse_obj_as_list(response) + + async def image_segmentation( + self, + image: ContentT, + *, + model: Optional[str] = None, + mask_threshold: Optional[float] = None, + overlap_mask_area_threshold: Optional[float] = None, + subtask: Optional["ImageSegmentationSubtask"] = None, + threshold: Optional[float] = None, + ) -> List[ImageSegmentationOutputElement]: + """ + Perform image segmentation on the given image using the specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to segment. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for image segmentation. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for image segmentation will be used. + mask_threshold (`float`, *optional*): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*): + Mask overlap threshold to eliminate small, disconnected segments. + subtask (`"ImageSegmentationSubtask"`, *optional*): + Segmentation task to be performed, depending on model capabilities. + threshold (`float`, *optional*): + Probability threshold to filter out predicted masks. + Returns: + `List[ImageSegmentationOutputElement]`: A list of [`ImageSegmentationOutputElement`] items containing the segmented masks and associated attributes. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.image_segmentation("cat.jpg") + [ImageSegmentationOutputElement(score=0.989008, label='LABEL_184', mask=), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-segmentation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "mask_threshold": mask_threshold, + "overlap_mask_area_threshold": overlap_mask_area_threshold, + "subtask": subtask, + "threshold": threshold, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + output = ImageSegmentationOutputElement.parse_obj_as_list(response) + for item in output: + item.mask = _b64_to_image(item.mask) # type: ignore [assignment] + return output + + async def image_to_image( + self, + image: ContentT, + prompt: Optional[str] = None, + *, + negative_prompt: Optional[str] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + model: Optional[str] = None, + target_size: Optional[ImageToImageTargetSize] = None, + **kwargs, + ) -> "Image": + """ + Perform image-to-image translation using a specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for translation. It can be raw bytes, an image file, or a URL to an online image. + prompt (`str`, *optional*): + The text prompt to guide the image generation. + negative_prompt (`str`, *optional*): + One prompt to guide what NOT to include in image generation. + num_inference_steps (`int`, *optional*): + For diffusion models. The number of denoising steps. More denoising steps usually lead to a higher + quality image at the expense of slower inference. + guidance_scale (`float`, *optional*): + For diffusion models. A higher guidance scale value encourages the model to generate images closely + linked to the text prompt at the expense of lower image quality. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + target_size (`ImageToImageTargetSize`, *optional*): + The size in pixel of the output image. + + Returns: + `Image`: The translated image. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> image = await client.image_to_image("cat.jpg", prompt="turn the cat into a tiger") + >>> image.save("tiger.jpg") + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-to-image", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "prompt": prompt, + "negative_prompt": negative_prompt, + "target_size": target_size, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + **kwargs, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return _bytes_to_image(response) + + async def image_to_text(self, image: ContentT, *, model: Optional[str] = None) -> ImageToTextOutput: + """ + Takes an input image and return text. + + Models can have very different outputs depending on your use case (image captioning, optical character recognition + (OCR), Pix2Struct, etc). Please have a look to the model card to learn more about a model's specificities. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image to caption. It can be raw bytes, an image file, or a URL to an online image.. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + [`ImageToTextOutput`]: The generated text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.image_to_text("cat.jpg") + 'a cat standing in a grassy field ' + >>> await client.image_to_text("https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg") + 'a dog laying on the grass next to a flower pot ' + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="image-to-text", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + output = ImageToTextOutput.parse_obj(response) + return output[0] if isinstance(output, list) else output + + async def object_detection( + self, image: ContentT, *, model: Optional[str] = None, threshold: Optional[float] = None + ) -> List[ObjectDetectionOutputElement]: + """ + Perform object detection on the given image using the specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The image to detect objects on. It can be raw bytes, an image file, or a URL to an online image. + model (`str`, *optional*): + The model to use for object detection. Can be a model ID hosted on the Hugging Face Hub or a URL to a + deployed Inference Endpoint. If not provided, the default recommended model for object detection (DETR) will be used. + threshold (`float`, *optional*): + The probability necessary to make a prediction. + Returns: + `List[ObjectDetectionOutputElement]`: A list of [`ObjectDetectionOutputElement`] items containing the bounding boxes and associated attributes. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + `ValueError`: + If the request output is not a List. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.object_detection("people.jpg") + [ObjectDetectionOutputElement(score=0.9486683011054993, label='person', box=ObjectDetectionBoundingBox(xmin=59, ymin=39, xmax=420, ymax=510)), ...] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="object-detection", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"threshold": threshold}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return ObjectDetectionOutputElement.parse_obj_as_list(response) + + async def question_answering( + self, + question: str, + context: str, + *, + model: Optional[str] = None, + align_to_words: Optional[bool] = None, + doc_stride: Optional[int] = None, + handle_impossible_answer: Optional[bool] = None, + max_answer_len: Optional[int] = None, + max_question_len: Optional[int] = None, + max_seq_len: Optional[int] = None, + top_k: Optional[int] = None, + ) -> Union[QuestionAnsweringOutputElement, List[QuestionAnsweringOutputElement]]: + """ + Retrieve the answer to a question from a given text. + + Args: + question (`str`): + Question to be answered. + context (`str`): + The context of the question. + model (`str`): + The model to use for the question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. + align_to_words (`bool`, *optional*): + Attempts to align the answer to real words. Improves quality on space separated languages. Might hurt + on non-space-separated languages (like Japanese or Chinese) + doc_stride (`int`, *optional*): + If the context is too long to fit with the question for the model, it will be split in several chunks + with some overlap. This argument controls the size of that overlap. + handle_impossible_answer (`bool`, *optional*): + Whether to accept impossible as an answer. + max_answer_len (`int`, *optional*): + The maximum length of predicted answers (e.g., only answers with a shorter length are considered). + max_question_len (`int`, *optional*): + The maximum length of the question after tokenization. It will be truncated if needed. + max_seq_len (`int`, *optional*): + The maximum length of the total sentence (context + question) in tokens of each chunk passed to the + model. The context will be split in several chunks (using docStride as overlap) if needed. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Note that we return less than + topk answers if there are not enough options available within the context. + + Returns: + Union[`QuestionAnsweringOutputElement`, List[`QuestionAnsweringOutputElement`]]: + When top_k is 1 or not provided, it returns a single `QuestionAnsweringOutputElement`. + When top_k is greater than 1, it returns a list of `QuestionAnsweringOutputElement`. + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.question_answering(question="What's my name?", context="My name is Clara and I live in Berkeley.") + QuestionAnsweringOutputElement(answer='Clara', end=16, score=0.9326565265655518, start=11) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={ + "align_to_words": align_to_words, + "doc_stride": doc_stride, + "handle_impossible_answer": handle_impossible_answer, + "max_answer_len": max_answer_len, + "max_question_len": max_question_len, + "max_seq_len": max_seq_len, + "top_k": top_k, + }, + extra_payload={"question": question, "context": context}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + # Parse the response as a single `QuestionAnsweringOutputElement` when top_k is 1 or not provided, or a list of `QuestionAnsweringOutputElement` to ensure backward compatibility. + output = QuestionAnsweringOutputElement.parse_obj(response) + return output + + async def sentence_similarity( + self, sentence: str, other_sentences: List[str], *, model: Optional[str] = None + ) -> List[float]: + """ + Compute the semantic similarity between a sentence and a list of other sentences by comparing their embeddings. + + Args: + sentence (`str`): + The main sentence to compare to others. + other_sentences (`List[str]`): + The list of sentences to compare to. + model (`str`, *optional*): + The model to use for the sentence similarity task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended sentence similarity model will be used. + Defaults to None. + + Returns: + `List[float]`: The embedding representing the input text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.sentence_similarity( + ... "Machine learning is so easy.", + ... other_sentences=[ + ... "Deep learning is so straightforward.", + ... "This is so difficult, like rocket science.", + ... "I can't believe how much I struggled with this.", + ... ], + ... ) + [0.7785726189613342, 0.45876261591911316, 0.2906220555305481] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="sentence-similarity", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs={"source_sentence": sentence, "sentences": other_sentences}, + parameters={}, + extra_payload={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return _bytes_to_list(response) + + async def summarization( + self, + text: str, + *, + model: Optional[str] = None, + clean_up_tokenization_spaces: Optional[bool] = None, + generate_parameters: Optional[Dict[str, Any]] = None, + truncation: Optional["SummarizationTruncationStrategy"] = None, + ) -> SummarizationOutput: + """ + Generate a summary of a given text using a specified model. + + Args: + text (`str`): + The input text to summarize. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended model for summarization will be used. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether to clean up the potential extra spaces in the text output. + generate_parameters (`Dict[str, Any]`, *optional*): + Additional parametrization of the text generation algorithm. + truncation (`"SummarizationTruncationStrategy"`, *optional*): + The truncation strategy to use. + Returns: + [`SummarizationOutput`]: The generated summary text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.summarization("The Eiffel tower...") + SummarizationOutput(generated_text="The Eiffel tower is one of the most famous landmarks in the world....") + ``` + """ + parameters = { + "clean_up_tokenization_spaces": clean_up_tokenization_spaces, + "generate_parameters": generate_parameters, + "truncation": truncation, + } + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="summarization", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters=parameters, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return SummarizationOutput.parse_obj_as_list(response)[0] + + async def table_question_answering( + self, + table: Dict[str, Any], + query: str, + *, + model: Optional[str] = None, + padding: Optional["Padding"] = None, + sequential: Optional[bool] = None, + truncation: Optional[bool] = None, + ) -> TableQuestionAnsweringOutputElement: + """ + Retrieve the answer to a question from information given in a table. + + Args: + table (`str`): + A table of data represented as a dict of lists where entries are headers and the lists are all the + values, all lists must have the same size. + query (`str`): + The query in plain text that you want to ask the table. + model (`str`): + The model to use for the table-question-answering task. Can be a model ID hosted on the Hugging Face + Hub or a URL to a deployed Inference Endpoint. + padding (`"Padding"`, *optional*): + Activates and controls padding. + sequential (`bool`, *optional*): + Whether to do inference sequentially or as a batch. Batching is faster, but models like SQA require the + inference to be done sequentially to extract relations within sequences, given their conversational + nature. + truncation (`bool`, *optional*): + Activates and controls truncation. + + Returns: + [`TableQuestionAnsweringOutputElement`]: a table question answering output containing the answer, coordinates, cells and the aggregator used. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> query = "How many stars does the transformers repository have?" + >>> table = {"Repository": ["Transformers", "Datasets", "Tokenizers"], "Stars": ["36542", "4512", "3934"]} + >>> await client.table_question_answering(table, query, model="google/tapas-base-finetuned-wtq") + TableQuestionAnsweringOutputElement(answer='36542', coordinates=[[0, 1]], cells=['36542'], aggregator='AVERAGE') + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="table-question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={"model": model, "padding": padding, "sequential": sequential, "truncation": truncation}, + extra_payload={"query": query, "table": table}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return TableQuestionAnsweringOutputElement.parse_obj_as_instance(response) + + async def tabular_classification(self, table: Dict[str, Any], *, model: Optional[str] = None) -> List[str]: + """ + Classifying a target category (a group) based on a set of attributes. + + Args: + table (`Dict[str, Any]`): + Set of attributes to classify. + model (`str`, *optional*): + The model to use for the tabular classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended tabular classification model will be used. + Defaults to None. + + Returns: + `List`: a list of labels, one per row in the initial table. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> table = { + ... "fixed_acidity": ["7.4", "7.8", "10.3"], + ... "volatile_acidity": ["0.7", "0.88", "0.32"], + ... "citric_acid": ["0", "0", "0.45"], + ... "residual_sugar": ["1.9", "2.6", "6.4"], + ... "chlorides": ["0.076", "0.098", "0.073"], + ... "free_sulfur_dioxide": ["11", "25", "5"], + ... "total_sulfur_dioxide": ["34", "67", "13"], + ... "density": ["0.9978", "0.9968", "0.9976"], + ... "pH": ["3.51", "3.2", "3.23"], + ... "sulphates": ["0.56", "0.68", "0.82"], + ... "alcohol": ["9.4", "9.8", "12.6"], + ... } + >>> await client.tabular_classification(table=table, model="julien-c/wine-quality") + ["5", "5", "5"] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="tabular-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + extra_payload={"table": table}, + parameters={}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return _bytes_to_list(response) + + async def tabular_regression(self, table: Dict[str, Any], *, model: Optional[str] = None) -> List[float]: + """ + Predicting a numerical target value given a set of attributes/features in a table. + + Args: + table (`Dict[str, Any]`): + Set of attributes stored in a table. The attributes used to predict the target can be both numerical and categorical. + model (`str`, *optional*): + The model to use for the tabular regression task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended tabular regression model will be used. + Defaults to None. + + Returns: + `List`: a list of predicted numerical target values. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> table = { + ... "Height": ["11.52", "12.48", "12.3778"], + ... "Length1": ["23.2", "24", "23.9"], + ... "Length2": ["25.4", "26.3", "26.5"], + ... "Length3": ["30", "31.2", "31.1"], + ... "Species": ["Bream", "Bream", "Bream"], + ... "Width": ["4.02", "4.3056", "4.6961"], + ... } + >>> await client.tabular_regression(table, model="scikit-learn/Fish-Weight") + [110, 120, 130] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="tabular-regression", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=None, + parameters={}, + extra_payload={"table": table}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return _bytes_to_list(response) + + async def text_classification( + self, + text: str, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + function_to_apply: Optional["TextClassificationOutputTransform"] = None, + ) -> List[TextClassificationOutputElement]: + """ + Perform text classification (e.g. sentiment-analysis) on the given text. + + Args: + text (`str`): + A string to be classified. + model (`str`, *optional*): + The model to use for the text classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended text classification model will be used. + Defaults to None. + top_k (`int`, *optional*): + When specified, limits the output to the top K most probable classes. + function_to_apply (`"TextClassificationOutputTransform"`, *optional*): + The function to apply to the model outputs in order to retrieve the scores. + + Returns: + `List[TextClassificationOutputElement]`: a list of [`TextClassificationOutputElement`] items containing the predicted label and associated probability. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.text_classification("I like you") + [ + TextClassificationOutputElement(label='POSITIVE', score=0.9998695850372314), + TextClassificationOutputElement(label='NEGATIVE', score=0.0001304351753788069), + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "function_to_apply": function_to_apply, + "top_k": top_k, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return TextClassificationOutputElement.parse_obj_as_list(response)[0] # type: ignore [return-value] + + @overload + async def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[False] = ..., + stream: Literal[False] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> str: ... + + @overload + async def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: Literal[False] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> TextGenerationOutput: ... + + @overload + async def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[False] = ..., + stream: Literal[True] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> AsyncIterable[str]: ... + + @overload + async def text_generation( # type: ignore + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: Literal[True] = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> AsyncIterable[TextGenerationStreamOutput]: ... + + @overload + async def text_generation( + self, + prompt: str, + *, + details: Literal[True] = ..., + stream: bool = ..., + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Union[TextGenerationOutput, AsyncIterable[TextGenerationStreamOutput]]: ... + + async def text_generation( + self, + prompt: str, + *, + details: bool = False, + stream: bool = False, + model: Optional[str] = None, + # Parameters from `TextGenerationInputGenerateParameters` (maintained manually) + adapter_id: Optional[str] = None, + best_of: Optional[int] = None, + decoder_input_details: Optional[bool] = None, + do_sample: Optional[bool] = False, # Manual default value + frequency_penalty: Optional[float] = None, + grammar: Optional[TextGenerationInputGrammarType] = None, + max_new_tokens: Optional[int] = None, + repetition_penalty: Optional[float] = None, + return_full_text: Optional[bool] = False, # Manual default value + seed: Optional[int] = None, + stop: Optional[List[str]] = None, + stop_sequences: Optional[List[str]] = None, # Deprecated, use `stop` instead + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_n_tokens: Optional[int] = None, + top_p: Optional[float] = None, + truncate: Optional[int] = None, + typical_p: Optional[float] = None, + watermark: Optional[bool] = None, + ) -> Union[str, TextGenerationOutput, AsyncIterable[str], AsyncIterable[TextGenerationStreamOutput]]: + """ + Given a prompt, generate the following text. + + + + If you want to generate a response from chat messages, you should use the [`InferenceClient.chat_completion`] method. + It accepts a list of messages instead of a single text prompt and handles the chat templating for you. + + + + Args: + prompt (`str`): + Input text. + details (`bool`, *optional*): + By default, text_generation returns a string. Pass `details=True` if you want a detailed output (tokens, + probabilities, seed, finish reason, etc.). Only available for models running on with the + `text-generation-inference` backend. + stream (`bool`, *optional*): + By default, text_generation returns the full generated text. Pass `stream=True` if you want a stream of + tokens to be returned. Only available for models running on with the `text-generation-inference` + backend. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + adapter_id (`str`, *optional*): + Lora adapter id. + best_of (`int`, *optional*): + Generate best_of sequences and return the one if the highest token logprobs. + decoder_input_details (`bool`, *optional*): + Return the decoder input token logprobs and ids. You must set `details=True` as well for it to be taken + into account. Defaults to `False`. + do_sample (`bool`, *optional*): + Activate logits sampling + frequency_penalty (`float`, *optional*): + Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in + the text so far, decreasing the model's likelihood to repeat the same line verbatim. + grammar ([`TextGenerationInputGrammarType`], *optional*): + Grammar constraints. Can be either a JSONSchema or a regex. + max_new_tokens (`int`, *optional*): + Maximum number of generated tokens. Defaults to 100. + repetition_penalty (`float`, *optional*): + The parameter for repetition penalty. 1.0 means no penalty. See [this + paper](https://arxiv.org/pdf/1909.05858.pdf) for more details. + return_full_text (`bool`, *optional*): + Whether to prepend the prompt to the generated text + seed (`int`, *optional*): + Random sampling seed + stop (`List[str]`, *optional*): + Stop generating tokens if a member of `stop` is generated. + stop_sequences (`List[str]`, *optional*): + Deprecated argument. Use `stop` instead. + temperature (`float`, *optional*): + The value used to module the logits distribution. + top_n_tokens (`int`, *optional*): + Return information about the `top_n_tokens` most likely tokens at each generation step, instead of + just the sampled token. + top_k (`int`, *optional`): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + top_p (`float`, *optional`): + If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or + higher are kept for generation. + truncate (`int`, *optional`): + Truncate inputs tokens to the given size. + typical_p (`float`, *optional`): + Typical Decoding mass + See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information + watermark (`bool`, *optional`): + Watermarking with [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226) + + Returns: + `Union[str, TextGenerationOutput, Iterable[str], Iterable[TextGenerationStreamOutput]]`: + Generated text returned from the server: + - if `stream=False` and `details=False`, the generated text is returned as a `str` (default) + - if `stream=True` and `details=False`, the generated text is returned token by token as a `Iterable[str]` + - if `stream=False` and `details=True`, the generated text is returned with more details as a [`~huggingface_hub.TextGenerationOutput`] + - if `details=True` and `stream=True`, the generated text is returned token by token as a iterable of [`~huggingface_hub.TextGenerationStreamOutput`] + + Raises: + `ValidationError`: + If input values are not valid. No HTTP call is made to the server. + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + + # Case 1: generate text + >>> await client.text_generation("The huggingface_hub library is ", max_new_tokens=12) + '100% open source and built to be easy to use.' + + # Case 2: iterate over the generated tokens. Useful for large generation. + >>> async for token in await client.text_generation("The huggingface_hub library is ", max_new_tokens=12, stream=True): + ... print(token) + 100 + % + open + source + and + built + to + be + easy + to + use + . + + # Case 3: get more details about the generation process. + >>> await client.text_generation("The huggingface_hub library is ", max_new_tokens=12, details=True) + TextGenerationOutput( + generated_text='100% open source and built to be easy to use.', + details=TextGenerationDetails( + finish_reason='length', + generated_tokens=12, + seed=None, + prefill=[ + TextGenerationPrefillOutputToken(id=487, text='The', logprob=None), + TextGenerationPrefillOutputToken(id=53789, text=' hugging', logprob=-13.171875), + (...) + TextGenerationPrefillOutputToken(id=204, text=' ', logprob=-7.0390625) + ], + tokens=[ + TokenElement(id=1425, text='100', logprob=-1.0175781, special=False), + TokenElement(id=16, text='%', logprob=-0.0463562, special=False), + (...) + TokenElement(id=25, text='.', logprob=-0.5703125, special=False) + ], + best_of_sequences=None + ) + ) + + # Case 4: iterate over the generated tokens with more details. + # Last object is more complete, containing the full generated text and the finish reason. + >>> async for details in await client.text_generation("The huggingface_hub library is ", max_new_tokens=12, details=True, stream=True): + ... print(details) + ... + TextGenerationStreamOutput(token=TokenElement(id=1425, text='100', logprob=-1.0175781, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=16, text='%', logprob=-0.0463562, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=1314, text=' open', logprob=-1.3359375, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=3178, text=' source', logprob=-0.28100586, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=273, text=' and', logprob=-0.5961914, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=3426, text=' built', logprob=-1.9423828, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=271, text=' to', logprob=-1.4121094, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=314, text=' be', logprob=-1.5224609, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=1833, text=' easy', logprob=-2.1132812, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=271, text=' to', logprob=-0.08520508, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement(id=745, text=' use', logprob=-0.39453125, special=False), generated_text=None, details=None) + TextGenerationStreamOutput(token=TokenElement( + id=25, + text='.', + logprob=-0.5703125, + special=False), + generated_text='100% open source and built to be easy to use.', + details=TextGenerationStreamOutputStreamDetails(finish_reason='length', generated_tokens=12, seed=None) + ) + + # Case 5: generate constrained output using grammar + >>> response = await client.text_generation( + ... prompt="I saw a puppy a cat and a raccoon during my bike ride in the park", + ... model="HuggingFaceH4/zephyr-orpo-141b-A35b-v0.1", + ... max_new_tokens=100, + ... repetition_penalty=1.3, + ... grammar={ + ... "type": "json", + ... "value": { + ... "properties": { + ... "location": {"type": "string"}, + ... "activity": {"type": "string"}, + ... "animals_seen": {"type": "integer", "minimum": 1, "maximum": 5}, + ... "animals": {"type": "array", "items": {"type": "string"}}, + ... }, + ... "required": ["location", "activity", "animals_seen", "animals"], + ... }, + ... }, + ... ) + >>> json.loads(response) + { + "activity": "bike riding", + "animals": ["puppy", "cat", "raccoon"], + "animals_seen": 3, + "location": "park" + } + ``` + """ + if decoder_input_details and not details: + warnings.warn( + "`decoder_input_details=True` has been passed to the server but `details=False` is set meaning that" + " the output from the server will be truncated." + ) + decoder_input_details = False + + if stop_sequences is not None: + warnings.warn( + "`stop_sequences` is a deprecated argument for `text_generation` task" + " and will be removed in version '0.28.0'. Use `stop` instead.", + FutureWarning, + ) + if stop is None: + stop = stop_sequences # use deprecated arg if provided + + # Build payload + parameters = { + "adapter_id": adapter_id, + "best_of": best_of, + "decoder_input_details": decoder_input_details, + "details": details, + "do_sample": do_sample, + "frequency_penalty": frequency_penalty, + "grammar": grammar, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + "return_full_text": return_full_text, + "seed": seed, + "stop": stop if stop is not None else [], + "temperature": temperature, + "top_k": top_k, + "top_n_tokens": top_n_tokens, + "top_p": top_p, + "truncate": truncate, + "typical_p": typical_p, + "watermark": watermark, + } + + # Remove some parameters if not a TGI server + unsupported_kwargs = _get_unsupported_text_generation_kwargs(model) + if len(unsupported_kwargs) > 0: + # The server does not support some parameters + # => means it is not a TGI server + # => remove unsupported parameters and warn the user + + ignored_parameters = [] + for key in unsupported_kwargs: + if parameters.get(key): + ignored_parameters.append(key) + parameters.pop(key, None) + if len(ignored_parameters) > 0: + warnings.warn( + "API endpoint/model for text-generation is not served via TGI. Ignoring following parameters:" + f" {', '.join(ignored_parameters)}.", + UserWarning, + ) + if details: + warnings.warn( + "API endpoint/model for text-generation is not served via TGI. Parameter `details=True` will" + " be ignored meaning only the generated text will be returned.", + UserWarning, + ) + details = False + if stream: + raise ValueError( + "API endpoint/model for text-generation is not served via TGI. Cannot return output as a stream." + " Please pass `stream=False` as input." + ) + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-generation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters=parameters, + extra_payload={"stream": stream}, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + + # Handle errors separately for more precise error messages + try: + bytes_output = await self._inner_post(request_parameters, stream=stream) + except _import_aiohttp().ClientResponseError as e: + match = MODEL_KWARGS_NOT_USED_REGEX.search(e.response_error_payload["error"]) + if e.status == 400 and match: + unused_params = [kwarg.strip("' ") for kwarg in match.group(1).split(",")] + _set_unsupported_text_generation_kwargs(model, unused_params) + return await self.text_generation( # type: ignore + prompt=prompt, + details=details, + stream=stream, + model=model_id, + adapter_id=adapter_id, + best_of=best_of, + decoder_input_details=decoder_input_details, + do_sample=do_sample, + frequency_penalty=frequency_penalty, + grammar=grammar, + max_new_tokens=max_new_tokens, + repetition_penalty=repetition_penalty, + return_full_text=return_full_text, + seed=seed, + stop=stop, + temperature=temperature, + top_k=top_k, + top_n_tokens=top_n_tokens, + top_p=top_p, + truncate=truncate, + typical_p=typical_p, + watermark=watermark, + ) + raise_text_generation_error(e) + + # Parse output + if stream: + return _async_stream_text_generation_response(bytes_output, details) # type: ignore + + data = _bytes_to_dict(bytes_output) # type: ignore[arg-type] + + # Data can be a single element (dict) or an iterable of dicts where we select the first element of. + if isinstance(data, list): + data = data[0] + response = provider_helper.get_response(data, request_parameters) + return TextGenerationOutput.parse_obj_as_instance(response) if details else response["generated_text"] + + async def text_to_image( + self, + prompt: str, + *, + negative_prompt: Optional[str] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + model: Optional[str] = None, + scheduler: Optional[str] = None, + seed: Optional[int] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> "Image": + """ + Generate an image based on a given text using a specified model. + + + + You must have `PIL` installed if you want to work with images (`pip install Pillow`). + + + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + prompt (`str`): + The prompt to generate an image from. + negative_prompt (`str`, *optional*): + One prompt to guide what NOT to include in image generation. + height (`int`, *optional*): + The height in pixels of the output image + width (`int`, *optional*): + The width in pixels of the output image + num_inference_steps (`int`, *optional*): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*): + A higher guidance scale value encourages the model to generate images closely linked to the text + prompt, but values too high may cause saturation and other artifacts. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-image model will be used. + Defaults to None. + scheduler (`str`, *optional*): + Override the scheduler with a compatible one. + seed (`int`, *optional*): + Seed for the random number generator. + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + + Returns: + `Image`: The generated image. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + + >>> image = await client.text_to_image("An astronaut riding a horse on the moon.") + >>> image.save("astronaut.png") + + >>> image = await client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... negative_prompt="low resolution, blurry", + ... model="stabilityai/stable-diffusion-2-1", + ... ) + >>> image.save("better_astronaut.png") + ``` + Example using a third-party provider directly. Usage will be billed on your fal.ai account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="fal-ai", # Use fal.ai provider + ... api_key="fal-ai-api-key", # Pass your fal.ai API key + ... ) + >>> image = client.text_to_image( + ... "A majestic lion in a fantasy forest", + ... model="black-forest-labs/FLUX.1-schnell", + ... ) + >>> image.save("lion.png") + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> image = client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... model="black-forest-labs/FLUX.1-dev", + ... ) + >>> image.save("astronaut.png") + ``` + + Example using Replicate provider with extra parameters + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> image = client.text_to_image( + ... "An astronaut riding a horse on the moon.", + ... model="black-forest-labs/FLUX.1-schnell", + ... extra_body={"output_quality": 100}, + ... ) + >>> image.save("astronaut.png") + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-image", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters={ + "negative_prompt": negative_prompt, + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "scheduler": scheduler, + "seed": seed, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + response = provider_helper.get_response(response) + return _bytes_to_image(response) + + async def text_to_video( + self, + prompt: str, + *, + model: Optional[str] = None, + guidance_scale: Optional[float] = None, + negative_prompt: Optional[List[str]] = None, + num_frames: Optional[float] = None, + num_inference_steps: Optional[int] = None, + seed: Optional[int] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> bytes: + """ + Generate a video based on a given text. + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + prompt (`str`): + The prompt to generate a video from. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-video model will be used. + Defaults to None. + guidance_scale (`float`, *optional*): + A higher guidance scale value encourages the model to generate videos closely linked to the text + prompt, but values too high may cause saturation and other artifacts. + negative_prompt (`List[str]`, *optional*): + One or several prompt to guide what NOT to include in video generation. + num_frames (`float`, *optional*): + The num_frames parameter determines how many video frames are generated. + num_inference_steps (`int`, *optional*): + The number of denoising steps. More denoising steps usually lead to a higher quality video at the + expense of slower inference. + seed (`int`, *optional*): + Seed for the random number generator. + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + + Returns: + `bytes`: The generated video. + + Example: + + Example using a third-party provider directly. Usage will be billed on your fal.ai account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="fal-ai", # Using fal.ai provider + ... api_key="fal-ai-api-key", # Pass your fal.ai API key + ... ) + >>> video = client.text_to_video( + ... "A majestic lion running in a fantasy forest", + ... model="tencent/HunyuanVideo", + ... ) + >>> with open("lion.mp4", "wb") as file: + ... file.write(video) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Using replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> video = client.text_to_video( + ... "A cat running in a park", + ... model="genmo/mochi-1-preview", + ... ) + >>> with open("cat.mp4", "wb") as file: + ... file.write(video) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-video", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=prompt, + parameters={ + "guidance_scale": guidance_scale, + "negative_prompt": negative_prompt, + "num_frames": num_frames, + "num_inference_steps": num_inference_steps, + "seed": seed, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + response = provider_helper.get_response(response, request_parameters) + return response + + async def text_to_speech( + self, + text: str, + *, + model: Optional[str] = None, + do_sample: Optional[bool] = None, + early_stopping: Optional[Union[bool, "TextToSpeechEarlyStoppingEnum"]] = None, + epsilon_cutoff: Optional[float] = None, + eta_cutoff: Optional[float] = None, + max_length: Optional[int] = None, + max_new_tokens: Optional[int] = None, + min_length: Optional[int] = None, + min_new_tokens: Optional[int] = None, + num_beam_groups: Optional[int] = None, + num_beams: Optional[int] = None, + penalty_alpha: Optional[float] = None, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + use_cache: Optional[bool] = None, + extra_body: Optional[Dict[str, Any]] = None, + ) -> bytes: + """ + Synthesize an audio of a voice pronouncing a given text. + + + You can pass provider-specific parameters to the model by using the `extra_body` argument. + + + Args: + text (`str`): + The text to synthesize. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. If not provided, the default recommended text-to-speech model will be used. + Defaults to None. + do_sample (`bool`, *optional*): + Whether to use sampling instead of greedy decoding when generating new tokens. + early_stopping (`Union[bool, "TextToSpeechEarlyStoppingEnum"]`, *optional*): + Controls the stopping condition for beam-based methods. + epsilon_cutoff (`float`, *optional*): + If set to float strictly between 0 and 1, only tokens with a conditional probability greater than + epsilon_cutoff will be sampled. In the paper, suggested values range from 3e-4 to 9e-4, depending on + the size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://hf.co/papers/2210.15191) for more details. + eta_cutoff (`float`, *optional*): + Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to float strictly + between 0 and 1, a token is only considered if it is greater than either eta_cutoff or sqrt(eta_cutoff) + * exp(-entropy(softmax(next_token_logits))). The latter term is intuitively the expected next token + probability, scaled by sqrt(eta_cutoff). In the paper, suggested values range from 3e-4 to 2e-3, + depending on the size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://hf.co/papers/2210.15191) for more details. + max_length (`int`, *optional*): + The maximum length (in tokens) of the generated text, including the input. + max_new_tokens (`int`, *optional*): + The maximum number of tokens to generate. Takes precedence over max_length. + min_length (`int`, *optional*): + The minimum length (in tokens) of the generated text, including the input. + min_new_tokens (`int`, *optional*): + The minimum number of tokens to generate. Takes precedence over min_length. + num_beam_groups (`int`, *optional*): + Number of groups to divide num_beams into in order to ensure diversity among different groups of beams. + See [this paper](https://hf.co/papers/1610.02424) for more details. + num_beams (`int`, *optional*): + Number of beams to use for beam search. + penalty_alpha (`float`, *optional*): + The value balances the model confidence and the degeneration penalty in contrastive search decoding. + temperature (`float`, *optional*): + The value used to modulate the next token probabilities. + top_k (`int`, *optional*): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + top_p (`float`, *optional*): + If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to + top_p or higher are kept for generation. + typical_p (`float`, *optional*): + Local typicality measures how similar the conditional probability of predicting a target token next is + to the expected conditional probability of predicting a random token next, given the partial text + already generated. If set to float < 1, the smallest set of the most locally typical tokens with + probabilities that add up to typical_p or higher are kept for generation. See [this + paper](https://hf.co/papers/2202.00666) for more details. + use_cache (`bool`, *optional*): + Whether the model should use the past last key/values attentions to speed up decoding + extra_body (`Dict[str, Any]`, *optional*): + Additional provider-specific parameters to pass to the model. Refer to the provider's documentation + for supported parameters. + Returns: + `bytes`: The generated audio. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from pathlib import Path + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + + >>> audio = await client.text_to_speech("Hello world") + >>> Path("hello_world.flac").write_bytes(audio) + ``` + + Example using a third-party provider directly. Usage will be billed on your Replicate account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", + ... api_key="your-replicate-api-key", # Pass your Replicate API key directly + ... ) + >>> audio = client.text_to_speech( + ... text="Hello world", + ... model="OuteAI/OuteTTS-0.3-500M", + ... ) + >>> Path("hello_world.flac").write_bytes(audio) + ``` + + Example using a third-party provider through Hugging Face Routing. Usage will be billed on your Hugging Face account. + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", + ... api_key="hf_...", # Pass your HF token + ... ) + >>> audio =client.text_to_speech( + ... text="Hello world", + ... model="OuteAI/OuteTTS-0.3-500M", + ... ) + >>> Path("hello_world.flac").write_bytes(audio) + ``` + Example using Replicate provider with extra parameters + ```py + >>> from huggingface_hub import InferenceClient + >>> client = InferenceClient( + ... provider="replicate", # Use replicate provider + ... api_key="hf_...", # Pass your HF token + ... ) + >>> audio = client.text_to_speech( + ... "Hello, my name is Kororo, an awesome text-to-speech model.", + ... model="hexgrad/Kokoro-82M", + ... extra_body={"voice": "af_nicole"}, + ... ) + >>> Path("hello.flac").write_bytes(audio) + ``` + + Example music-gen using "YuE-s1-7B-anneal-en-cot" on fal.ai + ```py + >>> from huggingface_hub import InferenceClient + >>> lyrics = ''' + ... [verse] + ... In the town where I was born + ... Lived a man who sailed to sea + ... And he told us of his life + ... In the land of submarines + ... So we sailed on to the sun + ... 'Til we found a sea of green + ... And we lived beneath the waves + ... In our yellow submarine + + ... [chorus] + ... We all live in a yellow submarine + ... Yellow submarine, yellow submarine + ... We all live in a yellow submarine + ... Yellow submarine, yellow submarine + ... ''' + >>> genres = "pavarotti-style tenor voice" + >>> client = InferenceClient( + ... provider="fal-ai", + ... model="m-a-p/YuE-s1-7B-anneal-en-cot", + ... api_key=..., + ... ) + >>> audio = client.text_to_speech(lyrics, extra_body={"genres": genres}) + >>> with open("output.mp3", "wb") as f: + ... f.write(audio) + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="text-to-speech", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "do_sample": do_sample, + "early_stopping": early_stopping, + "epsilon_cutoff": epsilon_cutoff, + "eta_cutoff": eta_cutoff, + "max_length": max_length, + "max_new_tokens": max_new_tokens, + "min_length": min_length, + "min_new_tokens": min_new_tokens, + "num_beam_groups": num_beam_groups, + "num_beams": num_beams, + "penalty_alpha": penalty_alpha, + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + "typical_p": typical_p, + "use_cache": use_cache, + **(extra_body or {}), + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + response = provider_helper.get_response(response) + return response + + async def token_classification( + self, + text: str, + *, + model: Optional[str] = None, + aggregation_strategy: Optional["TokenClassificationAggregationStrategy"] = None, + ignore_labels: Optional[List[str]] = None, + stride: Optional[int] = None, + ) -> List[TokenClassificationOutputElement]: + """ + Perform token classification on the given text. + Usually used for sentence parsing, either grammatical, or Named Entity Recognition (NER) to understand keywords contained within text. + + Args: + text (`str`): + A string to be classified. + model (`str`, *optional*): + The model to use for the token classification task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended token classification model will be used. + Defaults to None. + aggregation_strategy (`"TokenClassificationAggregationStrategy"`, *optional*): + The strategy used to fuse tokens based on model predictions + ignore_labels (`List[str`, *optional*): + A list of labels to ignore + stride (`int`, *optional*): + The number of overlapping tokens between chunks when splitting the input text. + + Returns: + `List[TokenClassificationOutputElement]`: List of [`TokenClassificationOutputElement`] items containing the entity group, confidence score, word, start and end index. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.token_classification("My name is Sarah Jessica Parker but you can call me Jessica") + [ + TokenClassificationOutputElement( + entity_group='PER', + score=0.9971321225166321, + word='Sarah Jessica Parker', + start=11, + end=31, + ), + TokenClassificationOutputElement( + entity_group='PER', + score=0.9773476123809814, + word='Jessica', + start=52, + end=59, + ) + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="token-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "aggregation_strategy": aggregation_strategy, + "ignore_labels": ignore_labels, + "stride": stride, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return TokenClassificationOutputElement.parse_obj_as_list(response) + + async def translation( + self, + text: str, + *, + model: Optional[str] = None, + src_lang: Optional[str] = None, + tgt_lang: Optional[str] = None, + clean_up_tokenization_spaces: Optional[bool] = None, + truncation: Optional["TranslationTruncationStrategy"] = None, + generate_parameters: Optional[Dict[str, Any]] = None, + ) -> TranslationOutput: + """ + Convert text from one language to another. + + Check out https://huggingface.co/tasks/translation for more information on how to choose the best model for + your specific use case. Source and target languages usually depend on the model. + However, it is possible to specify source and target languages for certain models. If you are working with one of these models, + you can use `src_lang` and `tgt_lang` arguments to pass the relevant information. + + Args: + text (`str`): + A string to be translated. + model (`str`, *optional*): + The model to use for the translation task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended translation model will be used. + Defaults to None. + src_lang (`str`, *optional*): + The source language of the text. Required for models that can translate from multiple languages. + tgt_lang (`str`, *optional*): + Target language to translate to. Required for models that can translate to multiple languages. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether to clean up the potential extra spaces in the text output. + truncation (`"TranslationTruncationStrategy"`, *optional*): + The truncation strategy to use. + generate_parameters (`Dict[str, Any]`, *optional*): + Additional parametrization of the text generation algorithm. + + Returns: + [`TranslationOutput`]: The generated translated text. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + `ValueError`: + If only one of the `src_lang` and `tgt_lang` arguments are provided. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.translation("My name is Wolfgang and I live in Berlin") + 'Mein Name ist Wolfgang und ich lebe in Berlin.' + >>> await client.translation("My name is Wolfgang and I live in Berlin", model="Helsinki-NLP/opus-mt-en-fr") + TranslationOutput(translation_text='Je m'appelle Wolfgang et je vis à Berlin.') + ``` + + Specifying languages: + ```py + >>> client.translation("My name is Sarah Jessica Parker but you can call me Jessica", model="facebook/mbart-large-50-many-to-many-mmt", src_lang="en_XX", tgt_lang="fr_XX") + "Mon nom est Sarah Jessica Parker mais vous pouvez m'appeler Jessica" + ``` + """ + # Throw error if only one of `src_lang` and `tgt_lang` was given + if src_lang is not None and tgt_lang is None: + raise ValueError("You cannot specify `src_lang` without specifying `tgt_lang`.") + + if src_lang is None and tgt_lang is not None: + raise ValueError("You cannot specify `tgt_lang` without specifying `src_lang`.") + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="translation", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "src_lang": src_lang, + "tgt_lang": tgt_lang, + "clean_up_tokenization_spaces": clean_up_tokenization_spaces, + "truncation": truncation, + "generate_parameters": generate_parameters, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return TranslationOutput.parse_obj_as_list(response)[0] + + async def visual_question_answering( + self, + image: ContentT, + question: str, + *, + model: Optional[str] = None, + top_k: Optional[int] = None, + ) -> List[VisualQuestionAnsweringOutputElement]: + """ + Answering open-ended questions based on an image. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image for the context. It can be raw bytes, an image file, or a URL to an online image. + question (`str`): + Question to be answered. + model (`str`, *optional*): + The model to use for the visual question answering task. Can be a model ID hosted on the Hugging Face Hub or a URL to + a deployed Inference Endpoint. If not provided, the default recommended visual question answering model will be used. + Defaults to None. + top_k (`int`, *optional*): + The number of answers to return (will be chosen by order of likelihood). Note that we return less than + topk answers if there are not enough options available within the context. + Returns: + `List[VisualQuestionAnsweringOutputElement]`: a list of [`VisualQuestionAnsweringOutputElement`] items containing the predicted label and associated probability. + + Raises: + `InferenceTimeoutError`: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.visual_question_answering( + ... image="https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg", + ... question="What is the animal doing?" + ... ) + [ + VisualQuestionAnsweringOutputElement(score=0.778609573841095, answer='laying down'), + VisualQuestionAnsweringOutputElement(score=0.6957435607910156, answer='sitting'), + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="visual-question-answering", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={"top_k": top_k}, + headers=self.headers, + model=model_id, + api_key=self.token, + extra_payload={"question": question, "image": _b64_encode(image)}, + ) + response = await self._inner_post(request_parameters) + return VisualQuestionAnsweringOutputElement.parse_obj_as_list(response) + + async def zero_shot_classification( + self, + text: str, + candidate_labels: List[str], + *, + multi_label: Optional[bool] = False, + hypothesis_template: Optional[str] = None, + model: Optional[str] = None, + ) -> List[ZeroShotClassificationOutputElement]: + """ + Provide as input a text and a set of candidate labels to classify the input text. + + Args: + text (`str`): + The input text to classify. + candidate_labels (`List[str]`): + The set of possible class labels to classify the text into. + labels (`List[str]`, *optional*): + (deprecated) List of strings. Each string is the verbalization of a possible label for the input text. + multi_label (`bool`, *optional*): + Whether multiple candidate labels can be true. If false, the scores are normalized such that the sum of + the label likelihoods for each sequence is 1. If true, the labels are considered independent and + probabilities are normalized for each candidate. + hypothesis_template (`str`, *optional*): + The sentence used in conjunction with `candidate_labels` to attempt the text classification by + replacing the placeholder with the candidate labels. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. If not provided, the default recommended zero-shot classification model will be used. + + + Returns: + `List[ZeroShotClassificationOutputElement]`: List of [`ZeroShotClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example with `multi_label=False`: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> text = ( + ... "A new model offers an explanation for how the Galilean satellites formed around the solar system's" + ... "largest world. Konstantin Batygin did not set out to solve one of the solar system's most puzzling" + ... " mysteries when he went for a run up a hill in Nice, France." + ... ) + >>> labels = ["space & cosmos", "scientific discovery", "microbiology", "robots", "archeology"] + >>> await client.zero_shot_classification(text, labels) + [ + ZeroShotClassificationOutputElement(label='scientific discovery', score=0.7961668968200684), + ZeroShotClassificationOutputElement(label='space & cosmos', score=0.18570658564567566), + ZeroShotClassificationOutputElement(label='microbiology', score=0.00730885099619627), + ZeroShotClassificationOutputElement(label='archeology', score=0.006258360575884581), + ZeroShotClassificationOutputElement(label='robots', score=0.004559356719255447), + ] + >>> await client.zero_shot_classification(text, labels, multi_label=True) + [ + ZeroShotClassificationOutputElement(label='scientific discovery', score=0.9829297661781311), + ZeroShotClassificationOutputElement(label='space & cosmos', score=0.755190908908844), + ZeroShotClassificationOutputElement(label='microbiology', score=0.0005462635890580714), + ZeroShotClassificationOutputElement(label='archeology', score=0.00047131875180639327), + ZeroShotClassificationOutputElement(label='robots', score=0.00030448526376858354), + ] + ``` + + Example with `multi_label=True` and a custom `hypothesis_template`: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.zero_shot_classification( + ... text="I really like our dinner and I'm very happy. I don't like the weather though.", + ... labels=["positive", "negative", "pessimistic", "optimistic"], + ... multi_label=True, + ... hypothesis_template="This text is {} towards the weather" + ... ) + [ + ZeroShotClassificationOutputElement(label='negative', score=0.9231801629066467), + ZeroShotClassificationOutputElement(label='pessimistic', score=0.8760990500450134), + ZeroShotClassificationOutputElement(label='optimistic', score=0.0008674879791215062), + ZeroShotClassificationOutputElement(label='positive', score=0.0005250611575320363) + ] + ``` + """ + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="zero-shot-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=text, + parameters={ + "candidate_labels": candidate_labels, + "multi_label": multi_label, + "hypothesis_template": hypothesis_template, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + output = _bytes_to_dict(response) + return [ + ZeroShotClassificationOutputElement.parse_obj_as_instance({"label": label, "score": score}) + for label, score in zip(output["labels"], output["scores"]) + ] + + async def zero_shot_image_classification( + self, + image: ContentT, + candidate_labels: List[str], + *, + model: Optional[str] = None, + hypothesis_template: Optional[str] = None, + # deprecated argument + labels: List[str] = None, # type: ignore + ) -> List[ZeroShotImageClassificationOutputElement]: + """ + Provide input image and text labels to predict text labels for the image. + + Args: + image (`Union[str, Path, bytes, BinaryIO]`): + The input image to caption. It can be raw bytes, an image file, or a URL to an online image. + candidate_labels (`List[str]`): + The candidate labels for this image + labels (`List[str]`, *optional*): + (deprecated) List of string possible labels. There must be at least 2 labels. + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. If not provided, the default recommended zero-shot image classification model will be used. + hypothesis_template (`str`, *optional*): + The sentence used in conjunction with `candidate_labels` to attempt the image classification by + replacing the placeholder with the candidate labels. + + Returns: + `List[ZeroShotImageClassificationOutputElement]`: List of [`ZeroShotImageClassificationOutputElement`] items containing the predicted labels and their confidence. + + Raises: + [`InferenceTimeoutError`]: + If the model is unavailable or the request times out. + `aiohttp.ClientResponseError`: + If the request fails with an HTTP error status code other than HTTP 503. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + + >>> await client.zero_shot_image_classification( + ... "https://upload.wikimedia.org/wikipedia/commons/thumb/4/43/Cute_dog.jpg/320px-Cute_dog.jpg", + ... labels=["dog", "cat", "horse"], + ... ) + [ZeroShotImageClassificationOutputElement(label='dog', score=0.956),...] + ``` + """ + # Raise ValueError if input is less than 2 labels + if len(candidate_labels) < 2: + raise ValueError("You must specify at least 2 classes to compare.") + + model_id = model or self.model + provider_helper = get_provider_helper(self.provider, task="zero-shot-image-classification", model=model_id) + request_parameters = provider_helper.prepare_request( + inputs=image, + parameters={ + "candidate_labels": candidate_labels, + "hypothesis_template": hypothesis_template, + }, + headers=self.headers, + model=model_id, + api_key=self.token, + ) + response = await self._inner_post(request_parameters) + return ZeroShotImageClassificationOutputElement.parse_obj_as_list(response) + + @_deprecate_method( + version="0.33.0", + message=( + "HF Inference API is getting revamped and will only support warm models in the future (no cold start allowed)." + " Use `HfApi.list_models(..., inference_provider='...')` to list warm models per provider." + ), + ) + async def list_deployed_models( + self, frameworks: Union[None, str, Literal["all"], List[str]] = None + ) -> Dict[str, List[str]]: + """ + List models deployed on the HF Serverless Inference API service. + + This helper checks deployed models framework by framework. By default, it will check the 4 main frameworks that + are supported and account for 95% of the hosted models. However, if you want a complete list of models you can + specify `frameworks="all"` as input. Alternatively, if you know before-hand which framework you are interested + in, you can also restrict to search to this one (e.g. `frameworks="text-generation-inference"`). The more + frameworks are checked, the more time it will take. + + + + This endpoint method does not return a live list of all models available for the HF Inference API service. + It searches over a cached list of models that were recently available and the list may not be up to date. + If you want to know the live status of a specific model, use [`~InferenceClient.get_model_status`]. + + + + + + This endpoint method is mostly useful for discoverability. If you already know which model you want to use and want to + check its availability, you can directly use [`~InferenceClient.get_model_status`]. + + + + Args: + frameworks (`Literal["all"]` or `List[str]` or `str`, *optional*): + The frameworks to filter on. By default only a subset of the available frameworks are tested. If set to + "all", all available frameworks will be tested. It is also possible to provide a single framework or a + custom set of frameworks to check. + + Returns: + `Dict[str, List[str]]`: A dictionary mapping task names to a sorted list of model IDs. + + Example: + ```py + # Must be run in an async contextthon + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + + # Discover zero-shot-classification models currently deployed + >>> models = await client.list_deployed_models() + >>> models["zero-shot-classification"] + ['Narsil/deberta-large-mnli-zero-cls', 'facebook/bart-large-mnli', ...] + + # List from only 1 framework + >>> await client.list_deployed_models("text-generation-inference") + {'text-generation': ['bigcode/starcoder', 'meta-llama/Llama-2-70b-chat-hf', ...], ...} + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Listing deployed models is not supported on '{self.provider}'.") + + # Resolve which frameworks to check + if frameworks is None: + frameworks = constants.MAIN_INFERENCE_API_FRAMEWORKS + elif frameworks == "all": + frameworks = constants.ALL_INFERENCE_API_FRAMEWORKS + elif isinstance(frameworks, str): + frameworks = [frameworks] + frameworks = list(set(frameworks)) + + # Fetch them iteratively + models_by_task: Dict[str, List[str]] = {} + + def _unpack_response(framework: str, items: List[Dict]) -> None: + for model in items: + if framework == "sentence-transformers": + # Model running with the `sentence-transformers` framework can work with both tasks even if not + # branded as such in the API response + models_by_task.setdefault("feature-extraction", []).append(model["model_id"]) + models_by_task.setdefault("sentence-similarity", []).append(model["model_id"]) + else: + models_by_task.setdefault(model["task"], []).append(model["model_id"]) + + for framework in frameworks: + response = get_session().get( + f"{constants.INFERENCE_ENDPOINT}/framework/{framework}", headers=build_hf_headers(token=self.token) + ) + hf_raise_for_status(response) + _unpack_response(framework, response.json()) + + # Sort alphabetically for discoverability and return + for task, models in models_by_task.items(): + models_by_task[task] = sorted(set(models), key=lambda x: x.lower()) + return models_by_task + + def _get_client_session(self, headers: Optional[Dict] = None) -> "ClientSession": + aiohttp = _import_aiohttp() + client_headers = self.headers.copy() + if headers is not None: + client_headers.update(headers) + + # Return a new aiohttp ClientSession with correct settings. + session = aiohttp.ClientSession( + headers=client_headers, + cookies=self.cookies, + timeout=aiohttp.ClientTimeout(self.timeout), + trust_env=self.trust_env, + ) + + # Keep track of sessions to close them later + self._sessions[session] = set() + + # Override the `._request` method to register responses to be closed + session._wrapped_request = session._request + + async def _request(method, url, **kwargs): + response = await session._wrapped_request(method, url, **kwargs) + self._sessions[session].add(response) + return response + + session._request = _request + + # Override the 'close' method to + # 1. close ongoing responses + # 2. deregister the session when closed + session._close = session.close + + async def close_session(): + for response in self._sessions[session]: + response.close() + await session._close() + self._sessions.pop(session, None) + + session.close = close_session + return session + + async def get_endpoint_info(self, *, model: Optional[str] = None) -> Dict[str, Any]: + """ + Get information about the deployed endpoint. + + This endpoint is only available on endpoints powered by Text-Generation-Inference (TGI) or Text-Embedding-Inference (TEI). + Endpoints powered by `transformers` return an empty payload. + + Args: + model (`str`, *optional*): + The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed + Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + `Dict[str, Any]`: Information about the endpoint. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient("meta-llama/Meta-Llama-3-70B-Instruct") + >>> await client.get_endpoint_info() + { + 'model_id': 'meta-llama/Meta-Llama-3-70B-Instruct', + 'model_sha': None, + 'model_dtype': 'torch.float16', + 'model_device_type': 'cuda', + 'model_pipeline_tag': None, + 'max_concurrent_requests': 128, + 'max_best_of': 2, + 'max_stop_sequences': 4, + 'max_input_length': 8191, + 'max_total_tokens': 8192, + 'waiting_served_ratio': 0.3, + 'max_batch_total_tokens': 1259392, + 'max_waiting_tokens': 20, + 'max_batch_size': None, + 'validation_workers': 32, + 'max_client_batch_size': 4, + 'version': '2.0.2', + 'sha': 'dccab72549635c7eb5ddb17f43f0b7cdff07c214', + 'docker_label': 'sha-dccab72' + } + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Getting endpoint info is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if model.startswith(("http://", "https://")): + url = model.rstrip("/") + "/info" + else: + url = f"{constants.INFERENCE_ENDPOINT}/models/{model}/info" + + async with self._get_client_session(headers=build_hf_headers(token=self.token)) as client: + response = await client.get(url, proxy=self.proxies) + response.raise_for_status() + return await response.json() + + async def health_check(self, model: Optional[str] = None) -> bool: + """ + Check the health of the deployed endpoint. + + Health check is only available with Inference Endpoints powered by Text-Generation-Inference (TGI) or Text-Embedding-Inference (TEI). + For Inference API, please use [`InferenceClient.get_model_status`] instead. + + Args: + model (`str`, *optional*): + URL of the Inference Endpoint. This parameter overrides the model defined at the instance level. Defaults to None. + + Returns: + `bool`: True if everything is working fine. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient("https://jzgu0buei5.us-east-1.aws.endpoints.huggingface.cloud") + >>> await client.health_check() + True + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Health check is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if not model.startswith(("http://", "https://")): + raise ValueError( + "Model must be an Inference Endpoint URL. For serverless Inference API, please use `InferenceClient.get_model_status`." + ) + url = model.rstrip("/") + "/health" + + async with self._get_client_session(headers=build_hf_headers(token=self.token)) as client: + response = await client.get(url, proxy=self.proxies) + return response.status == 200 + + @_deprecate_method( + version="0.33.0", + message=( + "HF Inference API is getting revamped and will only support warm models in the future (no cold start allowed)." + " Use `HfApi.model_info` to get the model status both with HF Inference API and external providers." + ), + ) + async def get_model_status(self, model: Optional[str] = None) -> ModelStatus: + """ + Get the status of a model hosted on the HF Inference API. + + + + This endpoint is mostly useful when you already know which model you want to use and want to check its + availability. If you want to discover already deployed models, you should rather use [`~InferenceClient.list_deployed_models`]. + + + + Args: + model (`str`, *optional*): + Identifier of the model for witch the status gonna be checked. If model is not provided, + the model associated with this instance of [`InferenceClient`] will be used. Only HF Inference API service can be checked so the + identifier cannot be a URL. + + + Returns: + [`ModelStatus`]: An instance of ModelStatus dataclass, containing information, + about the state of the model: load, state, compute type and framework. + + Example: + ```py + # Must be run in an async context + >>> from huggingface_hub import AsyncInferenceClient + >>> client = AsyncInferenceClient() + >>> await client.get_model_status("meta-llama/Meta-Llama-3-8B-Instruct") + ModelStatus(loaded=True, state='Loaded', compute_type='gpu', framework='text-generation-inference') + ``` + """ + if self.provider != "hf-inference": + raise ValueError(f"Getting model status is not supported on '{self.provider}'.") + + model = model or self.model + if model is None: + raise ValueError("Model id not provided.") + if model.startswith("https://"): + raise NotImplementedError("Model status is only available for Inference API endpoints.") + url = f"{constants.INFERENCE_ENDPOINT}/status/{model}" + + async with self._get_client_session(headers=build_hf_headers(token=self.token)) as client: + response = await client.get(url, proxy=self.proxies) + response.raise_for_status() + response_data = await response.json() + + if "error" in response_data: + raise ValueError(response_data["error"]) + + return ModelStatus( + loaded=response_data["loaded"], + state=response_data["state"], + compute_type=response_data["compute_type"], + framework=response_data["framework"], + ) + + @property + def chat(self) -> "ProxyClientChat": + return ProxyClientChat(self) + + +class _ProxyClient: + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + def __init__(self, client: AsyncInferenceClient): + self._client = client + + +class ProxyClientChat(_ProxyClient): + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + @property + def completions(self) -> "ProxyClientChatCompletions": + return ProxyClientChatCompletions(self._client) + + +class ProxyClientChatCompletions(_ProxyClient): + """Proxy class to be able to call `client.chat.completion.create(...)` as OpenAI client.""" + + @property + def create(self): + return self._client.chat_completion diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__init__.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92f286792b67e21a1bfdef27823343cb906ff26e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__init__.py @@ -0,0 +1,188 @@ +# This file is auto-generated by `utils/generate_inference_types.py`. +# Do not modify it manually. +# +# ruff: noqa: F401 + +from .audio_classification import ( + AudioClassificationInput, + AudioClassificationOutputElement, + AudioClassificationOutputTransform, + AudioClassificationParameters, +) +from .audio_to_audio import AudioToAudioInput, AudioToAudioOutputElement +from .automatic_speech_recognition import ( + AutomaticSpeechRecognitionEarlyStoppingEnum, + AutomaticSpeechRecognitionGenerationParameters, + AutomaticSpeechRecognitionInput, + AutomaticSpeechRecognitionOutput, + AutomaticSpeechRecognitionOutputChunk, + AutomaticSpeechRecognitionParameters, +) +from .base import BaseInferenceType +from .chat_completion import ( + ChatCompletionInput, + ChatCompletionInputFunctionDefinition, + ChatCompletionInputFunctionName, + ChatCompletionInputGrammarType, + ChatCompletionInputGrammarTypeType, + ChatCompletionInputMessage, + ChatCompletionInputMessageChunk, + ChatCompletionInputMessageChunkType, + ChatCompletionInputStreamOptions, + ChatCompletionInputTool, + ChatCompletionInputToolCall, + ChatCompletionInputToolChoiceClass, + ChatCompletionInputToolChoiceEnum, + ChatCompletionInputURL, + ChatCompletionOutput, + ChatCompletionOutputComplete, + ChatCompletionOutputFunctionDefinition, + ChatCompletionOutputLogprob, + ChatCompletionOutputLogprobs, + ChatCompletionOutputMessage, + ChatCompletionOutputToolCall, + ChatCompletionOutputTopLogprob, + ChatCompletionOutputUsage, + ChatCompletionStreamOutput, + ChatCompletionStreamOutputChoice, + ChatCompletionStreamOutputDelta, + ChatCompletionStreamOutputDeltaToolCall, + ChatCompletionStreamOutputFunction, + ChatCompletionStreamOutputLogprob, + ChatCompletionStreamOutputLogprobs, + ChatCompletionStreamOutputTopLogprob, + ChatCompletionStreamOutputUsage, +) +from .depth_estimation import DepthEstimationInput, DepthEstimationOutput +from .document_question_answering import ( + DocumentQuestionAnsweringInput, + DocumentQuestionAnsweringInputData, + DocumentQuestionAnsweringOutputElement, + DocumentQuestionAnsweringParameters, +) +from .feature_extraction import FeatureExtractionInput, FeatureExtractionInputTruncationDirection +from .fill_mask import FillMaskInput, FillMaskOutputElement, FillMaskParameters +from .image_classification import ( + ImageClassificationInput, + ImageClassificationOutputElement, + ImageClassificationOutputTransform, + ImageClassificationParameters, +) +from .image_segmentation import ( + ImageSegmentationInput, + ImageSegmentationOutputElement, + ImageSegmentationParameters, + ImageSegmentationSubtask, +) +from .image_to_image import ImageToImageInput, ImageToImageOutput, ImageToImageParameters, ImageToImageTargetSize +from .image_to_text import ( + ImageToTextEarlyStoppingEnum, + ImageToTextGenerationParameters, + ImageToTextInput, + ImageToTextOutput, + ImageToTextParameters, +) +from .object_detection import ( + ObjectDetectionBoundingBox, + ObjectDetectionInput, + ObjectDetectionOutputElement, + ObjectDetectionParameters, +) +from .question_answering import ( + QuestionAnsweringInput, + QuestionAnsweringInputData, + QuestionAnsweringOutputElement, + QuestionAnsweringParameters, +) +from .sentence_similarity import SentenceSimilarityInput, SentenceSimilarityInputData +from .summarization import ( + SummarizationInput, + SummarizationOutput, + SummarizationParameters, + SummarizationTruncationStrategy, +) +from .table_question_answering import ( + Padding, + TableQuestionAnsweringInput, + TableQuestionAnsweringInputData, + TableQuestionAnsweringOutputElement, + TableQuestionAnsweringParameters, +) +from .text2text_generation import ( + Text2TextGenerationInput, + Text2TextGenerationOutput, + Text2TextGenerationParameters, + Text2TextGenerationTruncationStrategy, +) +from .text_classification import ( + TextClassificationInput, + TextClassificationOutputElement, + TextClassificationOutputTransform, + TextClassificationParameters, +) +from .text_generation import ( + TextGenerationInput, + TextGenerationInputGenerateParameters, + TextGenerationInputGrammarType, + TextGenerationOutput, + TextGenerationOutputBestOfSequence, + TextGenerationOutputDetails, + TextGenerationOutputFinishReason, + TextGenerationOutputPrefillToken, + TextGenerationOutputToken, + TextGenerationStreamOutput, + TextGenerationStreamOutputStreamDetails, + TextGenerationStreamOutputToken, + TypeEnum, +) +from .text_to_audio import ( + TextToAudioEarlyStoppingEnum, + TextToAudioGenerationParameters, + TextToAudioInput, + TextToAudioOutput, + TextToAudioParameters, +) +from .text_to_image import TextToImageInput, TextToImageOutput, TextToImageParameters +from .text_to_speech import ( + TextToSpeechEarlyStoppingEnum, + TextToSpeechGenerationParameters, + TextToSpeechInput, + TextToSpeechOutput, + TextToSpeechParameters, +) +from .text_to_video import TextToVideoInput, TextToVideoOutput, TextToVideoParameters +from .token_classification import ( + TokenClassificationAggregationStrategy, + TokenClassificationInput, + TokenClassificationOutputElement, + TokenClassificationParameters, +) +from .translation import TranslationInput, TranslationOutput, TranslationParameters, TranslationTruncationStrategy +from .video_classification import ( + VideoClassificationInput, + VideoClassificationOutputElement, + VideoClassificationOutputTransform, + VideoClassificationParameters, +) +from .visual_question_answering import ( + VisualQuestionAnsweringInput, + VisualQuestionAnsweringInputData, + VisualQuestionAnsweringOutputElement, + VisualQuestionAnsweringParameters, +) +from .zero_shot_classification import ( + ZeroShotClassificationInput, + ZeroShotClassificationOutputElement, + ZeroShotClassificationParameters, +) +from .zero_shot_image_classification import ( + ZeroShotImageClassificationInput, + ZeroShotImageClassificationOutputElement, + ZeroShotImageClassificationParameters, +) +from .zero_shot_object_detection import ( + ZeroShotObjectDetectionBoundingBox, + ZeroShotObjectDetectionInput, + ZeroShotObjectDetectionOutputElement, + ZeroShotObjectDetectionParameters, +) diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df5eb6dcc2801abd0672c2b6f2b7bb3d11eeb38e Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d4a934e1df2d4e5c0cd30f36241bb0f2823f18d Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_to_audio.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_to_audio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0363cf7590e172e391d7d2f88f0b688dfae5d4cf Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/audio_to_audio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/automatic_speech_recognition.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/automatic_speech_recognition.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8ec71e8fc64a2190982c73a10eb41050b133a0b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/automatic_speech_recognition.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ac9683628ea3b908fe3227c5126b6cf3e39c23a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/chat_completion.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/chat_completion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..935604056a4a50c73fabc8f4827d0682f166715d Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/chat_completion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/depth_estimation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/depth_estimation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b54c60fef85d84ecac6c67dbc9437eee7ac91045 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/depth_estimation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/document_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/document_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4363ff57efa9efa12935122005f3a2651e6d3d9 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/document_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/feature_extraction.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/feature_extraction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1936d7e69b9fccafe80d64502ce608afcc6e188f Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/feature_extraction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/fill_mask.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/fill_mask.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..567a6a895ce2df6e85aed3a67ee5107347cc55a0 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/fill_mask.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de8ce2fd1be19623f4673a6c7827a835fa992f45 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_segmentation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_segmentation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ed7c6c4d27e8cca0b68fe3a894b983dedc6449b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_segmentation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_image.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_image.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4ff8114093a9799ae5bc2995f102885392bb383 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_image.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_text.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_text.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2682f739b81d5d2b8ef0b2b85eed19fe6e389b59 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/image_to_text.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/object_detection.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/object_detection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35f2e43482a9c8a57dbacb36eb55321abbe32534 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/object_detection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/question_answering.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e38f09a408fb10789ff547a29d72aeb21756fe2a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/sentence_similarity.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/sentence_similarity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9802f58072b0002db0633ed467929b13e29fcebf Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/sentence_similarity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/summarization.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/summarization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf8242c72496e3bafb5366747a82d1af4b87d6fb Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/summarization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/table_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/table_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1be1434826e0c894226198a2239bd41a7ac31732 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/table_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text2text_generation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text2text_generation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d09a7ae6d44a55b8124a9b9c436508229146be04 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text2text_generation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d25f8fce9a0e004cf411b81a36d7b2a1bf4dcad Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_generation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_generation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffc1f61c623e95d5fd80c82763a318ab19f8d03c Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_generation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_audio.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_audio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6dbc18a78099e163ad3e2c4448e0bacd89906f9 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_audio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_image.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_image.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..949b190d5199df8117bb847ff4e312cf98f38dd4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_image.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_speech.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_speech.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79a5d2a8d655160c4758dd3f5ae66027626426a6 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_speech.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_video.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_video.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..609ded816975071cca90c722b98a6c31a099d603 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/text_to_video.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/token_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/token_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44d343e7279aae9f72c2e45e94881e0c0325dc68 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/token_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/translation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/translation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0e66e8bfe687bb6a208c24fa1a9d41983aa57b5 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/translation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/video_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/video_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb3d1004dc3abe6b1a51fa5b051cad1d880595a4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/video_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/visual_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/visual_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c900a487f6e85a72aeaae9af2f56287aa781041b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/visual_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33761fcac73e67708f073a3885b97f9a9f43ab2e Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_image_classification.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_image_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..567eb2d360a48db393748be9cb4b89d7600db73b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_image_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_object_detection.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_object_detection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb54d1bde0c71ef1c1bfa0a6e5d0c4bc1362c662 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__pycache__/zero_shot_object_detection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..053055787bce933e1fbd393cfbc00d81c43c8c2d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_classification.py @@ -0,0 +1,43 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +AudioClassificationOutputTransform = Literal["sigmoid", "softmax", "none"] + + +@dataclass_with_extra +class AudioClassificationParameters(BaseInferenceType): + """Additional inference parameters for Audio Classification""" + + function_to_apply: Optional["AudioClassificationOutputTransform"] = None + """The function to apply to the model outputs in order to retrieve the scores.""" + top_k: Optional[int] = None + """When specified, limits the output to the top K most probable classes.""" + + +@dataclass_with_extra +class AudioClassificationInput(BaseInferenceType): + """Inputs for Audio Classification inference""" + + inputs: str + """The input audio data as a base64-encoded string. If no `parameters` are provided, you can + also provide the audio data as a raw bytes payload. + """ + parameters: Optional[AudioClassificationParameters] = None + """Additional inference parameters for Audio Classification""" + + +@dataclass_with_extra +class AudioClassificationOutputElement(BaseInferenceType): + """Outputs for Audio Classification inference""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_to_audio.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_to_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..43f376b5345fab6b854b028d1c17416c020d7bc1 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_to_audio.py @@ -0,0 +1,30 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class AudioToAudioInput(BaseInferenceType): + """Inputs for Audio to Audio inference""" + + inputs: Any + """The input audio data""" + + +@dataclass_with_extra +class AudioToAudioOutputElement(BaseInferenceType): + """Outputs of inference for the Audio To Audio task + A generated audio file with its label. + """ + + blob: Any + """The generated audio file.""" + content_type: str + """The content type of audio file.""" + label: str + """The label of the audio file.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/automatic_speech_recognition.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/automatic_speech_recognition.py new file mode 100644 index 0000000000000000000000000000000000000000..f6bfd28256c82309b160f337aba5a54e2dd11872 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/automatic_speech_recognition.py @@ -0,0 +1,113 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +AutomaticSpeechRecognitionEarlyStoppingEnum = Literal["never"] + + +@dataclass_with_extra +class AutomaticSpeechRecognitionGenerationParameters(BaseInferenceType): + """Parametrization of the text generation process""" + + do_sample: Optional[bool] = None + """Whether to use sampling instead of greedy decoding when generating new tokens.""" + early_stopping: Optional[Union[bool, "AutomaticSpeechRecognitionEarlyStoppingEnum"]] = None + """Controls the stopping condition for beam-based methods.""" + epsilon_cutoff: Optional[float] = None + """If set to float strictly between 0 and 1, only tokens with a conditional probability + greater than epsilon_cutoff will be sampled. In the paper, suggested values range from + 3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language + Model Desmoothing](https://hf.co/papers/2210.15191) for more details. + """ + eta_cutoff: Optional[float] = None + """Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to + float strictly between 0 and 1, a token is only considered if it is greater than either + eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter + term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In + the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model. + See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191) + for more details. + """ + max_length: Optional[int] = None + """The maximum length (in tokens) of the generated text, including the input.""" + max_new_tokens: Optional[int] = None + """The maximum number of tokens to generate. Takes precedence over max_length.""" + min_length: Optional[int] = None + """The minimum length (in tokens) of the generated text, including the input.""" + min_new_tokens: Optional[int] = None + """The minimum number of tokens to generate. Takes precedence over min_length.""" + num_beam_groups: Optional[int] = None + """Number of groups to divide num_beams into in order to ensure diversity among different + groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details. + """ + num_beams: Optional[int] = None + """Number of beams to use for beam search.""" + penalty_alpha: Optional[float] = None + """The value balances the model confidence and the degeneration penalty in contrastive + search decoding. + """ + temperature: Optional[float] = None + """The value used to modulate the next token probabilities.""" + top_k: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-k-filtering.""" + top_p: Optional[float] = None + """If set to float < 1, only the smallest set of most probable tokens with probabilities + that add up to top_p or higher are kept for generation. + """ + typical_p: Optional[float] = None + """Local typicality measures how similar the conditional probability of predicting a target + token next is to the expected conditional probability of predicting a random token next, + given the partial text already generated. If set to float < 1, the smallest set of the + most locally typical tokens with probabilities that add up to typical_p or higher are + kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details. + """ + use_cache: Optional[bool] = None + """Whether the model should use the past last key/values attentions to speed up decoding""" + + +@dataclass_with_extra +class AutomaticSpeechRecognitionParameters(BaseInferenceType): + """Additional inference parameters for Automatic Speech Recognition""" + + generation_parameters: Optional[AutomaticSpeechRecognitionGenerationParameters] = None + """Parametrization of the text generation process""" + return_timestamps: Optional[bool] = None + """Whether to output corresponding timestamps with the generated text""" + + +@dataclass_with_extra +class AutomaticSpeechRecognitionInput(BaseInferenceType): + """Inputs for Automatic Speech Recognition inference""" + + inputs: str + """The input audio data as a base64-encoded string. If no `parameters` are provided, you can + also provide the audio data as a raw bytes payload. + """ + parameters: Optional[AutomaticSpeechRecognitionParameters] = None + """Additional inference parameters for Automatic Speech Recognition""" + + +@dataclass_with_extra +class AutomaticSpeechRecognitionOutputChunk(BaseInferenceType): + text: str + """A chunk of text identified by the model""" + timestamp: List[float] + """The start and end timestamps corresponding with the text""" + + +@dataclass_with_extra +class AutomaticSpeechRecognitionOutput(BaseInferenceType): + """Outputs of inference for the Automatic Speech Recognition task""" + + text: str + """The recognized text.""" + chunks: Optional[List[AutomaticSpeechRecognitionOutputChunk]] = None + """When returnTimestamps is enabled, chunks contains a list of audio chunks identified by + the model. + """ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/base.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1f0c4687ceccbfb738da3f38c583c2516d065a01 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/base.py @@ -0,0 +1,161 @@ +# 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. +"""Contains a base class for all inference types.""" + +import inspect +import json +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Type, TypeVar, Union, get_args + + +T = TypeVar("T", bound="BaseInferenceType") + + +def _repr_with_extra(self): + fields = list(self.__dataclass_fields__.keys()) + other_fields = list(k for k in self.__dict__ if k not in fields) + return f"{self.__class__.__name__}({', '.join(f'{k}={self.__dict__[k]!r}' for k in fields + other_fields)})" + + +def dataclass_with_extra(cls: Type[T]) -> Type[T]: + """Decorator to add a custom __repr__ method to a dataclass, showing all fields, including extra ones. + + This decorator only works with dataclasses that inherit from `BaseInferenceType`. + """ + cls = dataclass(cls) + cls.__repr__ = _repr_with_extra # type: ignore[method-assign] + return cls + + +@dataclass +class BaseInferenceType(dict): + """Base class for all inference types. + + Object is a dataclass and a dict for backward compatibility but plan is to remove the dict part in the future. + + Handle parsing from dict, list and json strings in a permissive way to ensure future-compatibility (e.g. all fields + are made optional, and non-expected fields are added as dict attributes). + """ + + @classmethod + def parse_obj_as_list(cls: Type[T], data: Union[bytes, str, List, Dict]) -> List[T]: + """Alias to parse server response and return a single instance. + + See `parse_obj` for more details. + """ + output = cls.parse_obj(data) + if not isinstance(output, list): + raise ValueError(f"Invalid input data for {cls}. Expected a list, but got {type(output)}.") + return output + + @classmethod + def parse_obj_as_instance(cls: Type[T], data: Union[bytes, str, List, Dict]) -> T: + """Alias to parse server response and return a single instance. + + See `parse_obj` for more details. + """ + output = cls.parse_obj(data) + if isinstance(output, list): + raise ValueError(f"Invalid input data for {cls}. Expected a single instance, but got a list.") + return output + + @classmethod + def parse_obj(cls: Type[T], data: Union[bytes, str, List, Dict]) -> Union[List[T], T]: + """Parse server response as a dataclass or list of dataclasses. + + To enable future-compatibility, we want to handle cases where the server return more fields than expected. + In such cases, we don't want to raise an error but still create the dataclass object. Remaining fields are + added as dict attributes. + """ + # Parse server response (from bytes) + if isinstance(data, bytes): + data = data.decode() + if isinstance(data, str): + data = json.loads(data) + + # If a list, parse each item individually + if isinstance(data, List): + return [cls.parse_obj(d) for d in data] # type: ignore [misc] + + # At this point, we expect a dict + if not isinstance(data, dict): + raise ValueError(f"Invalid data type: {type(data)}") + + init_values = {} + other_values = {} + for key, value in data.items(): + key = normalize_key(key) + if key in cls.__dataclass_fields__ and cls.__dataclass_fields__[key].init: + if isinstance(value, dict) or isinstance(value, list): + field_type = cls.__dataclass_fields__[key].type + + # if `field_type` is a `BaseInferenceType`, parse it + if inspect.isclass(field_type) and issubclass(field_type, BaseInferenceType): + value = field_type.parse_obj(value) + + # otherwise, recursively parse nested dataclasses (if possible) + # `get_args` returns handle Union and Optional for us + else: + expected_types = get_args(field_type) + for expected_type in expected_types: + if getattr(expected_type, "_name", None) == "List": + expected_type = get_args(expected_type)[ + 0 + ] # assume same type for all items in the list + if inspect.isclass(expected_type) and issubclass(expected_type, BaseInferenceType): + value = expected_type.parse_obj(value) + break + init_values[key] = value + else: + other_values[key] = value + + # Make all missing fields default to None + # => ensure that dataclass initialization will never fail even if the server does not return all fields. + for key in cls.__dataclass_fields__: + if key not in init_values: + init_values[key] = None + + # Initialize dataclass with expected values + item = cls(**init_values) + + # Add remaining fields as dict attributes + item.update(other_values) + + # Add remaining fields as extra dataclass fields. + # They won't be part of the dataclass fields but will be accessible as attributes. + # Use @dataclass_with_extra to show them in __repr__. + item.__dict__.update(other_values) + return item + + def __post_init__(self): + self.update(asdict(self)) + + def __setitem__(self, __key: Any, __value: Any) -> None: + # Hacky way to keep dataclass values in sync when dict is updated + super().__setitem__(__key, __value) + if __key in self.__dataclass_fields__ and getattr(self, __key, None) != __value: + self.__setattr__(__key, __value) + return + + def __setattr__(self, __name: str, __value: Any) -> None: + # Hacky way to keep dict values is sync when dataclass is updated + super().__setattr__(__name, __value) + if self.get(__name) != __value: + self[__name] = __value + return + + +def normalize_key(key: str) -> str: + # e.g "content-type" -> "content_type", "Accept" -> "accept" + return key.replace("-", "_").replace(" ", "_").lower() diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/chat_completion.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/chat_completion.py new file mode 100644 index 0000000000000000000000000000000000000000..9978c0a5a9f4b86cf168ff1508f8c2a35513e9f2 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/chat_completion.py @@ -0,0 +1,311 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, List, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ChatCompletionInputURL(BaseInferenceType): + url: str + + +ChatCompletionInputMessageChunkType = Literal["text", "image_url"] + + +@dataclass_with_extra +class ChatCompletionInputMessageChunk(BaseInferenceType): + type: "ChatCompletionInputMessageChunkType" + image_url: Optional[ChatCompletionInputURL] = None + text: Optional[str] = None + + +@dataclass_with_extra +class ChatCompletionInputFunctionDefinition(BaseInferenceType): + name: str + parameters: Any + description: Optional[str] = None + + +@dataclass_with_extra +class ChatCompletionInputToolCall(BaseInferenceType): + function: ChatCompletionInputFunctionDefinition + id: str + type: str + + +@dataclass_with_extra +class ChatCompletionInputMessage(BaseInferenceType): + role: str + content: Optional[Union[List[ChatCompletionInputMessageChunk], str]] = None + name: Optional[str] = None + tool_calls: Optional[List[ChatCompletionInputToolCall]] = None + + +ChatCompletionInputGrammarTypeType = Literal["json", "regex", "json_schema"] + + +@dataclass_with_extra +class ChatCompletionInputGrammarType(BaseInferenceType): + type: "ChatCompletionInputGrammarTypeType" + value: Any + """A string that represents a [JSON Schema](https://json-schema.org/). + JSON Schema is a declarative language that allows to annotate JSON documents + with types and descriptions. + """ + + +@dataclass_with_extra +class ChatCompletionInputStreamOptions(BaseInferenceType): + include_usage: Optional[bool] = None + """If set, an additional chunk will be streamed before the data: [DONE] message. The usage + field on this chunk shows the token usage statistics for the entire request, and the + choices field will always be an empty array. All other chunks will also include a usage + field, but with a null value. + """ + + +@dataclass_with_extra +class ChatCompletionInputFunctionName(BaseInferenceType): + name: str + + +@dataclass_with_extra +class ChatCompletionInputToolChoiceClass(BaseInferenceType): + function: ChatCompletionInputFunctionName + + +ChatCompletionInputToolChoiceEnum = Literal["auto", "none", "required"] + + +@dataclass_with_extra +class ChatCompletionInputTool(BaseInferenceType): + function: ChatCompletionInputFunctionDefinition + type: str + + +@dataclass_with_extra +class ChatCompletionInput(BaseInferenceType): + """Chat Completion Input. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + messages: List[ChatCompletionInputMessage] + """A list of messages comprising the conversation so far.""" + frequency_penalty: Optional[float] = None + """Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing + frequency in the text so far, + decreasing the model's likelihood to repeat the same line verbatim. + """ + logit_bias: Optional[List[float]] = None + """UNUSED + Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON + object that maps tokens + (specified by their token ID in the tokenizer) to an associated bias value from -100 to + 100. Mathematically, + the bias is added to the logits generated by the model prior to sampling. The exact + effect will vary per model, + but values between -1 and 1 should decrease or increase likelihood of selection; values + like -100 or 100 should + result in a ban or exclusive selection of the relevant token. + """ + logprobs: Optional[bool] = None + """Whether to return log probabilities of the output tokens or not. If true, returns the log + probabilities of each + output token returned in the content of message. + """ + max_tokens: Optional[int] = None + """The maximum number of tokens that can be generated in the chat completion.""" + model: Optional[str] = None + """[UNUSED] ID of the model to use. See the model endpoint compatibility table for details + on which models work with the Chat API. + """ + n: Optional[int] = None + """UNUSED + How many chat completion choices to generate for each input message. Note that you will + be charged based on the + number of generated tokens across all of the choices. Keep n as 1 to minimize costs. + """ + presence_penalty: Optional[float] = None + """Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they + appear in the text so far, + increasing the model's likelihood to talk about new topics + """ + response_format: Optional[ChatCompletionInputGrammarType] = None + seed: Optional[int] = None + stop: Optional[List[str]] = None + """Up to 4 sequences where the API will stop generating further tokens.""" + stream: Optional[bool] = None + stream_options: Optional[ChatCompletionInputStreamOptions] = None + temperature: Optional[float] = None + """What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the + output more random, while + lower values like 0.2 will make it more focused and deterministic. + We generally recommend altering this or `top_p` but not both. + """ + tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None + tool_prompt: Optional[str] = None + """A prompt to be appended before the tools""" + tools: Optional[List[ChatCompletionInputTool]] = None + """A list of tools the model may call. Currently, only functions are supported as a tool. + Use this to provide a list of + functions the model may generate JSON inputs for. + """ + top_logprobs: Optional[int] = None + """An integer between 0 and 5 specifying the number of most likely tokens to return at each + token position, each with + an associated log probability. logprobs must be set to true if this parameter is used. + """ + top_p: Optional[float] = None + """An alternative to sampling with temperature, called nucleus sampling, where the model + considers the results of the + tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% + probability mass are considered. + """ + + +@dataclass_with_extra +class ChatCompletionOutputTopLogprob(BaseInferenceType): + logprob: float + token: str + + +@dataclass_with_extra +class ChatCompletionOutputLogprob(BaseInferenceType): + logprob: float + token: str + top_logprobs: List[ChatCompletionOutputTopLogprob] + + +@dataclass_with_extra +class ChatCompletionOutputLogprobs(BaseInferenceType): + content: List[ChatCompletionOutputLogprob] + + +@dataclass_with_extra +class ChatCompletionOutputFunctionDefinition(BaseInferenceType): + arguments: str + name: str + description: Optional[str] = None + + +@dataclass_with_extra +class ChatCompletionOutputToolCall(BaseInferenceType): + function: ChatCompletionOutputFunctionDefinition + id: str + type: str + + +@dataclass_with_extra +class ChatCompletionOutputMessage(BaseInferenceType): + role: str + content: Optional[str] = None + tool_call_id: Optional[str] = None + tool_calls: Optional[List[ChatCompletionOutputToolCall]] = None + + +@dataclass_with_extra +class ChatCompletionOutputComplete(BaseInferenceType): + finish_reason: str + index: int + message: ChatCompletionOutputMessage + logprobs: Optional[ChatCompletionOutputLogprobs] = None + + +@dataclass_with_extra +class ChatCompletionOutputUsage(BaseInferenceType): + completion_tokens: int + prompt_tokens: int + total_tokens: int + + +@dataclass_with_extra +class ChatCompletionOutput(BaseInferenceType): + """Chat Completion Output. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + choices: List[ChatCompletionOutputComplete] + created: int + id: str + model: str + system_fingerprint: str + usage: ChatCompletionOutputUsage + + +@dataclass_with_extra +class ChatCompletionStreamOutputFunction(BaseInferenceType): + arguments: str + name: Optional[str] = None + + +@dataclass_with_extra +class ChatCompletionStreamOutputDeltaToolCall(BaseInferenceType): + function: ChatCompletionStreamOutputFunction + id: str + index: int + type: str + + +@dataclass_with_extra +class ChatCompletionStreamOutputDelta(BaseInferenceType): + role: str + content: Optional[str] = None + tool_call_id: Optional[str] = None + tool_calls: Optional[List[ChatCompletionStreamOutputDeltaToolCall]] = None + + +@dataclass_with_extra +class ChatCompletionStreamOutputTopLogprob(BaseInferenceType): + logprob: float + token: str + + +@dataclass_with_extra +class ChatCompletionStreamOutputLogprob(BaseInferenceType): + logprob: float + token: str + top_logprobs: List[ChatCompletionStreamOutputTopLogprob] + + +@dataclass_with_extra +class ChatCompletionStreamOutputLogprobs(BaseInferenceType): + content: List[ChatCompletionStreamOutputLogprob] + + +@dataclass_with_extra +class ChatCompletionStreamOutputChoice(BaseInferenceType): + delta: ChatCompletionStreamOutputDelta + index: int + finish_reason: Optional[str] = None + logprobs: Optional[ChatCompletionStreamOutputLogprobs] = None + + +@dataclass_with_extra +class ChatCompletionStreamOutputUsage(BaseInferenceType): + completion_tokens: int + prompt_tokens: int + total_tokens: int + + +@dataclass_with_extra +class ChatCompletionStreamOutput(BaseInferenceType): + """Chat Completion Stream Output. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + choices: List[ChatCompletionStreamOutputChoice] + created: int + id: str + model: str + system_fingerprint: str + usage: Optional[ChatCompletionStreamOutputUsage] = None diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/depth_estimation.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/depth_estimation.py new file mode 100644 index 0000000000000000000000000000000000000000..1e09bdffa194f97444e484de6e930f67ac030207 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/depth_estimation.py @@ -0,0 +1,28 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Dict, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class DepthEstimationInput(BaseInferenceType): + """Inputs for Depth Estimation inference""" + + inputs: Any + """The input image data""" + parameters: Optional[Dict[str, Any]] = None + """Additional inference parameters for Depth Estimation""" + + +@dataclass_with_extra +class DepthEstimationOutput(BaseInferenceType): + """Outputs of inference for the Depth Estimation task""" + + depth: Any + """The predicted depth as an image""" + predicted_depth: Any + """The predicted depth as a tensor""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/document_question_answering.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/document_question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..2457d2c8c237f055f660e0e8291d846bb036949d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/document_question_answering.py @@ -0,0 +1,80 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, List, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class DocumentQuestionAnsweringInputData(BaseInferenceType): + """One (document, question) pair to answer""" + + image: Any + """The image on which the question is asked""" + question: str + """A question to ask of the document""" + + +@dataclass_with_extra +class DocumentQuestionAnsweringParameters(BaseInferenceType): + """Additional inference parameters for Document Question Answering""" + + doc_stride: Optional[int] = None + """If the words in the document are too long to fit with the question for the model, it will + be split in several chunks with some overlap. This argument controls the size of that + overlap. + """ + handle_impossible_answer: Optional[bool] = None + """Whether to accept impossible as an answer""" + lang: Optional[str] = None + """Language to use while running OCR. Defaults to english.""" + max_answer_len: Optional[int] = None + """The maximum length of predicted answers (e.g., only answers with a shorter length are + considered). + """ + max_question_len: Optional[int] = None + """The maximum length of the question after tokenization. It will be truncated if needed.""" + max_seq_len: Optional[int] = None + """The maximum length of the total sentence (context + question) in tokens of each chunk + passed to the model. The context will be split in several chunks (using doc_stride as + overlap) if needed. + """ + top_k: Optional[int] = None + """The number of answers to return (will be chosen by order of likelihood). Can return less + than top_k answers if there are not enough options available within the context. + """ + word_boxes: Optional[List[Union[List[float], str]]] = None + """A list of words and bounding boxes (normalized 0->1000). If provided, the inference will + skip the OCR step and use the provided bounding boxes instead. + """ + + +@dataclass_with_extra +class DocumentQuestionAnsweringInput(BaseInferenceType): + """Inputs for Document Question Answering inference""" + + inputs: DocumentQuestionAnsweringInputData + """One (document, question) pair to answer""" + parameters: Optional[DocumentQuestionAnsweringParameters] = None + """Additional inference parameters for Document Question Answering""" + + +@dataclass_with_extra +class DocumentQuestionAnsweringOutputElement(BaseInferenceType): + """Outputs of inference for the Document Question Answering task""" + + answer: str + """The answer to the question.""" + end: int + """The end word index of the answer (in the OCR’d version of the input or provided word + boxes). + """ + score: float + """The probability associated to the answer.""" + start: int + """The start word index of the answer (in the OCR’d version of the input or provided word + boxes). + """ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/feature_extraction.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/feature_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..e965ddbac2af0a5bf73e662a7c18c847611d18a1 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/feature_extraction.py @@ -0,0 +1,36 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +FeatureExtractionInputTruncationDirection = Literal["Left", "Right"] + + +@dataclass_with_extra +class FeatureExtractionInput(BaseInferenceType): + """Feature Extraction Input. + Auto-generated from TEI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts. + """ + + inputs: Union[List[str], str] + """The text or list of texts to embed.""" + normalize: Optional[bool] = None + prompt_name: Optional[str] = None + """The name of the prompt that should be used by for encoding. If not set, no prompt + will be applied. + Must be a key in the `sentence-transformers` configuration `prompts` dictionary. + For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ", + ...}, + then the sentence "What is the capital of France?" will be encoded as + "query: What is the capital of France?" because the prompt text will be prepended before + any text to encode. + """ + truncate: Optional[bool] = None + truncation_direction: Optional["FeatureExtractionInputTruncationDirection"] = None diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/fill_mask.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/fill_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..dfcdc56bc507e50280d38e0f63b024ada6a7ea94 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/fill_mask.py @@ -0,0 +1,47 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, List, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class FillMaskParameters(BaseInferenceType): + """Additional inference parameters for Fill Mask""" + + targets: Optional[List[str]] = None + """When passed, the model will limit the scores to the passed targets instead of looking up + in the whole vocabulary. If the provided targets are not in the model vocab, they will be + tokenized and the first resulting token will be used (with a warning, and that might be + slower). + """ + top_k: Optional[int] = None + """When passed, overrides the number of predictions to return.""" + + +@dataclass_with_extra +class FillMaskInput(BaseInferenceType): + """Inputs for Fill Mask inference""" + + inputs: str + """The text with masked tokens""" + parameters: Optional[FillMaskParameters] = None + """Additional inference parameters for Fill Mask""" + + +@dataclass_with_extra +class FillMaskOutputElement(BaseInferenceType): + """Outputs of inference for the Fill Mask task""" + + score: float + """The corresponding probability""" + sequence: str + """The corresponding input with the mask token prediction.""" + token: int + """The predicted token id (to replace the masked one).""" + token_str: Any + fill_mask_output_token_str: Optional[str] = None + """The predicted token (to replace the masked one).""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..0fdda6c83ff4c7aee5dc7794f0530e89d6b43047 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_classification.py @@ -0,0 +1,43 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +ImageClassificationOutputTransform = Literal["sigmoid", "softmax", "none"] + + +@dataclass_with_extra +class ImageClassificationParameters(BaseInferenceType): + """Additional inference parameters for Image Classification""" + + function_to_apply: Optional["ImageClassificationOutputTransform"] = None + """The function to apply to the model outputs in order to retrieve the scores.""" + top_k: Optional[int] = None + """When specified, limits the output to the top K most probable classes.""" + + +@dataclass_with_extra +class ImageClassificationInput(BaseInferenceType): + """Inputs for Image Classification inference""" + + inputs: str + """The input image data as a base64-encoded string. If no `parameters` are provided, you can + also provide the image data as a raw bytes payload. + """ + parameters: Optional[ImageClassificationParameters] = None + """Additional inference parameters for Image Classification""" + + +@dataclass_with_extra +class ImageClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Image Classification task""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_segmentation.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_segmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..3dbf61db83ec2ae6ceafd901c4425567cd2e5b03 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_segmentation.py @@ -0,0 +1,51 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +ImageSegmentationSubtask = Literal["instance", "panoptic", "semantic"] + + +@dataclass_with_extra +class ImageSegmentationParameters(BaseInferenceType): + """Additional inference parameters for Image Segmentation""" + + mask_threshold: Optional[float] = None + """Threshold to use when turning the predicted masks into binary values.""" + overlap_mask_area_threshold: Optional[float] = None + """Mask overlap threshold to eliminate small, disconnected segments.""" + subtask: Optional["ImageSegmentationSubtask"] = None + """Segmentation task to be performed, depending on model capabilities.""" + threshold: Optional[float] = None + """Probability threshold to filter out predicted masks.""" + + +@dataclass_with_extra +class ImageSegmentationInput(BaseInferenceType): + """Inputs for Image Segmentation inference""" + + inputs: str + """The input image data as a base64-encoded string. If no `parameters` are provided, you can + also provide the image data as a raw bytes payload. + """ + parameters: Optional[ImageSegmentationParameters] = None + """Additional inference parameters for Image Segmentation""" + + +@dataclass_with_extra +class ImageSegmentationOutputElement(BaseInferenceType): + """Outputs of inference for the Image Segmentation task + A predicted mask / segment + """ + + label: str + """The label of the predicted segment.""" + mask: str + """The corresponding mask as a black-and-white image (base64-encoded).""" + score: Optional[float] = None + """The score or confidence degree the model has.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_image.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_image.py new file mode 100644 index 0000000000000000000000000000000000000000..e99e719f8b837c76392b58d33ca19f9b615e857e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_image.py @@ -0,0 +1,56 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ImageToImageTargetSize(BaseInferenceType): + """The size in pixel of the output image.""" + + height: int + width: int + + +@dataclass_with_extra +class ImageToImageParameters(BaseInferenceType): + """Additional inference parameters for Image To Image""" + + guidance_scale: Optional[float] = None + """For diffusion models. A higher guidance scale value encourages the model to generate + images closely linked to the text prompt at the expense of lower image quality. + """ + negative_prompt: Optional[str] = None + """One prompt to guide what NOT to include in image generation.""" + num_inference_steps: Optional[int] = None + """For diffusion models. The number of denoising steps. More denoising steps usually lead to + a higher quality image at the expense of slower inference. + """ + prompt: Optional[str] = None + """The text prompt to guide the image generation.""" + target_size: Optional[ImageToImageTargetSize] = None + """The size in pixel of the output image.""" + + +@dataclass_with_extra +class ImageToImageInput(BaseInferenceType): + """Inputs for Image To Image inference""" + + inputs: str + """The input image data as a base64-encoded string. If no `parameters` are provided, you can + also provide the image data as a raw bytes payload. + """ + parameters: Optional[ImageToImageParameters] = None + """Additional inference parameters for Image To Image""" + + +@dataclass_with_extra +class ImageToImageOutput(BaseInferenceType): + """Outputs of inference for the Image To Image task""" + + image: Any + """The output image returned as raw bytes in the payload.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_text.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..b65e0e0068e80dbcab5a4706fb5d49be2538c4ca --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_text.py @@ -0,0 +1,100 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +ImageToTextEarlyStoppingEnum = Literal["never"] + + +@dataclass_with_extra +class ImageToTextGenerationParameters(BaseInferenceType): + """Parametrization of the text generation process""" + + do_sample: Optional[bool] = None + """Whether to use sampling instead of greedy decoding when generating new tokens.""" + early_stopping: Optional[Union[bool, "ImageToTextEarlyStoppingEnum"]] = None + """Controls the stopping condition for beam-based methods.""" + epsilon_cutoff: Optional[float] = None + """If set to float strictly between 0 and 1, only tokens with a conditional probability + greater than epsilon_cutoff will be sampled. In the paper, suggested values range from + 3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language + Model Desmoothing](https://hf.co/papers/2210.15191) for more details. + """ + eta_cutoff: Optional[float] = None + """Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to + float strictly between 0 and 1, a token is only considered if it is greater than either + eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter + term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In + the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model. + See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191) + for more details. + """ + max_length: Optional[int] = None + """The maximum length (in tokens) of the generated text, including the input.""" + max_new_tokens: Optional[int] = None + """The maximum number of tokens to generate. Takes precedence over max_length.""" + min_length: Optional[int] = None + """The minimum length (in tokens) of the generated text, including the input.""" + min_new_tokens: Optional[int] = None + """The minimum number of tokens to generate. Takes precedence over min_length.""" + num_beam_groups: Optional[int] = None + """Number of groups to divide num_beams into in order to ensure diversity among different + groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details. + """ + num_beams: Optional[int] = None + """Number of beams to use for beam search.""" + penalty_alpha: Optional[float] = None + """The value balances the model confidence and the degeneration penalty in contrastive + search decoding. + """ + temperature: Optional[float] = None + """The value used to modulate the next token probabilities.""" + top_k: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-k-filtering.""" + top_p: Optional[float] = None + """If set to float < 1, only the smallest set of most probable tokens with probabilities + that add up to top_p or higher are kept for generation. + """ + typical_p: Optional[float] = None + """Local typicality measures how similar the conditional probability of predicting a target + token next is to the expected conditional probability of predicting a random token next, + given the partial text already generated. If set to float < 1, the smallest set of the + most locally typical tokens with probabilities that add up to typical_p or higher are + kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details. + """ + use_cache: Optional[bool] = None + """Whether the model should use the past last key/values attentions to speed up decoding""" + + +@dataclass_with_extra +class ImageToTextParameters(BaseInferenceType): + """Additional inference parameters for Image To Text""" + + generation_parameters: Optional[ImageToTextGenerationParameters] = None + """Parametrization of the text generation process""" + max_new_tokens: Optional[int] = None + """The amount of maximum tokens to generate.""" + + +@dataclass_with_extra +class ImageToTextInput(BaseInferenceType): + """Inputs for Image To Text inference""" + + inputs: Any + """The input image data""" + parameters: Optional[ImageToTextParameters] = None + """Additional inference parameters for Image To Text""" + + +@dataclass_with_extra +class ImageToTextOutput(BaseInferenceType): + """Outputs of inference for the Image To Text task""" + + generated_text: Any + image_to_text_output_generated_text: Optional[str] = None + """The generated text.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/object_detection.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/object_detection.py new file mode 100644 index 0000000000000000000000000000000000000000..75f3ebcfe1199462d0df60879b5ba6e517f7001e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/object_detection.py @@ -0,0 +1,58 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ObjectDetectionParameters(BaseInferenceType): + """Additional inference parameters for Object Detection""" + + threshold: Optional[float] = None + """The probability necessary to make a prediction.""" + + +@dataclass_with_extra +class ObjectDetectionInput(BaseInferenceType): + """Inputs for Object Detection inference""" + + inputs: str + """The input image data as a base64-encoded string. If no `parameters` are provided, you can + also provide the image data as a raw bytes payload. + """ + parameters: Optional[ObjectDetectionParameters] = None + """Additional inference parameters for Object Detection""" + + +@dataclass_with_extra +class ObjectDetectionBoundingBox(BaseInferenceType): + """The predicted bounding box. Coordinates are relative to the top left corner of the input + image. + """ + + xmax: int + """The x-coordinate of the bottom-right corner of the bounding box.""" + xmin: int + """The x-coordinate of the top-left corner of the bounding box.""" + ymax: int + """The y-coordinate of the bottom-right corner of the bounding box.""" + ymin: int + """The y-coordinate of the top-left corner of the bounding box.""" + + +@dataclass_with_extra +class ObjectDetectionOutputElement(BaseInferenceType): + """Outputs of inference for the Object Detection task""" + + box: ObjectDetectionBoundingBox + """The predicted bounding box. Coordinates are relative to the top left corner of the input + image. + """ + label: str + """The predicted label for the bounding box.""" + score: float + """The associated score / probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/question_answering.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..014ab41893c560a2c266bc04a1d60bc933be31c7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/question_answering.py @@ -0,0 +1,74 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class QuestionAnsweringInputData(BaseInferenceType): + """One (context, question) pair to answer""" + + context: str + """The context to be used for answering the question""" + question: str + """The question to be answered""" + + +@dataclass_with_extra +class QuestionAnsweringParameters(BaseInferenceType): + """Additional inference parameters for Question Answering""" + + align_to_words: Optional[bool] = None + """Attempts to align the answer to real words. Improves quality on space separated + languages. Might hurt on non-space-separated languages (like Japanese or Chinese) + """ + doc_stride: Optional[int] = None + """If the context is too long to fit with the question for the model, it will be split in + several chunks with some overlap. This argument controls the size of that overlap. + """ + handle_impossible_answer: Optional[bool] = None + """Whether to accept impossible as an answer.""" + max_answer_len: Optional[int] = None + """The maximum length of predicted answers (e.g., only answers with a shorter length are + considered). + """ + max_question_len: Optional[int] = None + """The maximum length of the question after tokenization. It will be truncated if needed.""" + max_seq_len: Optional[int] = None + """The maximum length of the total sentence (context + question) in tokens of each chunk + passed to the model. The context will be split in several chunks (using docStride as + overlap) if needed. + """ + top_k: Optional[int] = None + """The number of answers to return (will be chosen by order of likelihood). Note that we + return less than topk answers if there are not enough options available within the + context. + """ + + +@dataclass_with_extra +class QuestionAnsweringInput(BaseInferenceType): + """Inputs for Question Answering inference""" + + inputs: QuestionAnsweringInputData + """One (context, question) pair to answer""" + parameters: Optional[QuestionAnsweringParameters] = None + """Additional inference parameters for Question Answering""" + + +@dataclass_with_extra +class QuestionAnsweringOutputElement(BaseInferenceType): + """Outputs of inference for the Question Answering task""" + + answer: str + """The answer to the question.""" + end: int + """The character position in the input where the answer ends.""" + score: float + """The probability associated to the answer.""" + start: int + """The character position in the input where the answer begins.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/sentence_similarity.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/sentence_similarity.py new file mode 100644 index 0000000000000000000000000000000000000000..66e8bb4d9322d4847556b7a17dc17bd208a37d0c --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/sentence_similarity.py @@ -0,0 +1,27 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Dict, List, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class SentenceSimilarityInputData(BaseInferenceType): + sentences: List[str] + """A list of strings which will be compared against the source_sentence.""" + source_sentence: str + """The string that you wish to compare the other strings with. This can be a phrase, + sentence, or longer passage, depending on the model being used. + """ + + +@dataclass_with_extra +class SentenceSimilarityInput(BaseInferenceType): + """Inputs for Sentence similarity inference""" + + inputs: SentenceSimilarityInputData + parameters: Optional[Dict[str, Any]] = None + """Additional inference parameters for Sentence Similarity""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/summarization.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/summarization.py new file mode 100644 index 0000000000000000000000000000000000000000..33eae6fcba0e8724babf145f93be005868429c33 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/summarization.py @@ -0,0 +1,41 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Dict, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +SummarizationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"] + + +@dataclass_with_extra +class SummarizationParameters(BaseInferenceType): + """Additional inference parameters for summarization.""" + + clean_up_tokenization_spaces: Optional[bool] = None + """Whether to clean up the potential extra spaces in the text output.""" + generate_parameters: Optional[Dict[str, Any]] = None + """Additional parametrization of the text generation algorithm.""" + truncation: Optional["SummarizationTruncationStrategy"] = None + """The truncation strategy to use.""" + + +@dataclass_with_extra +class SummarizationInput(BaseInferenceType): + """Inputs for Summarization inference""" + + inputs: str + """The input text to summarize.""" + parameters: Optional[SummarizationParameters] = None + """Additional inference parameters for summarization.""" + + +@dataclass_with_extra +class SummarizationOutput(BaseInferenceType): + """Outputs of inference for the Summarization task""" + + summary_text: str + """The summarized text.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/table_question_answering.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/table_question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..10e208eeeb50a689d2826a160432a2b005ec006c --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/table_question_answering.py @@ -0,0 +1,62 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Dict, List, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class TableQuestionAnsweringInputData(BaseInferenceType): + """One (table, question) pair to answer""" + + question: str + """The question to be answered about the table""" + table: Dict[str, List[str]] + """The table to serve as context for the questions""" + + +Padding = Literal["do_not_pad", "longest", "max_length"] + + +@dataclass_with_extra +class TableQuestionAnsweringParameters(BaseInferenceType): + """Additional inference parameters for Table Question Answering""" + + padding: Optional["Padding"] = None + """Activates and controls padding.""" + sequential: Optional[bool] = None + """Whether to do inference sequentially or as a batch. Batching is faster, but models like + SQA require the inference to be done sequentially to extract relations within sequences, + given their conversational nature. + """ + truncation: Optional[bool] = None + """Activates and controls truncation.""" + + +@dataclass_with_extra +class TableQuestionAnsweringInput(BaseInferenceType): + """Inputs for Table Question Answering inference""" + + inputs: TableQuestionAnsweringInputData + """One (table, question) pair to answer""" + parameters: Optional[TableQuestionAnsweringParameters] = None + """Additional inference parameters for Table Question Answering""" + + +@dataclass_with_extra +class TableQuestionAnsweringOutputElement(BaseInferenceType): + """Outputs of inference for the Table Question Answering task""" + + answer: str + """The answer of the question given the table. If there is an aggregator, the answer will be + preceded by `AGGREGATOR >`. + """ + cells: List[str] + """List of strings made up of the answer cell values.""" + coordinates: List[List[int]] + """Coordinates of the cells of the answers.""" + aggregator: Optional[str] = None + """If the model has an aggregator, this returns the aggregator.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text2text_generation.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text2text_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..34ac74e21e8a30d889f1a251f648d4c365325be6 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text2text_generation.py @@ -0,0 +1,42 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Dict, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +Text2TextGenerationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"] + + +@dataclass_with_extra +class Text2TextGenerationParameters(BaseInferenceType): + """Additional inference parameters for Text2text Generation""" + + clean_up_tokenization_spaces: Optional[bool] = None + """Whether to clean up the potential extra spaces in the text output.""" + generate_parameters: Optional[Dict[str, Any]] = None + """Additional parametrization of the text generation algorithm""" + truncation: Optional["Text2TextGenerationTruncationStrategy"] = None + """The truncation strategy to use""" + + +@dataclass_with_extra +class Text2TextGenerationInput(BaseInferenceType): + """Inputs for Text2text Generation inference""" + + inputs: str + """The input text data""" + parameters: Optional[Text2TextGenerationParameters] = None + """Additional inference parameters for Text2text Generation""" + + +@dataclass_with_extra +class Text2TextGenerationOutput(BaseInferenceType): + """Outputs of inference for the Text2text Generation task""" + + generated_text: Any + text2_text_generation_output_generated_text: Optional[str] = None + """The generated text.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..9a172b23f844fa58f757a644d52138a18e7b6ddb --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_classification.py @@ -0,0 +1,41 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +TextClassificationOutputTransform = Literal["sigmoid", "softmax", "none"] + + +@dataclass_with_extra +class TextClassificationParameters(BaseInferenceType): + """Additional inference parameters for Text Classification""" + + function_to_apply: Optional["TextClassificationOutputTransform"] = None + """The function to apply to the model outputs in order to retrieve the scores.""" + top_k: Optional[int] = None + """When specified, limits the output to the top K most probable classes.""" + + +@dataclass_with_extra +class TextClassificationInput(BaseInferenceType): + """Inputs for Text Classification inference""" + + inputs: str + """The text to classify""" + parameters: Optional[TextClassificationParameters] = None + """Additional inference parameters for Text Classification""" + + +@dataclass_with_extra +class TextClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Text Classification task""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_generation.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..9b79cc691dce3a6d42aef716d4a93a719f2d600c --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_generation.py @@ -0,0 +1,168 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, List, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +TypeEnum = Literal["json", "regex", "json_schema"] + + +@dataclass_with_extra +class TextGenerationInputGrammarType(BaseInferenceType): + type: "TypeEnum" + value: Any + """A string that represents a [JSON Schema](https://json-schema.org/). + JSON Schema is a declarative language that allows to annotate JSON documents + with types and descriptions. + """ + + +@dataclass_with_extra +class TextGenerationInputGenerateParameters(BaseInferenceType): + adapter_id: Optional[str] = None + """Lora adapter id""" + best_of: Optional[int] = None + """Generate best_of sequences and return the one if the highest token logprobs.""" + decoder_input_details: Optional[bool] = None + """Whether to return decoder input token logprobs and ids.""" + details: Optional[bool] = None + """Whether to return generation details.""" + do_sample: Optional[bool] = None + """Activate logits sampling.""" + frequency_penalty: Optional[float] = None + """The parameter for frequency penalty. 1.0 means no penalty + Penalize new tokens based on their existing frequency in the text so far, + decreasing the model's likelihood to repeat the same line verbatim. + """ + grammar: Optional[TextGenerationInputGrammarType] = None + max_new_tokens: Optional[int] = None + """Maximum number of tokens to generate.""" + repetition_penalty: Optional[float] = None + """The parameter for repetition penalty. 1.0 means no penalty. + See [this paper](https://arxiv.org/pdf/1909.05858.pdf) for more details. + """ + return_full_text: Optional[bool] = None + """Whether to prepend the prompt to the generated text""" + seed: Optional[int] = None + """Random sampling seed.""" + stop: Optional[List[str]] = None + """Stop generating tokens if a member of `stop` is generated.""" + temperature: Optional[float] = None + """The value used to module the logits distribution.""" + top_k: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-k-filtering.""" + top_n_tokens: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-n-filtering.""" + top_p: Optional[float] = None + """Top-p value for nucleus sampling.""" + truncate: Optional[int] = None + """Truncate inputs tokens to the given size.""" + typical_p: Optional[float] = None + """Typical Decoding mass + See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) + for more information. + """ + watermark: Optional[bool] = None + """Watermarking with [A Watermark for Large Language + Models](https://arxiv.org/abs/2301.10226). + """ + + +@dataclass_with_extra +class TextGenerationInput(BaseInferenceType): + """Text Generation Input. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + inputs: str + parameters: Optional[TextGenerationInputGenerateParameters] = None + stream: Optional[bool] = None + + +TextGenerationOutputFinishReason = Literal["length", "eos_token", "stop_sequence"] + + +@dataclass_with_extra +class TextGenerationOutputPrefillToken(BaseInferenceType): + id: int + logprob: float + text: str + + +@dataclass_with_extra +class TextGenerationOutputToken(BaseInferenceType): + id: int + logprob: float + special: bool + text: str + + +@dataclass_with_extra +class TextGenerationOutputBestOfSequence(BaseInferenceType): + finish_reason: "TextGenerationOutputFinishReason" + generated_text: str + generated_tokens: int + prefill: List[TextGenerationOutputPrefillToken] + tokens: List[TextGenerationOutputToken] + seed: Optional[int] = None + top_tokens: Optional[List[List[TextGenerationOutputToken]]] = None + + +@dataclass_with_extra +class TextGenerationOutputDetails(BaseInferenceType): + finish_reason: "TextGenerationOutputFinishReason" + generated_tokens: int + prefill: List[TextGenerationOutputPrefillToken] + tokens: List[TextGenerationOutputToken] + best_of_sequences: Optional[List[TextGenerationOutputBestOfSequence]] = None + seed: Optional[int] = None + top_tokens: Optional[List[List[TextGenerationOutputToken]]] = None + + +@dataclass_with_extra +class TextGenerationOutput(BaseInferenceType): + """Text Generation Output. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + generated_text: str + details: Optional[TextGenerationOutputDetails] = None + + +@dataclass_with_extra +class TextGenerationStreamOutputStreamDetails(BaseInferenceType): + finish_reason: "TextGenerationOutputFinishReason" + generated_tokens: int + input_length: int + seed: Optional[int] = None + + +@dataclass_with_extra +class TextGenerationStreamOutputToken(BaseInferenceType): + id: int + logprob: float + special: bool + text: str + + +@dataclass_with_extra +class TextGenerationStreamOutput(BaseInferenceType): + """Text Generation Stream Output. + Auto-generated from TGI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts. + """ + + index: int + token: TextGenerationStreamOutputToken + details: Optional[TextGenerationStreamOutputStreamDetails] = None + generated_text: Optional[str] = None + top_tokens: Optional[List[TextGenerationStreamOutputToken]] = None diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_audio.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_audio.py new file mode 100644 index 0000000000000000000000000000000000000000..87af80a598af70800b8386f034c65de0b397479e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_audio.py @@ -0,0 +1,99 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +TextToAudioEarlyStoppingEnum = Literal["never"] + + +@dataclass_with_extra +class TextToAudioGenerationParameters(BaseInferenceType): + """Parametrization of the text generation process""" + + do_sample: Optional[bool] = None + """Whether to use sampling instead of greedy decoding when generating new tokens.""" + early_stopping: Optional[Union[bool, "TextToAudioEarlyStoppingEnum"]] = None + """Controls the stopping condition for beam-based methods.""" + epsilon_cutoff: Optional[float] = None + """If set to float strictly between 0 and 1, only tokens with a conditional probability + greater than epsilon_cutoff will be sampled. In the paper, suggested values range from + 3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language + Model Desmoothing](https://hf.co/papers/2210.15191) for more details. + """ + eta_cutoff: Optional[float] = None + """Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to + float strictly between 0 and 1, a token is only considered if it is greater than either + eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter + term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In + the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model. + See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191) + for more details. + """ + max_length: Optional[int] = None + """The maximum length (in tokens) of the generated text, including the input.""" + max_new_tokens: Optional[int] = None + """The maximum number of tokens to generate. Takes precedence over max_length.""" + min_length: Optional[int] = None + """The minimum length (in tokens) of the generated text, including the input.""" + min_new_tokens: Optional[int] = None + """The minimum number of tokens to generate. Takes precedence over min_length.""" + num_beam_groups: Optional[int] = None + """Number of groups to divide num_beams into in order to ensure diversity among different + groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details. + """ + num_beams: Optional[int] = None + """Number of beams to use for beam search.""" + penalty_alpha: Optional[float] = None + """The value balances the model confidence and the degeneration penalty in contrastive + search decoding. + """ + temperature: Optional[float] = None + """The value used to modulate the next token probabilities.""" + top_k: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-k-filtering.""" + top_p: Optional[float] = None + """If set to float < 1, only the smallest set of most probable tokens with probabilities + that add up to top_p or higher are kept for generation. + """ + typical_p: Optional[float] = None + """Local typicality measures how similar the conditional probability of predicting a target + token next is to the expected conditional probability of predicting a random token next, + given the partial text already generated. If set to float < 1, the smallest set of the + most locally typical tokens with probabilities that add up to typical_p or higher are + kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details. + """ + use_cache: Optional[bool] = None + """Whether the model should use the past last key/values attentions to speed up decoding""" + + +@dataclass_with_extra +class TextToAudioParameters(BaseInferenceType): + """Additional inference parameters for Text To Audio""" + + generation_parameters: Optional[TextToAudioGenerationParameters] = None + """Parametrization of the text generation process""" + + +@dataclass_with_extra +class TextToAudioInput(BaseInferenceType): + """Inputs for Text To Audio inference""" + + inputs: str + """The input text data""" + parameters: Optional[TextToAudioParameters] = None + """Additional inference parameters for Text To Audio""" + + +@dataclass_with_extra +class TextToAudioOutput(BaseInferenceType): + """Outputs of inference for the Text To Audio task""" + + audio: Any + """The generated audio waveform.""" + sampling_rate: float + """The sampling rate of the generated audio waveform.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_image.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_image.py new file mode 100644 index 0000000000000000000000000000000000000000..20c963731371339975019ca5d40c95303d79209b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_image.py @@ -0,0 +1,50 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class TextToImageParameters(BaseInferenceType): + """Additional inference parameters for Text To Image""" + + guidance_scale: Optional[float] = None + """A higher guidance scale value encourages the model to generate images closely linked to + the text prompt, but values too high may cause saturation and other artifacts. + """ + height: Optional[int] = None + """The height in pixels of the output image""" + negative_prompt: Optional[str] = None + """One prompt to guide what NOT to include in image generation.""" + num_inference_steps: Optional[int] = None + """The number of denoising steps. More denoising steps usually lead to a higher quality + image at the expense of slower inference. + """ + scheduler: Optional[str] = None + """Override the scheduler with a compatible one.""" + seed: Optional[int] = None + """Seed for the random number generator.""" + width: Optional[int] = None + """The width in pixels of the output image""" + + +@dataclass_with_extra +class TextToImageInput(BaseInferenceType): + """Inputs for Text To Image inference""" + + inputs: str + """The input text data (sometimes called "prompt")""" + parameters: Optional[TextToImageParameters] = None + """Additional inference parameters for Text To Image""" + + +@dataclass_with_extra +class TextToImageOutput(BaseInferenceType): + """Outputs of inference for the Text To Image task""" + + image: Any + """The generated image returned as raw bytes in the payload.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_speech.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_speech.py new file mode 100644 index 0000000000000000000000000000000000000000..ce2db8f3f901cc99b5d2fcbb362c4b07b2a718e0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_speech.py @@ -0,0 +1,99 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Literal, Optional, Union + +from .base import BaseInferenceType, dataclass_with_extra + + +TextToSpeechEarlyStoppingEnum = Literal["never"] + + +@dataclass_with_extra +class TextToSpeechGenerationParameters(BaseInferenceType): + """Parametrization of the text generation process""" + + do_sample: Optional[bool] = None + """Whether to use sampling instead of greedy decoding when generating new tokens.""" + early_stopping: Optional[Union[bool, "TextToSpeechEarlyStoppingEnum"]] = None + """Controls the stopping condition for beam-based methods.""" + epsilon_cutoff: Optional[float] = None + """If set to float strictly between 0 and 1, only tokens with a conditional probability + greater than epsilon_cutoff will be sampled. In the paper, suggested values range from + 3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language + Model Desmoothing](https://hf.co/papers/2210.15191) for more details. + """ + eta_cutoff: Optional[float] = None + """Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to + float strictly between 0 and 1, a token is only considered if it is greater than either + eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter + term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In + the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model. + See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191) + for more details. + """ + max_length: Optional[int] = None + """The maximum length (in tokens) of the generated text, including the input.""" + max_new_tokens: Optional[int] = None + """The maximum number of tokens to generate. Takes precedence over max_length.""" + min_length: Optional[int] = None + """The minimum length (in tokens) of the generated text, including the input.""" + min_new_tokens: Optional[int] = None + """The minimum number of tokens to generate. Takes precedence over min_length.""" + num_beam_groups: Optional[int] = None + """Number of groups to divide num_beams into in order to ensure diversity among different + groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details. + """ + num_beams: Optional[int] = None + """Number of beams to use for beam search.""" + penalty_alpha: Optional[float] = None + """The value balances the model confidence and the degeneration penalty in contrastive + search decoding. + """ + temperature: Optional[float] = None + """The value used to modulate the next token probabilities.""" + top_k: Optional[int] = None + """The number of highest probability vocabulary tokens to keep for top-k-filtering.""" + top_p: Optional[float] = None + """If set to float < 1, only the smallest set of most probable tokens with probabilities + that add up to top_p or higher are kept for generation. + """ + typical_p: Optional[float] = None + """Local typicality measures how similar the conditional probability of predicting a target + token next is to the expected conditional probability of predicting a random token next, + given the partial text already generated. If set to float < 1, the smallest set of the + most locally typical tokens with probabilities that add up to typical_p or higher are + kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details. + """ + use_cache: Optional[bool] = None + """Whether the model should use the past last key/values attentions to speed up decoding""" + + +@dataclass_with_extra +class TextToSpeechParameters(BaseInferenceType): + """Additional inference parameters for Text To Speech""" + + generation_parameters: Optional[TextToSpeechGenerationParameters] = None + """Parametrization of the text generation process""" + + +@dataclass_with_extra +class TextToSpeechInput(BaseInferenceType): + """Inputs for Text To Speech inference""" + + inputs: str + """The input text data""" + parameters: Optional[TextToSpeechParameters] = None + """Additional inference parameters for Text To Speech""" + + +@dataclass_with_extra +class TextToSpeechOutput(BaseInferenceType): + """Outputs of inference for the Text To Speech task""" + + audio: Any + """The generated audio""" + sampling_rate: Optional[float] = None + """The sampling rate of the generated audio waveform.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_video.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_video.py new file mode 100644 index 0000000000000000000000000000000000000000..e54a1bc094e4aaf7132e502aa268bc052ab34f0a --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_video.py @@ -0,0 +1,46 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, List, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class TextToVideoParameters(BaseInferenceType): + """Additional inference parameters for Text To Video""" + + guidance_scale: Optional[float] = None + """A higher guidance scale value encourages the model to generate videos closely linked to + the text prompt, but values too high may cause saturation and other artifacts. + """ + negative_prompt: Optional[List[str]] = None + """One or several prompt to guide what NOT to include in video generation.""" + num_frames: Optional[float] = None + """The num_frames parameter determines how many video frames are generated.""" + num_inference_steps: Optional[int] = None + """The number of denoising steps. More denoising steps usually lead to a higher quality + video at the expense of slower inference. + """ + seed: Optional[int] = None + """Seed for the random number generator.""" + + +@dataclass_with_extra +class TextToVideoInput(BaseInferenceType): + """Inputs for Text To Video inference""" + + inputs: str + """The input text data (sometimes called "prompt")""" + parameters: Optional[TextToVideoParameters] = None + """Additional inference parameters for Text To Video""" + + +@dataclass_with_extra +class TextToVideoOutput(BaseInferenceType): + """Outputs of inference for the Text To Video task""" + + video: Any + """The generated video returned as raw bytes in the payload.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/token_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/token_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..e039b6a1db7dcd54dbc9434d3254da0770c6799e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/token_classification.py @@ -0,0 +1,51 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +TokenClassificationAggregationStrategy = Literal["none", "simple", "first", "average", "max"] + + +@dataclass_with_extra +class TokenClassificationParameters(BaseInferenceType): + """Additional inference parameters for Token Classification""" + + aggregation_strategy: Optional["TokenClassificationAggregationStrategy"] = None + """The strategy used to fuse tokens based on model predictions""" + ignore_labels: Optional[List[str]] = None + """A list of labels to ignore""" + stride: Optional[int] = None + """The number of overlapping tokens between chunks when splitting the input text.""" + + +@dataclass_with_extra +class TokenClassificationInput(BaseInferenceType): + """Inputs for Token Classification inference""" + + inputs: str + """The input text data""" + parameters: Optional[TokenClassificationParameters] = None + """Additional inference parameters for Token Classification""" + + +@dataclass_with_extra +class TokenClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Token Classification task""" + + end: int + """The character position in the input where this group ends.""" + score: float + """The associated score / probability""" + start: int + """The character position in the input where this group begins.""" + word: str + """The corresponding text""" + entity: Optional[str] = None + """The predicted label for a single token""" + entity_group: Optional[str] = None + """The predicted label for a group of one or more tokens""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/translation.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/translation.py new file mode 100644 index 0000000000000000000000000000000000000000..df95b7dbb1f4ce5b80cec034e004bb6e71387be8 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/translation.py @@ -0,0 +1,49 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Dict, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +TranslationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"] + + +@dataclass_with_extra +class TranslationParameters(BaseInferenceType): + """Additional inference parameters for Translation""" + + clean_up_tokenization_spaces: Optional[bool] = None + """Whether to clean up the potential extra spaces in the text output.""" + generate_parameters: Optional[Dict[str, Any]] = None + """Additional parametrization of the text generation algorithm.""" + src_lang: Optional[str] = None + """The source language of the text. Required for models that can translate from multiple + languages. + """ + tgt_lang: Optional[str] = None + """Target language to translate to. Required for models that can translate to multiple + languages. + """ + truncation: Optional["TranslationTruncationStrategy"] = None + """The truncation strategy to use.""" + + +@dataclass_with_extra +class TranslationInput(BaseInferenceType): + """Inputs for Translation inference""" + + inputs: str + """The text to translate.""" + parameters: Optional[TranslationParameters] = None + """Additional inference parameters for Translation""" + + +@dataclass_with_extra +class TranslationOutput(BaseInferenceType): + """Outputs of inference for the Translation task""" + + translation_text: str + """The translated text.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/video_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/video_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..e1d7a15bb4ee5fa63aa6ebc3750191bd38549212 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/video_classification.py @@ -0,0 +1,45 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Literal, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +VideoClassificationOutputTransform = Literal["sigmoid", "softmax", "none"] + + +@dataclass_with_extra +class VideoClassificationParameters(BaseInferenceType): + """Additional inference parameters for Video Classification""" + + frame_sampling_rate: Optional[int] = None + """The sampling rate used to select frames from the video.""" + function_to_apply: Optional["VideoClassificationOutputTransform"] = None + """The function to apply to the model outputs in order to retrieve the scores.""" + num_frames: Optional[int] = None + """The number of sampled frames to consider for classification.""" + top_k: Optional[int] = None + """When specified, limits the output to the top K most probable classes.""" + + +@dataclass_with_extra +class VideoClassificationInput(BaseInferenceType): + """Inputs for Video Classification inference""" + + inputs: Any + """The input video data""" + parameters: Optional[VideoClassificationParameters] = None + """Additional inference parameters for Video Classification""" + + +@dataclass_with_extra +class VideoClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Video Classification task""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/visual_question_answering.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/visual_question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..d368f1621289bc11a17be3e590cf8a040019d455 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/visual_question_answering.py @@ -0,0 +1,49 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import Any, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class VisualQuestionAnsweringInputData(BaseInferenceType): + """One (image, question) pair to answer""" + + image: Any + """The image.""" + question: str + """The question to answer based on the image.""" + + +@dataclass_with_extra +class VisualQuestionAnsweringParameters(BaseInferenceType): + """Additional inference parameters for Visual Question Answering""" + + top_k: Optional[int] = None + """The number of answers to return (will be chosen by order of likelihood). Note that we + return less than topk answers if there are not enough options available within the + context. + """ + + +@dataclass_with_extra +class VisualQuestionAnsweringInput(BaseInferenceType): + """Inputs for Visual Question Answering inference""" + + inputs: VisualQuestionAnsweringInputData + """One (image, question) pair to answer""" + parameters: Optional[VisualQuestionAnsweringParameters] = None + """Additional inference parameters for Visual Question Answering""" + + +@dataclass_with_extra +class VisualQuestionAnsweringOutputElement(BaseInferenceType): + """Outputs of inference for the Visual Question Answering task""" + + score: float + """The associated score / probability""" + answer: Optional[str] = None + """The answer to the question""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..47b32492e358edcc0de6aa09d53635b0a8156b25 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_classification.py @@ -0,0 +1,45 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ZeroShotClassificationParameters(BaseInferenceType): + """Additional inference parameters for Zero Shot Classification""" + + candidate_labels: List[str] + """The set of possible class labels to classify the text into.""" + hypothesis_template: Optional[str] = None + """The sentence used in conjunction with `candidate_labels` to attempt the text + classification by replacing the placeholder with the candidate labels. + """ + multi_label: Optional[bool] = None + """Whether multiple candidate labels can be true. If false, the scores are normalized such + that the sum of the label likelihoods for each sequence is 1. If true, the labels are + considered independent and probabilities are normalized for each candidate. + """ + + +@dataclass_with_extra +class ZeroShotClassificationInput(BaseInferenceType): + """Inputs for Zero Shot Classification inference""" + + inputs: str + """The text to classify""" + parameters: ZeroShotClassificationParameters + """Additional inference parameters for Zero Shot Classification""" + + +@dataclass_with_extra +class ZeroShotClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Zero Shot Classification task""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_image_classification.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_image_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..998d66b6b4e3356f0f09a0ad25ebdaf2e76cd03f --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_image_classification.py @@ -0,0 +1,40 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List, Optional + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ZeroShotImageClassificationParameters(BaseInferenceType): + """Additional inference parameters for Zero Shot Image Classification""" + + candidate_labels: List[str] + """The candidate labels for this image""" + hypothesis_template: Optional[str] = None + """The sentence used in conjunction with `candidate_labels` to attempt the image + classification by replacing the placeholder with the candidate labels. + """ + + +@dataclass_with_extra +class ZeroShotImageClassificationInput(BaseInferenceType): + """Inputs for Zero Shot Image Classification inference""" + + inputs: str + """The input image data to classify as a base64-encoded string.""" + parameters: ZeroShotImageClassificationParameters + """Additional inference parameters for Zero Shot Image Classification""" + + +@dataclass_with_extra +class ZeroShotImageClassificationOutputElement(BaseInferenceType): + """Outputs of inference for the Zero Shot Image Classification task""" + + label: str + """The predicted class label.""" + score: float + """The corresponding probability.""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_object_detection.py b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_object_detection.py new file mode 100644 index 0000000000000000000000000000000000000000..8ef76b5fcb93e8126266e4b1464934d01024b1b7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_object_detection.py @@ -0,0 +1,52 @@ +# Inference code generated from the JSON schema spec in @huggingface/tasks. +# +# See: +# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts +# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks. +from typing import List + +from .base import BaseInferenceType, dataclass_with_extra + + +@dataclass_with_extra +class ZeroShotObjectDetectionParameters(BaseInferenceType): + """Additional inference parameters for Zero Shot Object Detection""" + + candidate_labels: List[str] + """The candidate labels for this image""" + + +@dataclass_with_extra +class ZeroShotObjectDetectionInput(BaseInferenceType): + """Inputs for Zero Shot Object Detection inference""" + + inputs: str + """The input image data as a base64-encoded string.""" + parameters: ZeroShotObjectDetectionParameters + """Additional inference parameters for Zero Shot Object Detection""" + + +@dataclass_with_extra +class ZeroShotObjectDetectionBoundingBox(BaseInferenceType): + """The predicted bounding box. Coordinates are relative to the top left corner of the input + image. + """ + + xmax: int + xmin: int + ymax: int + ymin: int + + +@dataclass_with_extra +class ZeroShotObjectDetectionOutputElement(BaseInferenceType): + """Outputs of inference for the Zero Shot Object Detection task""" + + box: ZeroShotObjectDetectionBoundingBox + """The predicted bounding box. Coordinates are relative to the top left corner of the input + image. + """ + label: str + """A candidate label""" + score: float + """The associated score / probability""" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__init__.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73c3009904ecb595b41ce7b466eb472bb6ba2c51 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__init__.py @@ -0,0 +1,179 @@ +from typing import Dict, Literal, Optional, Union + +from huggingface_hub.utils import logging + +from ._common import TaskProviderHelper, _fetch_inference_provider_mapping +from .black_forest_labs import BlackForestLabsTextToImageTask +from .cerebras import CerebrasConversationalTask +from .cohere import CohereConversationalTask +from .fal_ai import ( + FalAIAutomaticSpeechRecognitionTask, + FalAITextToImageTask, + FalAITextToSpeechTask, + FalAITextToVideoTask, +) +from .fireworks_ai import FireworksAIConversationalTask +from .hf_inference import ( + HFInferenceBinaryInputTask, + HFInferenceConversational, + HFInferenceFeatureExtractionTask, + HFInferenceTask, +) +from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask +from .nebius import NebiusConversationalTask, NebiusTextGenerationTask, NebiusTextToImageTask +from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask +from .openai import OpenAIConversationalTask +from .replicate import ReplicateTask, ReplicateTextToImageTask, ReplicateTextToSpeechTask +from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask +from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask + + +logger = logging.get_logger(__name__) + + +PROVIDER_T = Literal[ + "black-forest-labs", + "cerebras", + "cohere", + "fal-ai", + "fireworks-ai", + "hf-inference", + "hyperbolic", + "nebius", + "novita", + "openai", + "replicate", + "sambanova", + "together", +] + +PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]] + +PROVIDERS: Dict[PROVIDER_T, Dict[str, TaskProviderHelper]] = { + "black-forest-labs": { + "text-to-image": BlackForestLabsTextToImageTask(), + }, + "cerebras": { + "conversational": CerebrasConversationalTask(), + }, + "cohere": { + "conversational": CohereConversationalTask(), + }, + "fal-ai": { + "automatic-speech-recognition": FalAIAutomaticSpeechRecognitionTask(), + "text-to-image": FalAITextToImageTask(), + "text-to-speech": FalAITextToSpeechTask(), + "text-to-video": FalAITextToVideoTask(), + }, + "fireworks-ai": { + "conversational": FireworksAIConversationalTask(), + }, + "hf-inference": { + "text-to-image": HFInferenceTask("text-to-image"), + "conversational": HFInferenceConversational(), + "text-generation": HFInferenceTask("text-generation"), + "text-classification": HFInferenceTask("text-classification"), + "question-answering": HFInferenceTask("question-answering"), + "audio-classification": HFInferenceBinaryInputTask("audio-classification"), + "automatic-speech-recognition": HFInferenceBinaryInputTask("automatic-speech-recognition"), + "fill-mask": HFInferenceTask("fill-mask"), + "feature-extraction": HFInferenceFeatureExtractionTask(), + "image-classification": HFInferenceBinaryInputTask("image-classification"), + "image-segmentation": HFInferenceBinaryInputTask("image-segmentation"), + "document-question-answering": HFInferenceTask("document-question-answering"), + "image-to-text": HFInferenceBinaryInputTask("image-to-text"), + "object-detection": HFInferenceBinaryInputTask("object-detection"), + "audio-to-audio": HFInferenceBinaryInputTask("audio-to-audio"), + "zero-shot-image-classification": HFInferenceBinaryInputTask("zero-shot-image-classification"), + "zero-shot-classification": HFInferenceTask("zero-shot-classification"), + "image-to-image": HFInferenceBinaryInputTask("image-to-image"), + "sentence-similarity": HFInferenceTask("sentence-similarity"), + "table-question-answering": HFInferenceTask("table-question-answering"), + "tabular-classification": HFInferenceTask("tabular-classification"), + "text-to-speech": HFInferenceTask("text-to-speech"), + "token-classification": HFInferenceTask("token-classification"), + "translation": HFInferenceTask("translation"), + "summarization": HFInferenceTask("summarization"), + "visual-question-answering": HFInferenceBinaryInputTask("visual-question-answering"), + }, + "hyperbolic": { + "text-to-image": HyperbolicTextToImageTask(), + "conversational": HyperbolicTextGenerationTask("conversational"), + "text-generation": HyperbolicTextGenerationTask("text-generation"), + }, + "nebius": { + "text-to-image": NebiusTextToImageTask(), + "conversational": NebiusConversationalTask(), + "text-generation": NebiusTextGenerationTask(), + }, + "novita": { + "text-generation": NovitaTextGenerationTask(), + "conversational": NovitaConversationalTask(), + "text-to-video": NovitaTextToVideoTask(), + }, + "openai": { + "conversational": OpenAIConversationalTask(), + }, + "replicate": { + "text-to-image": ReplicateTextToImageTask(), + "text-to-speech": ReplicateTextToSpeechTask(), + "text-to-video": ReplicateTask("text-to-video"), + }, + "sambanova": { + "conversational": SambanovaConversationalTask(), + "feature-extraction": SambanovaFeatureExtractionTask(), + }, + "together": { + "text-to-image": TogetherTextToImageTask(), + "conversational": TogetherConversationalTask(), + "text-generation": TogetherTextGenerationTask(), + }, +} + + +def get_provider_helper( + provider: Optional[PROVIDER_OR_POLICY_T], task: str, model: Optional[str] +) -> TaskProviderHelper: + """Get provider helper instance by name and task. + + Args: + provider (`str`, *optional*): name of the provider, or "auto" to automatically select the provider for the model. + task (`str`): Name of the task + model (`str`, *optional*): Name of the model + Returns: + TaskProviderHelper: Helper instance for the specified provider and task + + Raises: + ValueError: If provider or task is not supported + """ + + if (model is None and provider in (None, "auto")) or ( + model is not None and model.startswith(("http://", "https://")) + ): + provider = "hf-inference" + + if provider is None: + logger.info( + "Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers." + ) + provider = "auto" + + if provider == "auto": + if model is None: + raise ValueError("Specifying a model is required when provider is 'auto'") + provider_mapping = _fetch_inference_provider_mapping(model) + provider = next(iter(provider_mapping)) + + provider_tasks = PROVIDERS.get(provider) # type: ignore + if provider_tasks is None: + raise ValueError( + f"Provider '{provider}' not supported. Available values: 'auto' or any provider from {list(PROVIDERS.keys())}." + "Passing 'auto' (default value) will automatically select the first provider available for the model, sorted " + "by the user's order in https://hf.co/settings/inference-providers." + ) + + if task not in provider_tasks: + raise ValueError( + f"Task '{task}' not supported for provider '{provider}'. Available tasks: {list(provider_tasks.keys())}" + ) + return provider_tasks[task] diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d0b14fcd4948998efc2c92da65bc7c0a6f3dafb Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/_common.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..084d7a412d64116597b88bb6373f4dd78aef8d85 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/black_forest_labs.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/black_forest_labs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5058b79e828eb54cc34580d5941e4c4c1ec6abf Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/black_forest_labs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cerebras.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cerebras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d129281b60da5c84bbcd298f90936489a1cdd8dc Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cerebras.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cohere.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cohere.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96168ff4e0788a4f3ffdea5149dcd84f6b41fe2d Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/cohere.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fal_ai.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fal_ai.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..968cb2e71fc803f7312d140cc7d6735633a3b29d Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fal_ai.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fireworks_ai.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fireworks_ai.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b99647428f8296e87b0d88c032d31cfd0338570 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/fireworks_ai.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hf_inference.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hf_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27652aeec0a796610183ac3bd2ac93e0a38607cd Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hf_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hyperbolic.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hyperbolic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..057d4a793d386ad639f3f53a380a5dbb57f3a0d9 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/hyperbolic.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/nebius.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/nebius.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72b3a5717f754be4bf3c695ea1fac71ae7d1e7d4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/nebius.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/novita.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/novita.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..528e8e95bce5de63256508238ce681bd4c919282 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/novita.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/openai.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/openai.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d233b8d90743f102fb7e9fbf557e3a27198f4d14 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/openai.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/replicate.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/replicate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d29594fb3f7868c3f1ce460450bf7bcb6d1c1cda Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/replicate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/sambanova.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/sambanova.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1472242a463e033185ec3763b9657351f954d117 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/sambanova.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/together.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/together.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd359bfcc0562501babb9b8a1d76a00697067470 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__pycache__/together.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/_common.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..afea6fe29f6c6f50894cfd7008cb2e1aa1f0b9a7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/_common.py @@ -0,0 +1,260 @@ +from functools import lru_cache +from typing import Any, Dict, Optional, Union + +from huggingface_hub import constants +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters +from huggingface_hub.utils import build_hf_headers, get_token, logging + + +logger = logging.get_logger(__name__) + +# Dev purposes only. +# If you want to try to run inference for a new model locally before it's registered on huggingface.co +# for a given Inference Provider, you can add it to the following dictionary. +HARDCODED_MODEL_INFERENCE_MAPPING: Dict[str, Dict[str, InferenceProviderMapping]] = { + # "HF model ID" => InferenceProviderMapping object initialized with "Model ID on Inference Provider's side" + # + # Example: + # "Qwen/Qwen2.5-Coder-32B-Instruct": InferenceProviderMapping(hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct", + # provider_id="Qwen2.5-Coder-32B-Instruct", + # task="conversational", + # status="live") + "cerebras": {}, + "cohere": {}, + "fal-ai": {}, + "fireworks-ai": {}, + "hf-inference": {}, + "hyperbolic": {}, + "nebius": {}, + "replicate": {}, + "sambanova": {}, + "together": {}, +} + + +def filter_none(d: Dict[str, Any]) -> Dict[str, Any]: + return {k: v for k, v in d.items() if v is not None} + + +class TaskProviderHelper: + """Base class for task-specific provider helpers.""" + + def __init__(self, provider: str, base_url: str, task: str) -> None: + self.provider = provider + self.task = task + self.base_url = base_url + + def prepare_request( + self, + *, + inputs: Any, + parameters: Dict[str, Any], + headers: Dict, + model: Optional[str], + api_key: Optional[str], + extra_payload: Optional[Dict[str, Any]] = None, + ) -> RequestParameters: + """ + Prepare the request to be sent to the provider. + + Each step (api_key, model, headers, url, payload) can be customized in subclasses. + """ + # api_key from user, or local token, or raise error + api_key = self._prepare_api_key(api_key) + + # mapped model from HF model ID + provider_mapping_info = self._prepare_mapping_info(model) + + # default HF headers + user headers (to customize in subclasses) + headers = self._prepare_headers(headers, api_key) + + # routed URL if HF token, or direct URL (to customize in '_prepare_route' in subclasses) + url = self._prepare_url(api_key, provider_mapping_info.provider_id) + + # prepare payload (to customize in subclasses) + payload = self._prepare_payload_as_dict(inputs, parameters, provider_mapping_info=provider_mapping_info) + if payload is not None: + payload = recursive_merge(payload, extra_payload or {}) + + # body data (to customize in subclasses) + data = self._prepare_payload_as_bytes(inputs, parameters, provider_mapping_info, extra_payload) + + # check if both payload and data are set and return + if payload is not None and data is not None: + raise ValueError("Both payload and data cannot be set in the same request.") + if payload is None and data is None: + raise ValueError("Either payload or data must be set in the request.") + return RequestParameters( + url=url, task=self.task, model=provider_mapping_info.provider_id, json=payload, data=data, headers=headers + ) + + def get_response( + self, + response: Union[bytes, Dict], + request_params: Optional[RequestParameters] = None, + ) -> Any: + """ + Return the response in the expected format. + + Override this method in subclasses for customized response handling.""" + return response + + def _prepare_api_key(self, api_key: Optional[str]) -> str: + """Return the API key to use for the request. + + Usually not overwritten in subclasses.""" + if api_key is None: + api_key = get_token() + if api_key is None: + raise ValueError( + f"You must provide an api_key to work with {self.provider} API or log in with `huggingface-cli login`." + ) + return api_key + + def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping: + """Return the mapped model ID to use for the request. + + Usually not overwritten in subclasses.""" + if model is None: + raise ValueError(f"Please provide an HF model ID supported by {self.provider}.") + + # hardcoded mapping for local testing + if HARDCODED_MODEL_INFERENCE_MAPPING.get(self.provider, {}).get(model): + return HARDCODED_MODEL_INFERENCE_MAPPING[self.provider][model] + + provider_mapping = _fetch_inference_provider_mapping(model).get(self.provider) + if provider_mapping is None: + raise ValueError(f"Model {model} is not supported by provider {self.provider}.") + + if provider_mapping.task != self.task: + raise ValueError( + f"Model {model} is not supported for task {self.task} and provider {self.provider}. " + f"Supported task: {provider_mapping.task}." + ) + + if provider_mapping.status == "staging": + logger.warning( + f"Model {model} is in staging mode for provider {self.provider}. Meant for test purposes only." + ) + return provider_mapping + + def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: + """Return the headers to use for the request. + + Override this method in subclasses for customized headers. + """ + return {**build_hf_headers(token=api_key), **headers} + + def _prepare_url(self, api_key: str, mapped_model: str) -> str: + """Return the URL to use for the request. + + Usually not overwritten in subclasses.""" + base_url = self._prepare_base_url(api_key) + route = self._prepare_route(mapped_model, api_key) + return f"{base_url.rstrip('/')}/{route.lstrip('/')}" + + def _prepare_base_url(self, api_key: str) -> str: + """Return the base URL to use for the request. + + Usually not overwritten in subclasses.""" + # Route to the proxy if the api_key is a HF TOKEN + if api_key.startswith("hf_"): + logger.info(f"Calling '{self.provider}' provider through Hugging Face router.") + return constants.INFERENCE_PROXY_TEMPLATE.format(provider=self.provider) + else: + logger.info(f"Calling '{self.provider}' provider directly.") + return self.base_url + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + """Return the route to use for the request. + + Override this method in subclasses for customized routes. + """ + return "" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + """Return the payload to use for the request, as a dict. + + Override this method in subclasses for customized payloads. + Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value. + """ + return None + + def _prepare_payload_as_bytes( + self, + inputs: Any, + parameters: Dict, + provider_mapping_info: InferenceProviderMapping, + extra_payload: Optional[Dict], + ) -> Optional[bytes]: + """Return the body to use for the request, as bytes. + + Override this method in subclasses for customized body data. + Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value. + """ + return None + + +class BaseConversationalTask(TaskProviderHelper): + """ + Base class for conversational (chat completion) tasks. + The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/chat + """ + + def __init__(self, provider: str, base_url: str): + super().__init__(provider=provider, base_url=base_url, task="conversational") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/v1/chat/completions" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return {"messages": inputs, **filter_none(parameters), "model": provider_mapping_info.provider_id} + + +class BaseTextGenerationTask(TaskProviderHelper): + """ + Base class for text-generation (completion) tasks. + The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/completions + """ + + def __init__(self, provider: str, base_url: str): + super().__init__(provider=provider, base_url=base_url, task="text-generation") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/v1/completions" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return {"prompt": inputs, **filter_none(parameters), "model": provider_mapping_info.provider_id} + + +@lru_cache(maxsize=None) +def _fetch_inference_provider_mapping(model: str) -> Dict: + """ + Fetch provider mappings for a model from the Hub. + """ + from huggingface_hub.hf_api import HfApi + + info = HfApi().model_info(model, expand=["inferenceProviderMapping"]) + provider_mapping = info.inference_provider_mapping + if provider_mapping is None: + raise ValueError(f"No provider mapping found for model {model}") + return provider_mapping + + +def recursive_merge(dict1: Dict, dict2: Dict) -> Dict: + return { + **dict1, + **{ + key: recursive_merge(dict1[key], value) + if (key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict)) + else value + for key, value in dict2.items() + }, + } diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/black_forest_labs.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/black_forest_labs.py new file mode 100644 index 0000000000000000000000000000000000000000..afa8ed281d8a8e94a054b83b74ec6909f623e300 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/black_forest_labs.py @@ -0,0 +1,69 @@ +import time +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none +from huggingface_hub.utils import logging +from huggingface_hub.utils._http import get_session + + +logger = logging.get_logger(__name__) + +MAX_POLLING_ATTEMPTS = 6 +POLLING_INTERVAL = 1.0 + + +class BlackForestLabsTextToImageTask(TaskProviderHelper): + def __init__(self): + super().__init__(provider="black-forest-labs", base_url="https://api.us1.bfl.ai", task="text-to-image") + + def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: + headers = super()._prepare_headers(headers, api_key) + if not api_key.startswith("hf_"): + _ = headers.pop("authorization") + headers["X-Key"] = api_key + return headers + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return f"/v1/{mapped_model}" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + parameters = filter_none(parameters) + if "num_inference_steps" in parameters: + parameters["steps"] = parameters.pop("num_inference_steps") + if "guidance_scale" in parameters: + parameters["guidance"] = parameters.pop("guidance_scale") + + return {"prompt": inputs, **parameters} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + """ + Polling mechanism for Black Forest Labs since the API is asynchronous. + """ + url = _as_dict(response).get("polling_url") + session = get_session() + for _ in range(MAX_POLLING_ATTEMPTS): + time.sleep(POLLING_INTERVAL) + + response = session.get(url, headers={"Content-Type": "application/json"}) # type: ignore + response.raise_for_status() # type: ignore + response_json: Dict = response.json() # type: ignore + status = response_json.get("status") + logger.info( + f"Polling generation result from {url}. Current status: {status}. " + f"Will retry after {POLLING_INTERVAL} seconds if not ready." + ) + + if ( + status == "Ready" + and isinstance(response_json.get("result"), dict) + and (sample_url := response_json["result"].get("sample")) + ): + image_resp = session.get(sample_url) + image_resp.raise_for_status() + return image_resp.content + + raise TimeoutError(f"Failed to get the image URL after {MAX_POLLING_ATTEMPTS} attempts.") diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cerebras.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cerebras.py new file mode 100644 index 0000000000000000000000000000000000000000..12b181583257468ce2cba12c00dc1fb855774430 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cerebras.py @@ -0,0 +1,6 @@ +from huggingface_hub.inference._providers._common import BaseConversationalTask + + +class CerebrasConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider="cerebras", base_url="https://api.cerebras.ai") diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cohere.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cohere.py new file mode 100644 index 0000000000000000000000000000000000000000..0dc35c7e6cba26877cbbb7689148cfe76a5c0b2e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cohere.py @@ -0,0 +1,15 @@ +from huggingface_hub.inference._providers._common import ( + BaseConversationalTask, +) + + +_PROVIDER = "cohere" +_BASE_URL = "https://api.cohere.com" + + +class CohereConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL) + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/compatibility/v1/chat/completions" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fal_ai.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fal_ai.py new file mode 100644 index 0000000000000000000000000000000000000000..8dd463b6b1b3adcda4f954ace535b95b009eb17a --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fal_ai.py @@ -0,0 +1,172 @@ +import base64 +import time +from abc import ABC +from typing import Any, Dict, Optional, Union +from urllib.parse import urlparse + +from huggingface_hub import constants +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none +from huggingface_hub.utils import get_session, hf_raise_for_status +from huggingface_hub.utils.logging import get_logger + + +logger = get_logger(__name__) + +# Arbitrary polling interval +_POLLING_INTERVAL = 0.5 + + +class FalAITask(TaskProviderHelper, ABC): + def __init__(self, task: str): + super().__init__(provider="fal-ai", base_url="https://fal.run", task=task) + + def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: + headers = super()._prepare_headers(headers, api_key) + if not api_key.startswith("hf_"): + headers["authorization"] = f"Key {api_key}" + return headers + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return f"/{mapped_model}" + + +class FalAIAutomaticSpeechRecognitionTask(FalAITask): + def __init__(self): + super().__init__("automatic-speech-recognition") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + if isinstance(inputs, str) and inputs.startswith(("http://", "https://")): + # If input is a URL, pass it directly + audio_url = inputs + else: + # If input is a file path, read it first + if isinstance(inputs, str): + with open(inputs, "rb") as f: + inputs = f.read() + + audio_b64 = base64.b64encode(inputs).decode() + content_type = "audio/mpeg" + audio_url = f"data:{content_type};base64,{audio_b64}" + + return {"audio_url": audio_url, **filter_none(parameters)} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + text = _as_dict(response)["text"] + if not isinstance(text, str): + raise ValueError(f"Unexpected output format from FalAI API. Expected string, got {type(text)}.") + return text + + +class FalAITextToImageTask(FalAITask): + def __init__(self): + super().__init__("text-to-image") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + payload: Dict[str, Any] = { + "prompt": inputs, + **filter_none(parameters), + } + if "width" in payload and "height" in payload: + payload["image_size"] = { + "width": payload.pop("width"), + "height": payload.pop("height"), + } + if provider_mapping_info.adapter_weights_path is not None: + lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format( + repo_id=provider_mapping_info.hf_model_id, + revision="main", + filename=provider_mapping_info.adapter_weights_path, + ) + payload["loras"] = [{"path": lora_path, "scale": 1}] + if provider_mapping_info.provider_id == "fal-ai/lora": + # little hack: fal requires the base model for stable-diffusion-based loras but not for flux-based + # See payloads in https://fal.ai/models/fal-ai/lora/api vs https://fal.ai/models/fal-ai/flux-lora/api + payload["model_name"] = "stabilityai/stable-diffusion-xl-base-1.0" + + return payload + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + url = _as_dict(response)["images"][0]["url"] + return get_session().get(url).content + + +class FalAITextToSpeechTask(FalAITask): + def __init__(self): + super().__init__("text-to-speech") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return {"text": inputs, **filter_none(parameters)} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + url = _as_dict(response)["audio"]["url"] + return get_session().get(url).content + + +class FalAITextToVideoTask(FalAITask): + def __init__(self): + super().__init__("text-to-video") + + def _prepare_base_url(self, api_key: str) -> str: + if api_key.startswith("hf_"): + return super()._prepare_base_url(api_key) + else: + logger.info(f"Calling '{self.provider}' provider directly.") + return "https://queue.fal.run" + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + if api_key.startswith("hf_"): + # Use the queue subdomain for HF routing + return f"/{mapped_model}?_subdomain=queue" + return f"/{mapped_model}" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return {"prompt": inputs, **filter_none(parameters)} + + def get_response( + self, + response: Union[bytes, Dict], + request_params: Optional[RequestParameters] = None, + ) -> Any: + response_dict = _as_dict(response) + + request_id = response_dict.get("request_id") + if not request_id: + raise ValueError("No request ID found in the response") + if request_params is None: + raise ValueError( + "A `RequestParameters` object should be provided to get text-to-video responses with Fal AI." + ) + + # extract the base url and query params + parsed_url = urlparse(request_params.url) + # a bit hacky way to concatenate the provider name without parsing `parsed_url.path` + base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{'/fal-ai' if parsed_url.netloc == 'router.huggingface.co' else ''}" + query_param = f"?{parsed_url.query}" if parsed_url.query else "" + + # extracting the provider model id for status and result urls + # from the response as it might be different from the mapped model in `request_params.url` + model_id = urlparse(response_dict.get("response_url")).path + status_url = f"{base_url}{str(model_id)}/status{query_param}" + result_url = f"{base_url}{str(model_id)}{query_param}" + + status = response_dict.get("status") + logger.info("Generating the video.. this can take several minutes.") + while status != "COMPLETED": + time.sleep(_POLLING_INTERVAL) + status_response = get_session().get(status_url, headers=request_params.headers) + hf_raise_for_status(status_response) + status = status_response.json().get("status") + + response = get_session().get(result_url, headers=request_params.headers).json() + url = _as_dict(response)["video"]["url"] + return get_session().get(url).content diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fireworks_ai.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fireworks_ai.py new file mode 100644 index 0000000000000000000000000000000000000000..9fc9aba806a39154dde185bd84dff20e93203c9a --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fireworks_ai.py @@ -0,0 +1,9 @@ +from ._common import BaseConversationalTask + + +class FireworksAIConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/inference/v1/chat/completions" diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hf_inference.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hf_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..7923567be3297b0c4019016f3092f3640021a98c --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hf_inference.py @@ -0,0 +1,189 @@ +import json +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, Optional, Union + +from huggingface_hub import constants +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _b64_encode, _bytes_to_dict, _open_as_binary +from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none +from huggingface_hub.utils import build_hf_headers, get_session, get_token, hf_raise_for_status + + +class HFInferenceTask(TaskProviderHelper): + """Base class for HF Inference API tasks.""" + + def __init__(self, task: str): + super().__init__( + provider="hf-inference", + base_url=constants.INFERENCE_PROXY_TEMPLATE.format(provider="hf-inference"), + task=task, + ) + + def _prepare_api_key(self, api_key: Optional[str]) -> str: + # special case: for HF Inference we allow not providing an API key + return api_key or get_token() # type: ignore[return-value] + + def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping: + if model is not None and model.startswith(("http://", "https://")): + return InferenceProviderMapping(providerId=model, hf_model_id=model, task=self.task, status="live") + model_id = model if model is not None else _fetch_recommended_models().get(self.task) + if model_id is None: + raise ValueError( + f"Task {self.task} has no recommended model for HF Inference. Please specify a model" + " explicitly. Visit https://huggingface.co/tasks for more info." + ) + _check_supported_task(model_id, self.task) + return InferenceProviderMapping(providerId=model_id, hf_model_id=model_id, task=self.task, status="live") + + def _prepare_url(self, api_key: str, mapped_model: str) -> str: + # hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment) + if mapped_model.startswith(("http://", "https://")): + return mapped_model + return ( + # Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks. + f"{self.base_url}/models/{mapped_model}/pipeline/{self.task}" + if self.task in ("feature-extraction", "sentence-similarity") + # Otherwise, we use the default endpoint + else f"{self.base_url}/models/{mapped_model}" + ) + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + if isinstance(inputs, bytes): + raise ValueError(f"Unexpected binary input for task {self.task}.") + if isinstance(inputs, Path): + raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})") + return {"inputs": inputs, "parameters": filter_none(parameters)} + + +class HFInferenceBinaryInputTask(HFInferenceTask): + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return None + + def _prepare_payload_as_bytes( + self, + inputs: Any, + parameters: Dict, + provider_mapping_info: InferenceProviderMapping, + extra_payload: Optional[Dict], + ) -> Optional[bytes]: + parameters = filter_none({k: v for k, v in parameters.items() if v is not None}) + extra_payload = extra_payload or {} + has_parameters = len(parameters) > 0 or len(extra_payload) > 0 + + # Raise if not a binary object or a local path or a URL. + if not isinstance(inputs, (bytes, Path)) and not isinstance(inputs, str): + raise ValueError(f"Expected binary inputs or a local path or a URL. Got {inputs}") + + # Send inputs as raw content when no parameters are provided + if not has_parameters: + with _open_as_binary(inputs) as data: + data_as_bytes = data if isinstance(data, bytes) else data.read() + return data_as_bytes + + # Otherwise encode as b64 + return json.dumps({"inputs": _b64_encode(inputs), "parameters": parameters, **extra_payload}).encode("utf-8") + + +class HFInferenceConversational(HFInferenceTask): + def __init__(self): + super().__init__("conversational") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + mapped_model = provider_mapping_info.provider_id + payload_model = parameters.get("model") or mapped_model + + if payload_model is None or payload_model.startswith(("http://", "https://")): + payload_model = "dummy" + + return {**filter_none(parameters), "model": payload_model, "messages": inputs} + + def _prepare_url(self, api_key: str, mapped_model: str) -> str: + base_url = ( + mapped_model + if mapped_model.startswith(("http://", "https://")) + else f"{constants.INFERENCE_PROXY_TEMPLATE.format(provider='hf-inference')}/models/{mapped_model}" + ) + return _build_chat_completion_url(base_url) + + +def _build_chat_completion_url(model_url: str) -> str: + # Strip trailing / + model_url = model_url.rstrip("/") + + # Append /chat/completions if not already present + if model_url.endswith("/v1"): + model_url += "/chat/completions" + + # Append /v1/chat/completions if not already present + if not model_url.endswith("/chat/completions"): + model_url += "/v1/chat/completions" + + return model_url + + +@lru_cache(maxsize=1) +def _fetch_recommended_models() -> Dict[str, Optional[str]]: + response = get_session().get(f"{constants.ENDPOINT}/api/tasks", headers=build_hf_headers()) + hf_raise_for_status(response) + return {task: next(iter(details["widgetModels"]), None) for task, details in response.json().items()} + + +@lru_cache(maxsize=None) +def _check_supported_task(model: str, task: str) -> None: + from huggingface_hub.hf_api import HfApi + + model_info = HfApi().model_info(model) + pipeline_tag = model_info.pipeline_tag + tags = model_info.tags or [] + is_conversational = "conversational" in tags + if task in ("text-generation", "conversational"): + if pipeline_tag == "text-generation": + # text-generation + conversational tag -> both tasks allowed + if is_conversational: + return + # text-generation without conversational tag -> only text-generation allowed + if task == "text-generation": + return + raise ValueError(f"Model '{model}' doesn't support task '{task}'.") + + if pipeline_tag == "text2text-generation": + if task == "text-generation": + return + raise ValueError(f"Model '{model}' doesn't support task '{task}'.") + + if pipeline_tag == "image-text-to-text": + if is_conversational and task == "conversational": + return # Only conversational allowed if tagged as conversational + raise ValueError("Non-conversational image-text-to-text task is not supported.") + + if ( + task in ("feature-extraction", "sentence-similarity") + and pipeline_tag in ("feature-extraction", "sentence-similarity") + and task in tags + ): + # feature-extraction and sentence-similarity are interchangeable for HF Inference + return + + # For all other tasks, just check pipeline tag + if pipeline_tag != task: + raise ValueError( + f"Model '{model}' doesn't support task '{task}'. Supported tasks: '{pipeline_tag}', got: '{task}'" + ) + return + + +class HFInferenceFeatureExtractionTask(HFInferenceTask): + def __init__(self): + super().__init__("feature-extraction") + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + if isinstance(response, bytes): + return _bytes_to_dict(response) + return response diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hyperbolic.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hyperbolic.py new file mode 100644 index 0000000000000000000000000000000000000000..6dcb14cc275f6b80db5643361b9dfd3cbf8d91a2 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hyperbolic.py @@ -0,0 +1,47 @@ +import base64 +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none + + +class HyperbolicTextToImageTask(TaskProviderHelper): + def __init__(self): + super().__init__(provider="hyperbolic", base_url="https://api.hyperbolic.xyz", task="text-to-image") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/v1/images/generations" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + mapped_model = provider_mapping_info.provider_id + parameters = filter_none(parameters) + if "num_inference_steps" in parameters: + parameters["steps"] = parameters.pop("num_inference_steps") + if "guidance_scale" in parameters: + parameters["cfg_scale"] = parameters.pop("guidance_scale") + # For Hyperbolic, the width and height are required parameters + if "width" not in parameters: + parameters["width"] = 512 + if "height" not in parameters: + parameters["height"] = 512 + return {"prompt": inputs, "model_name": mapped_model, **parameters} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + response_dict = _as_dict(response) + return base64.b64decode(response_dict["images"][0]["image"]) + + +class HyperbolicTextGenerationTask(BaseConversationalTask): + """ + Special case for Hyperbolic, where text-generation task is handled as a conversational task. + """ + + def __init__(self, task: str): + super().__init__( + provider="hyperbolic", + base_url="https://api.hyperbolic.xyz", + ) + self.task = task diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/nebius.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/nebius.py new file mode 100644 index 0000000000000000000000000000000000000000..8593872a817136fc00665c40dca62e389441bb29 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/nebius.py @@ -0,0 +1,55 @@ +import base64 +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import ( + BaseConversationalTask, + BaseTextGenerationTask, + TaskProviderHelper, + filter_none, +) + + +class NebiusTextGenerationTask(BaseTextGenerationTask): + def __init__(self): + super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai") + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + output = _as_dict(response)["choices"][0] + return { + "generated_text": output["text"], + "details": { + "finish_reason": output.get("finish_reason"), + "seed": output.get("seed"), + }, + } + + +class NebiusConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai") + + +class NebiusTextToImageTask(TaskProviderHelper): + def __init__(self): + super().__init__(task="text-to-image", provider="nebius", base_url="https://api.studio.nebius.ai") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/v1/images/generations" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + mapped_model = provider_mapping_info.provider_id + parameters = filter_none(parameters) + if "guidance_scale" in parameters: + parameters.pop("guidance_scale") + if parameters.get("response_format") not in ("b64_json", "url"): + parameters["response_format"] = "b64_json" + + return {"prompt": inputs, **parameters, "model": mapped_model} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + response_dict = _as_dict(response) + return base64.b64decode(response_dict["data"][0]["b64_json"]) diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/novita.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/novita.py new file mode 100644 index 0000000000000000000000000000000000000000..44adc9017b456f487513cde251086075d84b69f0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/novita.py @@ -0,0 +1,69 @@ +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import ( + BaseConversationalTask, + BaseTextGenerationTask, + TaskProviderHelper, + filter_none, +) +from huggingface_hub.utils import get_session + + +_PROVIDER = "novita" +_BASE_URL = "https://api.novita.ai" + + +class NovitaTextGenerationTask(BaseTextGenerationTask): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL) + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + # there is no v1/ route for novita + return "/v3/openai/completions" + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + output = _as_dict(response)["choices"][0] + return { + "generated_text": output["text"], + "details": { + "finish_reason": output.get("finish_reason"), + "seed": output.get("seed"), + }, + } + + +class NovitaConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL) + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + # there is no v1/ route for novita + return "/v3/openai/chat/completions" + + +class NovitaTextToVideoTask(TaskProviderHelper): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task="text-to-video") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return f"/v3/hf/{mapped_model}" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + return {"prompt": inputs, **filter_none(parameters)} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + response_dict = _as_dict(response) + if not ( + isinstance(response_dict, dict) + and "video" in response_dict + and isinstance(response_dict["video"], dict) + and "video_url" in response_dict["video"] + ): + raise ValueError("Expected response format: { 'video': { 'video_url': string } }") + + video_url = response_dict["video"]["video_url"] + return get_session().get(video_url).content diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/openai.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/openai.py new file mode 100644 index 0000000000000000000000000000000000000000..4ea95f1643b9df9a795f35d6b0600e176089cde7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/openai.py @@ -0,0 +1,23 @@ +from typing import Optional + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._providers._common import BaseConversationalTask + + +class OpenAIConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider="openai", base_url="https://api.openai.com") + + def _prepare_api_key(self, api_key: Optional[str]) -> str: + if api_key is None: + raise ValueError("You must provide an api_key to work with OpenAI API.") + if api_key.startswith("hf_"): + raise ValueError( + "OpenAI provider is not available through Hugging Face routing, please use your own OpenAI API key." + ) + return api_key + + def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping: + if model is None: + raise ValueError("Please provide an OpenAI model ID, e.g. `gpt-4o` or `o1`.") + return InferenceProviderMapping(providerId=model, task="conversational", status="live", hf_model_id=model) diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/replicate.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..2ba312764735e289ebb6add72577f4a948a18dc3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/replicate.py @@ -0,0 +1,72 @@ +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none +from huggingface_hub.utils import get_session + + +_PROVIDER = "replicate" +_BASE_URL = "https://api.replicate.com" + + +class ReplicateTask(TaskProviderHelper): + def __init__(self, task: str): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task) + + def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: + headers = super()._prepare_headers(headers, api_key) + headers["Prefer"] = "wait" + return headers + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + if ":" in mapped_model: + return "/v1/predictions" + return f"/v1/models/{mapped_model}/predictions" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + mapped_model = provider_mapping_info.provider_id + payload: Dict[str, Any] = {"input": {"prompt": inputs, **filter_none(parameters)}} + if ":" in mapped_model: + version = mapped_model.split(":", 1)[1] + payload["version"] = version + return payload + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + response_dict = _as_dict(response) + if response_dict.get("output") is None: + raise TimeoutError( + f"Inference request timed out after 60 seconds. No output generated for model {response_dict.get('model')}" + "The model might be in cold state or starting up. Please try again later." + ) + output_url = ( + response_dict["output"] if isinstance(response_dict["output"], str) else response_dict["output"][0] + ) + return get_session().get(output_url).content + + +class ReplicateTextToImageTask(ReplicateTask): + def __init__(self): + super().__init__("text-to-image") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment] + if provider_mapping_info.adapter_weights_path is not None: + payload["input"]["lora_weights"] = f"https://huggingface.co/{provider_mapping_info.hf_model_id}" + return payload + + +class ReplicateTextToSpeechTask(ReplicateTask): + def __init__(self): + super().__init__("text-to-speech") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment] + payload["input"]["text"] = payload["input"].pop("prompt") # rename "prompt" to "text" for TTS + return payload diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/sambanova.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/sambanova.py new file mode 100644 index 0000000000000000000000000000000000000000..92bc95daa44f84ec233622550c672b487ed599b9 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/sambanova.py @@ -0,0 +1,28 @@ +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none + + +class SambanovaConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider="sambanova", base_url="https://api.sambanova.ai") + + +class SambanovaFeatureExtractionTask(TaskProviderHelper): + def __init__(self): + super().__init__(provider="sambanova", base_url="https://api.sambanova.ai", task="feature-extraction") + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + return "/v1/embeddings" + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + parameters = filter_none(parameters) + return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + embeddings = _as_dict(response)["data"] + return [embedding["embedding"] for embedding in embeddings] diff --git a/lib/python3.12/site-packages/huggingface_hub/inference/_providers/together.py b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/together.py new file mode 100644 index 0000000000000000000000000000000000000000..b27e3329383517faf586e566578f9d82483563d0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference/_providers/together.py @@ -0,0 +1,73 @@ +import base64 +from abc import ABC +from typing import Any, Dict, Optional, Union + +from huggingface_hub.hf_api import InferenceProviderMapping +from huggingface_hub.inference._common import RequestParameters, _as_dict +from huggingface_hub.inference._providers._common import ( + BaseConversationalTask, + BaseTextGenerationTask, + TaskProviderHelper, + filter_none, +) + + +_PROVIDER = "together" +_BASE_URL = "https://api.together.xyz" + + +class TogetherTask(TaskProviderHelper, ABC): + """Base class for Together API tasks.""" + + def __init__(self, task: str): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task) + + def _prepare_route(self, mapped_model: str, api_key: str) -> str: + if self.task == "text-to-image": + return "/v1/images/generations" + elif self.task == "conversational": + return "/v1/chat/completions" + elif self.task == "text-generation": + return "/v1/completions" + raise ValueError(f"Unsupported task '{self.task}' for Together API.") + + +class TogetherTextGenerationTask(BaseTextGenerationTask): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL) + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + output = _as_dict(response)["choices"][0] + return { + "generated_text": output["text"], + "details": { + "finish_reason": output.get("finish_reason"), + "seed": output.get("seed"), + }, + } + + +class TogetherConversationalTask(BaseConversationalTask): + def __init__(self): + super().__init__(provider=_PROVIDER, base_url=_BASE_URL) + + +class TogetherTextToImageTask(TogetherTask): + def __init__(self): + super().__init__("text-to-image") + + def _prepare_payload_as_dict( + self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping + ) -> Optional[Dict]: + mapped_model = provider_mapping_info.provider_id + parameters = filter_none(parameters) + if "num_inference_steps" in parameters: + parameters["steps"] = parameters.pop("num_inference_steps") + if "guidance_scale" in parameters: + parameters["guidance"] = parameters.pop("guidance_scale") + + return {"prompt": inputs, "response_format": "base64", **parameters, "model": mapped_model} + + def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any: + response_dict = _as_dict(response) + return base64.b64decode(response_dict["data"][0]["b64_json"]) diff --git a/lib/python3.12/site-packages/huggingface_hub/inference_api.py b/lib/python3.12/site-packages/huggingface_hub/inference_api.py new file mode 100644 index 0000000000000000000000000000000000000000..f895fcc61c3867838b013ecd3f6789cbc010b5b3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/inference_api.py @@ -0,0 +1,217 @@ +import io +from typing import Any, Dict, List, Optional, Union + +from . import constants +from .hf_api import HfApi +from .utils import build_hf_headers, get_session, is_pillow_available, logging, validate_hf_hub_args +from .utils._deprecation import _deprecate_method + + +logger = logging.get_logger(__name__) + + +ALL_TASKS = [ + # NLP + "text-classification", + "token-classification", + "table-question-answering", + "question-answering", + "zero-shot-classification", + "translation", + "summarization", + "conversational", + "feature-extraction", + "text-generation", + "text2text-generation", + "fill-mask", + "sentence-similarity", + # Audio + "text-to-speech", + "automatic-speech-recognition", + "audio-to-audio", + "audio-classification", + "voice-activity-detection", + # Computer vision + "image-classification", + "object-detection", + "image-segmentation", + "text-to-image", + "image-to-image", + # Others + "tabular-classification", + "tabular-regression", +] + + +class InferenceApi: + """Client to configure requests and make calls to the HuggingFace Inference API. + + Example: + + ```python + >>> from huggingface_hub.inference_api import InferenceApi + + >>> # Mask-fill example + >>> inference = InferenceApi("bert-base-uncased") + >>> inference(inputs="The goal of life is [MASK].") + [{'sequence': 'the goal of life is life.', 'score': 0.10933292657136917, 'token': 2166, 'token_str': 'life'}] + + >>> # Question Answering example + >>> inference = InferenceApi("deepset/roberta-base-squad2") + >>> inputs = { + ... "question": "What's my name?", + ... "context": "My name is Clara and I live in Berkeley.", + ... } + >>> inference(inputs) + {'score': 0.9326569437980652, 'start': 11, 'end': 16, 'answer': 'Clara'} + + >>> # Zero-shot example + >>> inference = InferenceApi("typeform/distilbert-base-uncased-mnli") + >>> inputs = "Hi, I recently bought a device from your company but it is not working as advertised and I would like to get reimbursed!" + >>> params = {"candidate_labels": ["refund", "legal", "faq"]} + >>> inference(inputs, params) + {'sequence': 'Hi, I recently bought a device from your company but it is not working as advertised and I would like to get reimbursed!', 'labels': ['refund', 'faq', 'legal'], 'scores': [0.9378499388694763, 0.04914155602455139, 0.013008488342165947]} + + >>> # Overriding configured task + >>> inference = InferenceApi("bert-base-uncased", task="feature-extraction") + + >>> # Text-to-image + >>> inference = InferenceApi("stabilityai/stable-diffusion-2-1") + >>> inference("cat") + + + >>> # Return as raw response to parse the output yourself + >>> inference = InferenceApi("mio/amadeus") + >>> response = inference("hello world", raw_response=True) + >>> response.headers + {"Content-Type": "audio/flac", ...} + >>> response.content # raw bytes from server + b'(...)' + ``` + """ + + @validate_hf_hub_args + @_deprecate_method( + version="1.0", + message=( + "`InferenceApi` client is deprecated in favor of the more feature-complete `InferenceClient`. Check out" + " this guide to learn how to convert your script to use it:" + " https://huggingface.co/docs/huggingface_hub/guides/inference#legacy-inferenceapi-client." + ), + ) + def __init__( + self, + repo_id: str, + task: Optional[str] = None, + token: Optional[str] = None, + gpu: bool = False, + ): + """Inits headers and API call information. + + Args: + repo_id (``str``): + Id of repository (e.g. `user/bert-base-uncased`). + task (``str``, `optional`, defaults ``None``): + Whether to force a task instead of using task specified in the + repository. + token (`str`, `optional`): + The API token to use as HTTP bearer authorization. This is not + the authentication token. You can find the token in + https://huggingface.co/settings/token. Alternatively, you can + find both your organizations and personal API tokens using + `HfApi().whoami(token)`. + gpu (`bool`, `optional`, defaults `False`): + Whether to use GPU instead of CPU for inference(requires Startup + plan at least). + """ + self.options = {"wait_for_model": True, "use_gpu": gpu} + self.headers = build_hf_headers(token=token) + + # Configure task + model_info = HfApi(token=token).model_info(repo_id=repo_id) + if not model_info.pipeline_tag and not task: + raise ValueError( + "Task not specified in the repository. Please add it to the model card" + " using pipeline_tag" + " (https://huggingface.co/docs#how-is-a-models-type-of-inference-api-and-widget-determined)" + ) + + if task and task != model_info.pipeline_tag: + if task not in ALL_TASKS: + raise ValueError(f"Invalid task {task}. Make sure it's valid.") + + logger.warning( + "You're using a different task than the one specified in the" + " repository. Be sure to know what you're doing :)" + ) + self.task = task + else: + assert model_info.pipeline_tag is not None, "Pipeline tag cannot be None" + self.task = model_info.pipeline_tag + + self.api_url = f"{constants.INFERENCE_ENDPOINT}/pipeline/{self.task}/{repo_id}" + + def __repr__(self): + # Do not add headers to repr to avoid leaking token. + return f"InferenceAPI(api_url='{self.api_url}', task='{self.task}', options={self.options})" + + def __call__( + self, + inputs: Optional[Union[str, Dict, List[str], List[List[str]]]] = None, + params: Optional[Dict] = None, + data: Optional[bytes] = None, + raw_response: bool = False, + ) -> Any: + """Make a call to the Inference API. + + Args: + inputs (`str` or `Dict` or `List[str]` or `List[List[str]]`, *optional*): + Inputs for the prediction. + params (`Dict`, *optional*): + Additional parameters for the models. Will be sent as `parameters` in the + payload. + data (`bytes`, *optional*): + Bytes content of the request. In this case, leave `inputs` and `params` empty. + raw_response (`bool`, defaults to `False`): + If `True`, the raw `Response` object is returned. You can parse its content + as preferred. By default, the content is parsed into a more practical format + (json dictionary or PIL Image for example). + """ + # Build payload + payload: Dict[str, Any] = { + "options": self.options, + } + if inputs: + payload["inputs"] = inputs + if params: + payload["parameters"] = params + + # Make API call + response = get_session().post(self.api_url, headers=self.headers, json=payload, data=data) + + # Let the user handle the response + if raw_response: + return response + + # By default, parse the response for the user. + content_type = response.headers.get("Content-Type") or "" + if content_type.startswith("image"): + if not is_pillow_available(): + raise ImportError( + f"Task '{self.task}' returned as image but Pillow is not installed." + " Please install it (`pip install Pillow`) or pass" + " `raw_response=True` to get the raw `Response` object and parse" + " the image by yourself." + ) + + from PIL import Image + + return Image.open(io.BytesIO(response.content)) + elif content_type == "application/json": + return response.json() + else: + raise NotImplementedError( + f"{content_type} output type is not implemented yet. You can pass" + " `raw_response=True` to get the raw `Response` object and parse the" + " output by yourself." + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/keras_mixin.py b/lib/python3.12/site-packages/huggingface_hub/keras_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..e1c7ad503e2f663b1f36da4c2ef5c954fde35606 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/keras_mixin.py @@ -0,0 +1,500 @@ +import collections.abc as collections +import json +import os +import warnings +from functools import wraps +from pathlib import Path +from shutil import copytree +from typing import Any, Dict, List, Optional, Union + +from huggingface_hub import ModelHubMixin, snapshot_download +from huggingface_hub.utils import ( + get_tf_version, + is_graphviz_available, + is_pydot_available, + is_tf_available, + yaml_dump, +) + +from . import constants +from .hf_api import HfApi +from .utils import SoftTemporaryDirectory, logging, validate_hf_hub_args +from .utils._typing import CallableT + + +logger = logging.get_logger(__name__) + +keras = None +if is_tf_available(): + # Depending on which version of TensorFlow is installed, we need to import + # keras from the correct location. + # See https://github.com/tensorflow/tensorflow/releases/tag/v2.16.1. + # Note: saving a keras model only works with Keras<3.0. + try: + import tf_keras as keras # type: ignore + except ImportError: + import tensorflow as tf # type: ignore + + keras = tf.keras + + +def _requires_keras_2_model(fn: CallableT) -> CallableT: + # Wrapper to raise if user tries to save a Keras 3.x model + @wraps(fn) + def _inner(model, *args, **kwargs): + if not hasattr(model, "history"): # hacky way to check if model is Keras 2.x + raise NotImplementedError( + f"Cannot use '{fn.__name__}': Keras 3.x is not supported." + " Please save models manually and upload them using `upload_folder` or `huggingface-cli upload`." + ) + return fn(model, *args, **kwargs) + + return _inner # type: ignore [return-value] + + +def _flatten_dict(dictionary, parent_key=""): + """Flatten a nested dictionary. + Reference: https://stackoverflow.com/a/6027615/10319735 + + Args: + dictionary (`dict`): + The nested dictionary to be flattened. + parent_key (`str`): + The parent key to be prefixed to the children keys. + Necessary for recursing over the nested dictionary. + + Returns: + The flattened dictionary. + """ + items = [] + for key, value in dictionary.items(): + new_key = f"{parent_key}.{key}" if parent_key else key + if isinstance(value, collections.MutableMapping): + items.extend( + _flatten_dict( + value, + new_key, + ).items() + ) + else: + items.append((new_key, value)) + return dict(items) + + +def _create_hyperparameter_table(model): + """Parse hyperparameter dictionary into a markdown table.""" + table = None + if model.optimizer is not None: + optimizer_params = model.optimizer.get_config() + # flatten the configuration + optimizer_params = _flatten_dict(optimizer_params) + optimizer_params["training_precision"] = keras.mixed_precision.global_policy().name + table = "| Hyperparameters | Value |\n| :-- | :-- |\n" + for key, value in optimizer_params.items(): + table += f"| {key} | {value} |\n" + return table + + +def _plot_network(model, save_directory): + keras.utils.plot_model( + model, + to_file=f"{save_directory}/model.png", + show_shapes=False, + show_dtype=False, + show_layer_names=True, + rankdir="TB", + expand_nested=False, + dpi=96, + layer_range=None, + ) + + +def _create_model_card( + model, + repo_dir: Path, + plot_model: bool = True, + metadata: Optional[dict] = None, +): + """ + Creates a model card for the repository. + + Do not overwrite an existing README.md file. + """ + readme_path = repo_dir / "README.md" + if readme_path.exists(): + return + + hyperparameters = _create_hyperparameter_table(model) + if plot_model and is_graphviz_available() and is_pydot_available(): + _plot_network(model, repo_dir) + if metadata is None: + metadata = {} + metadata["library_name"] = "keras" + model_card: str = "---\n" + model_card += yaml_dump(metadata, default_flow_style=False) + model_card += "---\n" + model_card += "\n## Model description\n\nMore information needed\n" + model_card += "\n## Intended uses & limitations\n\nMore information needed\n" + model_card += "\n## Training and evaluation data\n\nMore information needed\n" + if hyperparameters is not None: + model_card += "\n## Training procedure\n" + model_card += "\n### Training hyperparameters\n" + model_card += "\nThe following hyperparameters were used during training:\n\n" + model_card += hyperparameters + model_card += "\n" + if plot_model and os.path.exists(f"{repo_dir}/model.png"): + model_card += "\n ## Model Plot\n" + model_card += "\n
" + model_card += "\nView Model Plot\n" + path_to_plot = "./model.png" + model_card += f"\n![Model Image]({path_to_plot})\n" + model_card += "\n
" + + readme_path.write_text(model_card) + + +@_requires_keras_2_model +def save_pretrained_keras( + model, + save_directory: Union[str, Path], + config: Optional[Dict[str, Any]] = None, + include_optimizer: bool = False, + plot_model: bool = True, + tags: Optional[Union[list, str]] = None, + **model_save_kwargs, +): + """ + Saves a Keras model to save_directory in SavedModel format. Use this if + you're using the Functional or Sequential APIs. + + Args: + model (`Keras.Model`): + The [Keras + model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) + you'd like to save. The model must be compiled and built. + save_directory (`str` or `Path`): + Specify directory in which you want to save the Keras model. + config (`dict`, *optional*): + Configuration object to be saved alongside the model weights. + include_optimizer(`bool`, *optional*, defaults to `False`): + Whether or not to include optimizer in serialization. + plot_model (`bool`, *optional*, defaults to `True`): + Setting this to `True` will plot the model and put it in the model + card. Requires graphviz and pydot to be installed. + tags (Union[`str`,`list`], *optional*): + List of tags that are related to model or string of a single tag. See example tags + [here](https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1). + model_save_kwargs(`dict`, *optional*): + model_save_kwargs will be passed to + [`tf.keras.models.save_model()`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model). + """ + if keras is None: + raise ImportError("Called a Tensorflow-specific function but could not import it.") + + if not model.built: + raise ValueError("Model should be built before trying to save") + + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + + # saving config + if config: + if not isinstance(config, dict): + raise RuntimeError(f"Provided config to save_pretrained_keras should be a dict. Got: '{type(config)}'") + + with (save_directory / constants.CONFIG_NAME).open("w") as f: + json.dump(config, f) + + metadata = {} + if isinstance(tags, list): + metadata["tags"] = tags + elif isinstance(tags, str): + metadata["tags"] = [tags] + + task_name = model_save_kwargs.pop("task_name", None) + if task_name is not None: + warnings.warn( + "`task_name` input argument is deprecated. Pass `tags` instead.", + FutureWarning, + ) + if "tags" in metadata: + metadata["tags"].append(task_name) + else: + metadata["tags"] = [task_name] + + if model.history is not None: + if model.history.history != {}: + path = save_directory / "history.json" + if path.exists(): + warnings.warn( + "`history.json` file already exists, it will be overwritten by the history of this version.", + UserWarning, + ) + with path.open("w", encoding="utf-8") as f: + json.dump(model.history.history, f, indent=2, sort_keys=True) + + _create_model_card(model, save_directory, plot_model, metadata) + keras.models.save_model(model, save_directory, include_optimizer=include_optimizer, **model_save_kwargs) + + +def from_pretrained_keras(*args, **kwargs) -> "KerasModelHubMixin": + r""" + Instantiate a pretrained Keras model from a pre-trained model from the Hub. + The model is expected to be in `SavedModel` format. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + - A string, the `model id` of a pretrained model hosted inside a + model repo on huggingface.co. Valid model ids can be located + at the root-level, like `bert-base-uncased`, or namespaced + under a user or organization name, like + `dbmdz/bert-base-german-cased`. + - You can add `revision` by appending `@` at the end of model_id + simply like this: `dbmdz/bert-base-german-cased@main` Revision + is the specific model version to use. It can be a branch name, + a tag name, or a commit id, since we use a git-based system + for storing models and other artifacts on huggingface.co, so + `revision` can be any identifier allowed by git. + - A path to a `directory` containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., + `./my_model_directory/`. + - `None` if you are both providing the configuration and state + dictionary (resp. with keyword arguments `config` and + `state_dict`). + force_download (`bool`, *optional*, defaults to `False`): + Whether to force the (re-)download of the model weights and + configuration files, overriding the cached versions if they exist. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., + `{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}`. The + proxies are used on each request. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If + `True`, will use the token generated when running `transformers-cli + login` (stored in `~/.huggingface`). + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory in which a downloaded pretrained model + configuration should be cached if the standard cache should not be + used. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether to only look at local files (i.e., do not try to download + the model). + model_kwargs (`Dict`, *optional*): + model_kwargs will be passed to the model during initialization + + + + Passing `token=True` is required when you want to use a private + model. + + + """ + return KerasModelHubMixin.from_pretrained(*args, **kwargs) + + +@validate_hf_hub_args +@_requires_keras_2_model +def push_to_hub_keras( + model, + repo_id: str, + *, + config: Optional[dict] = None, + commit_message: str = "Push Keras model using huggingface_hub.", + private: Optional[bool] = None, + api_endpoint: Optional[str] = None, + token: Optional[str] = None, + branch: Optional[str] = None, + create_pr: Optional[bool] = None, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + delete_patterns: Optional[Union[List[str], str]] = None, + log_dir: Optional[str] = None, + include_optimizer: bool = False, + tags: Optional[Union[list, str]] = None, + plot_model: bool = True, + **model_save_kwargs, +): + """ + Upload model checkpoint to the Hub. + + Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use + `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more + details. + + Args: + model (`Keras.Model`): + The [Keras model](`https://www.tensorflow.org/api_docs/python/tf/keras/Model`) you'd like to push to the + Hub. The model must be compiled and built. + repo_id (`str`): + ID of the repository to push to (example: `"username/my-model"`). + commit_message (`str`, *optional*, defaults to "Add Keras model"): + Message to commit while pushing. + private (`bool`, *optional*): + Whether the repository created should be private. + If `None` (default), the repo will be public unless the organization's default is private. + api_endpoint (`str`, *optional*): + The API endpoint to use when pushing the model to the hub. + token (`str`, *optional*): + The token to use as HTTP bearer authorization for remote files. If + not set, will use the token set when logging in with + `huggingface-cli login` (stored in `~/.huggingface`). + branch (`str`, *optional*): + The git branch on which to push the model. This defaults to + the default branch as specified in your repository, which + defaults to `"main"`. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request from `branch` with that commit. + Defaults to `False`. + config (`dict`, *optional*): + Configuration object to be saved alongside the model weights. + allow_patterns (`List[str]` or `str`, *optional*): + If provided, only files matching at least one pattern are pushed. + ignore_patterns (`List[str]` or `str`, *optional*): + If provided, files matching any of the patterns are not pushed. + delete_patterns (`List[str]` or `str`, *optional*): + If provided, remote files matching any of the patterns will be deleted from the repo. + log_dir (`str`, *optional*): + TensorBoard logging directory to be pushed. The Hub automatically + hosts and displays a TensorBoard instance if log files are included + in the repository. + include_optimizer (`bool`, *optional*, defaults to `False`): + Whether or not to include optimizer during serialization. + tags (Union[`list`, `str`], *optional*): + List of tags that are related to model or string of a single tag. See example tags + [here](https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1). + plot_model (`bool`, *optional*, defaults to `True`): + Setting this to `True` will plot the model and put it in the model + card. Requires graphviz and pydot to be installed. + model_save_kwargs(`dict`, *optional*): + model_save_kwargs will be passed to + [`tf.keras.models.save_model()`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model). + + Returns: + The url of the commit of your model in the given repository. + """ + api = HfApi(endpoint=api_endpoint) + repo_id = api.create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True).repo_id + + # Push the files to the repo in a single commit + with SoftTemporaryDirectory() as tmp: + saved_path = Path(tmp) / repo_id + save_pretrained_keras( + model, + saved_path, + config=config, + include_optimizer=include_optimizer, + tags=tags, + plot_model=plot_model, + **model_save_kwargs, + ) + + # If `log_dir` provided, delete remote logs and upload new ones + if log_dir is not None: + delete_patterns = ( + [] + if delete_patterns is None + else ( + [delete_patterns] # convert `delete_patterns` to a list + if isinstance(delete_patterns, str) + else delete_patterns + ) + ) + delete_patterns.append("logs/*") + copytree(log_dir, saved_path / "logs") + + return api.upload_folder( + repo_type="model", + repo_id=repo_id, + folder_path=saved_path, + commit_message=commit_message, + token=token, + revision=branch, + create_pr=create_pr, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + delete_patterns=delete_patterns, + ) + + +class KerasModelHubMixin(ModelHubMixin): + """ + Implementation of [`ModelHubMixin`] to provide model Hub upload/download + capabilities to Keras models. + + + ```python + >>> import tensorflow as tf + >>> from huggingface_hub import KerasModelHubMixin + + + >>> class MyModel(tf.keras.Model, KerasModelHubMixin): + ... def __init__(self, **kwargs): + ... super().__init__() + ... self.config = kwargs.pop("config", None) + ... self.dummy_inputs = ... + ... self.layer = ... + + ... def call(self, *args): + ... return ... + + + >>> # Initialize and compile the model as you normally would + >>> model = MyModel() + >>> model.compile(...) + >>> # Build the graph by training it or passing dummy inputs + >>> _ = model(model.dummy_inputs) + >>> # Save model weights to local directory + >>> model.save_pretrained("my-awesome-model") + >>> # Push model weights to the Hub + >>> model.push_to_hub("my-awesome-model") + >>> # Download and initialize weights from the Hub + >>> model = MyModel.from_pretrained("username/super-cool-model") + ``` + """ + + def _save_pretrained(self, save_directory): + save_pretrained_keras(self, save_directory) + + @classmethod + def _from_pretrained( + cls, + model_id, + revision, + cache_dir, + force_download, + proxies, + resume_download, + local_files_only, + token, + config: Optional[Dict[str, Any]] = None, + **model_kwargs, + ): + """Here we just call [`from_pretrained_keras`] function so both the mixin and + functional APIs stay in sync. + + TODO - Some args above aren't used since we are calling + snapshot_download instead of hf_hub_download. + """ + if keras is None: + raise ImportError("Called a TensorFlow-specific function but could not import it.") + + # Root is either a local filepath matching model_id or a cached snapshot + if not os.path.isdir(model_id): + storage_folder = snapshot_download( + repo_id=model_id, + revision=revision, + cache_dir=cache_dir, + library_name="keras", + library_version=get_tf_version(), + ) + else: + storage_folder = model_id + + # TODO: change this in a future PR. We are not returning a KerasModelHubMixin instance here... + model = keras.models.load_model(storage_folder) + + # For now, we add a new attribute, config, to store the config loaded from the hub/a local dir. + model.config = config + + return model diff --git a/lib/python3.12/site-packages/huggingface_hub/lfs.py b/lib/python3.12/site-packages/huggingface_hub/lfs.py new file mode 100644 index 0000000000000000000000000000000000000000..c2d4f36829dfe941ce60c8b711c1cc912e8c324a --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/lfs.py @@ -0,0 +1,460 @@ +# coding=utf-8 +# Copyright 2019-present, the HuggingFace Inc. team. +# +# 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. +"""Git LFS related type definitions and utilities""" + +import inspect +import io +import re +import warnings +from dataclasses import dataclass +from math import ceil +from os.path import getsize +from pathlib import Path +from typing import TYPE_CHECKING, BinaryIO, Dict, Iterable, List, Optional, Tuple, TypedDict +from urllib.parse import unquote + +from huggingface_hub import constants + +from .utils import ( + build_hf_headers, + fix_hf_endpoint_in_url, + get_session, + hf_raise_for_status, + http_backoff, + logging, + tqdm, + validate_hf_hub_args, +) +from .utils._lfs import SliceFileObj +from .utils.sha import sha256, sha_fileobj +from .utils.tqdm import is_tqdm_disabled + + +if TYPE_CHECKING: + from ._commit_api import CommitOperationAdd + +logger = logging.get_logger(__name__) + +OID_REGEX = re.compile(r"^[0-9a-f]{40}$") + +LFS_MULTIPART_UPLOAD_COMMAND = "lfs-multipart-upload" + +LFS_HEADERS = { + "Accept": "application/vnd.git-lfs+json", + "Content-Type": "application/vnd.git-lfs+json", +} + + +@dataclass +class UploadInfo: + """ + Dataclass holding required information to determine whether a blob + should be uploaded to the hub using the LFS protocol or the regular protocol + + Args: + sha256 (`bytes`): + SHA256 hash of the blob + size (`int`): + Size in bytes of the blob + sample (`bytes`): + First 512 bytes of the blob + """ + + sha256: bytes + size: int + sample: bytes + + @classmethod + def from_path(cls, path: str): + size = getsize(path) + with io.open(path, "rb") as file: + sample = file.peek(512)[:512] + sha = sha_fileobj(file) + return cls(size=size, sha256=sha, sample=sample) + + @classmethod + def from_bytes(cls, data: bytes): + sha = sha256(data).digest() + return cls(size=len(data), sample=data[:512], sha256=sha) + + @classmethod + def from_fileobj(cls, fileobj: BinaryIO): + sample = fileobj.read(512) + fileobj.seek(0, io.SEEK_SET) + sha = sha_fileobj(fileobj) + size = fileobj.tell() + fileobj.seek(0, io.SEEK_SET) + return cls(size=size, sha256=sha, sample=sample) + + +@validate_hf_hub_args +def post_lfs_batch_info( + upload_infos: Iterable[UploadInfo], + token: Optional[str], + repo_type: str, + repo_id: str, + revision: Optional[str] = None, + endpoint: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, +) -> Tuple[List[dict], List[dict]]: + """ + Requests the LFS batch endpoint to retrieve upload instructions + + Learn more: https://github.com/git-lfs/git-lfs/blob/main/docs/api/batch.md + + Args: + upload_infos (`Iterable` of `UploadInfo`): + `UploadInfo` for the files that are being uploaded, typically obtained + from `CommitOperationAdd.upload_info` + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated + by a `/`. + revision (`str`, *optional*): + The git revision to upload to. + headers (`dict`, *optional*): + Additional headers to include in the request + + Returns: + `LfsBatchInfo`: 2-tuple: + - First element is the list of upload instructions from the server + - Second element is an list of errors, if any + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If an argument is invalid or the server response is malformed. + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + If the server returned an error. + """ + endpoint = endpoint if endpoint is not None else constants.ENDPOINT + url_prefix = "" + if repo_type in constants.REPO_TYPES_URL_PREFIXES: + url_prefix = constants.REPO_TYPES_URL_PREFIXES[repo_type] + batch_url = f"{endpoint}/{url_prefix}{repo_id}.git/info/lfs/objects/batch" + payload: Dict = { + "operation": "upload", + "transfers": ["basic", "multipart"], + "objects": [ + { + "oid": upload.sha256.hex(), + "size": upload.size, + } + for upload in upload_infos + ], + "hash_algo": "sha256", + } + if revision is not None: + payload["ref"] = {"name": unquote(revision)} # revision has been previously 'quoted' + + headers = { + **LFS_HEADERS, + **build_hf_headers(token=token), + **(headers or {}), + } + resp = get_session().post(batch_url, headers=headers, json=payload) + hf_raise_for_status(resp) + batch_info = resp.json() + + objects = batch_info.get("objects", None) + if not isinstance(objects, list): + raise ValueError("Malformed response from server") + + return ( + [_validate_batch_actions(obj) for obj in objects if "error" not in obj], + [_validate_batch_error(obj) for obj in objects if "error" in obj], + ) + + +class PayloadPartT(TypedDict): + partNumber: int + etag: str + + +class CompletionPayloadT(TypedDict): + """Payload that will be sent to the Hub when uploading multi-part.""" + + oid: str + parts: List[PayloadPartT] + + +def lfs_upload( + operation: "CommitOperationAdd", + lfs_batch_action: Dict, + token: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, + endpoint: Optional[str] = None, +) -> None: + """ + Handles uploading a given object to the Hub with the LFS protocol. + + Can be a No-op if the content of the file is already present on the hub large file storage. + + Args: + operation (`CommitOperationAdd`): + The add operation triggering this upload. + lfs_batch_action (`dict`): + Upload instructions from the LFS batch endpoint for this object. See [`~utils.lfs.post_lfs_batch_info`] for + more details. + headers (`dict`, *optional*): + Headers to include in the request, including authentication and user agent headers. + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `lfs_batch_action` is improperly formatted + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + If the upload resulted in an error + """ + # 0. If LFS file is already present, skip upload + _validate_batch_actions(lfs_batch_action) + actions = lfs_batch_action.get("actions") + if actions is None: + # The file was already uploaded + logger.debug(f"Content of file {operation.path_in_repo} is already present upstream - skipping upload") + return + + # 1. Validate server response (check required keys in dict) + upload_action = lfs_batch_action["actions"]["upload"] + _validate_lfs_action(upload_action) + verify_action = lfs_batch_action["actions"].get("verify") + if verify_action is not None: + _validate_lfs_action(verify_action) + + # 2. Upload file (either single part or multi-part) + header = upload_action.get("header", {}) + chunk_size = header.get("chunk_size") + upload_url = fix_hf_endpoint_in_url(upload_action["href"], endpoint=endpoint) + if chunk_size is not None: + try: + chunk_size = int(chunk_size) + except (ValueError, TypeError): + raise ValueError( + f"Malformed response from LFS batch endpoint: `chunk_size` should be an integer. Got '{chunk_size}'." + ) + _upload_multi_part(operation=operation, header=header, chunk_size=chunk_size, upload_url=upload_url) + else: + _upload_single_part(operation=operation, upload_url=upload_url) + + # 3. Verify upload went well + if verify_action is not None: + _validate_lfs_action(verify_action) + verify_url = fix_hf_endpoint_in_url(verify_action["href"], endpoint) + verify_resp = get_session().post( + verify_url, + headers=build_hf_headers(token=token, headers=headers), + json={"oid": operation.upload_info.sha256.hex(), "size": operation.upload_info.size}, + ) + hf_raise_for_status(verify_resp) + logger.debug(f"{operation.path_in_repo}: Upload successful") + + +def _validate_lfs_action(lfs_action: dict): + """validates response from the LFS batch endpoint""" + if not ( + isinstance(lfs_action.get("href"), str) + and (lfs_action.get("header") is None or isinstance(lfs_action.get("header"), dict)) + ): + raise ValueError("lfs_action is improperly formatted") + return lfs_action + + +def _validate_batch_actions(lfs_batch_actions: dict): + """validates response from the LFS batch endpoint""" + if not (isinstance(lfs_batch_actions.get("oid"), str) and isinstance(lfs_batch_actions.get("size"), int)): + raise ValueError("lfs_batch_actions is improperly formatted") + + upload_action = lfs_batch_actions.get("actions", {}).get("upload") + verify_action = lfs_batch_actions.get("actions", {}).get("verify") + if upload_action is not None: + _validate_lfs_action(upload_action) + if verify_action is not None: + _validate_lfs_action(verify_action) + return lfs_batch_actions + + +def _validate_batch_error(lfs_batch_error: dict): + """validates response from the LFS batch endpoint""" + if not (isinstance(lfs_batch_error.get("oid"), str) and isinstance(lfs_batch_error.get("size"), int)): + raise ValueError("lfs_batch_error is improperly formatted") + error_info = lfs_batch_error.get("error") + if not ( + isinstance(error_info, dict) + and isinstance(error_info.get("message"), str) + and isinstance(error_info.get("code"), int) + ): + raise ValueError("lfs_batch_error is improperly formatted") + return lfs_batch_error + + +def _upload_single_part(operation: "CommitOperationAdd", upload_url: str) -> None: + """ + Uploads `fileobj` as a single PUT HTTP request (basic LFS transfer protocol) + + Args: + upload_url (`str`): + The URL to PUT the file to. + fileobj: + The file-like object holding the data to upload. + + Returns: `requests.Response` + + Raises: + [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + If the upload resulted in an error. + """ + with operation.as_file(with_tqdm=True) as fileobj: + # S3 might raise a transient 500 error -> let's retry if that happens + response = http_backoff("PUT", upload_url, data=fileobj, retry_on_status_codes=(500, 502, 503, 504)) + hf_raise_for_status(response) + + +def _upload_multi_part(operation: "CommitOperationAdd", header: Dict, chunk_size: int, upload_url: str) -> None: + """ + Uploads file using HF multipart LFS transfer protocol. + """ + # 1. Get upload URLs for each part + sorted_parts_urls = _get_sorted_parts_urls(header=header, upload_info=operation.upload_info, chunk_size=chunk_size) + + # 2. Upload parts (either with hf_transfer or in pure Python) + use_hf_transfer = constants.HF_HUB_ENABLE_HF_TRANSFER + if ( + constants.HF_HUB_ENABLE_HF_TRANSFER + and not isinstance(operation.path_or_fileobj, str) + and not isinstance(operation.path_or_fileobj, Path) + ): + warnings.warn( + "hf_transfer is enabled but does not support uploading from bytes or BinaryIO, falling back to regular" + " upload" + ) + use_hf_transfer = False + + response_headers = ( + _upload_parts_hf_transfer(operation=operation, sorted_parts_urls=sorted_parts_urls, chunk_size=chunk_size) + if use_hf_transfer + else _upload_parts_iteratively(operation=operation, sorted_parts_urls=sorted_parts_urls, chunk_size=chunk_size) + ) + + # 3. Send completion request + completion_res = get_session().post( + upload_url, + json=_get_completion_payload(response_headers, operation.upload_info.sha256.hex()), + headers=LFS_HEADERS, + ) + hf_raise_for_status(completion_res) + + +def _get_sorted_parts_urls(header: Dict, upload_info: UploadInfo, chunk_size: int) -> List[str]: + sorted_part_upload_urls = [ + upload_url + for _, upload_url in sorted( + [ + (int(part_num, 10), upload_url) + for part_num, upload_url in header.items() + if part_num.isdigit() and len(part_num) > 0 + ], + key=lambda t: t[0], + ) + ] + num_parts = len(sorted_part_upload_urls) + if num_parts != ceil(upload_info.size / chunk_size): + raise ValueError("Invalid server response to upload large LFS file") + return sorted_part_upload_urls + + +def _get_completion_payload(response_headers: List[Dict], oid: str) -> CompletionPayloadT: + parts: List[PayloadPartT] = [] + for part_number, header in enumerate(response_headers): + etag = header.get("etag") + if etag is None or etag == "": + raise ValueError(f"Invalid etag (`{etag}`) returned for part {part_number + 1}") + parts.append( + { + "partNumber": part_number + 1, + "etag": etag, + } + ) + return {"oid": oid, "parts": parts} + + +def _upload_parts_iteratively( + operation: "CommitOperationAdd", sorted_parts_urls: List[str], chunk_size: int +) -> List[Dict]: + headers = [] + with operation.as_file(with_tqdm=True) as fileobj: + for part_idx, part_upload_url in enumerate(sorted_parts_urls): + with SliceFileObj( + fileobj, + seek_from=chunk_size * part_idx, + read_limit=chunk_size, + ) as fileobj_slice: + # S3 might raise a transient 500 error -> let's retry if that happens + part_upload_res = http_backoff( + "PUT", part_upload_url, data=fileobj_slice, retry_on_status_codes=(500, 502, 503, 504) + ) + hf_raise_for_status(part_upload_res) + headers.append(part_upload_res.headers) + return headers # type: ignore + + +def _upload_parts_hf_transfer( + operation: "CommitOperationAdd", sorted_parts_urls: List[str], chunk_size: int +) -> List[Dict]: + # Upload file using an external Rust-based package. Upload is faster but support less features (no progress bars). + try: + from hf_transfer import multipart_upload + except ImportError: + raise ValueError( + "Fast uploading using 'hf_transfer' is enabled (HF_HUB_ENABLE_HF_TRANSFER=1) but 'hf_transfer' package is" + " not available in your environment. Try `pip install hf_transfer`." + ) + + supports_callback = "callback" in inspect.signature(multipart_upload).parameters + if not supports_callback: + warnings.warn( + "You are using an outdated version of `hf_transfer`. Consider upgrading to latest version to enable progress bars using `pip install -U hf_transfer`." + ) + + total = operation.upload_info.size + desc = operation.path_in_repo + if len(desc) > 40: + desc = f"(…){desc[-40:]}" + + with tqdm( + unit="B", + unit_scale=True, + total=total, + initial=0, + desc=desc, + disable=is_tqdm_disabled(logger.getEffectiveLevel()), + name="huggingface_hub.lfs_upload", + ) as progress: + try: + output = multipart_upload( + file_path=operation.path_or_fileobj, + parts_urls=sorted_parts_urls, + chunk_size=chunk_size, + max_files=128, + parallel_failures=127, # could be removed + max_retries=5, + **({"callback": progress.update} if supports_callback else {}), + ) + except Exception as e: + raise RuntimeError( + "An error occurred while uploading using `hf_transfer`. Consider disabling HF_HUB_ENABLE_HF_TRANSFER for" + " better error handling." + ) from e + if not supports_callback: + progress.update(total) + return output diff --git a/lib/python3.12/site-packages/huggingface_hub/py.typed b/lib/python3.12/site-packages/huggingface_hub/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/huggingface_hub/repocard.py b/lib/python3.12/site-packages/huggingface_hub/repocard.py new file mode 100644 index 0000000000000000000000000000000000000000..83b22b2bf60a6aa52ecd1d66545fbf5fa6d45a0f --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/repocard.py @@ -0,0 +1,830 @@ +import os +import re +from pathlib import Path +from typing import Any, Dict, Literal, Optional, Type, Union + +import requests +import yaml + +from huggingface_hub.file_download import hf_hub_download +from huggingface_hub.hf_api import upload_file +from huggingface_hub.repocard_data import ( + CardData, + DatasetCardData, + EvalResult, + ModelCardData, + SpaceCardData, + eval_results_to_model_index, + model_index_to_eval_results, +) +from huggingface_hub.utils import get_session, is_jinja_available, yaml_dump + +from . import constants +from .errors import EntryNotFoundError +from .utils import SoftTemporaryDirectory, logging, validate_hf_hub_args + + +logger = logging.get_logger(__name__) + + +TEMPLATE_MODELCARD_PATH = Path(__file__).parent / "templates" / "modelcard_template.md" +TEMPLATE_DATASETCARD_PATH = Path(__file__).parent / "templates" / "datasetcard_template.md" + +# exact same regex as in the Hub server. Please keep in sync. +# See https://github.com/huggingface/moon-landing/blob/main/server/lib/ViewMarkdown.ts#L18 +REGEX_YAML_BLOCK = re.compile(r"^(\s*---[\r\n]+)([\S\s]*?)([\r\n]+---(\r\n|\n|$))") + + +class RepoCard: + card_data_class = CardData + default_template_path = TEMPLATE_MODELCARD_PATH + repo_type = "model" + + def __init__(self, content: str, ignore_metadata_errors: bool = False): + """Initialize a RepoCard from string content. The content should be a + Markdown file with a YAML block at the beginning and a Markdown body. + + Args: + content (`str`): The content of the Markdown file. + + Example: + ```python + >>> from huggingface_hub.repocard import RepoCard + >>> text = ''' + ... --- + ... language: en + ... license: mit + ... --- + ... + ... # My repo + ... ''' + >>> card = RepoCard(text) + >>> card.data.to_dict() + {'language': 'en', 'license': 'mit'} + >>> card.text + '\\n# My repo\\n' + + ``` + + Raises the following error: + + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + when the content of the repo card metadata is not a dictionary. + + + """ + + # Set the content of the RepoCard, as well as underlying .data and .text attributes. + # See the `content` property setter for more details. + self.ignore_metadata_errors = ignore_metadata_errors + self.content = content + + @property + def content(self): + """The content of the RepoCard, including the YAML block and the Markdown body.""" + line_break = _detect_line_ending(self._content) or "\n" + return f"---{line_break}{self.data.to_yaml(line_break=line_break, original_order=self._original_order)}{line_break}---{line_break}{self.text}" + + @content.setter + def content(self, content: str): + """Set the content of the RepoCard.""" + self._content = content + + match = REGEX_YAML_BLOCK.search(content) + if match: + # Metadata found in the YAML block + yaml_block = match.group(2) + self.text = content[match.end() :] + data_dict = yaml.safe_load(yaml_block) + + if data_dict is None: + data_dict = {} + + # The YAML block's data should be a dictionary + if not isinstance(data_dict, dict): + raise ValueError("repo card metadata block should be a dict") + else: + # Model card without metadata... create empty metadata + logger.warning("Repo card metadata block was not found. Setting CardData to empty.") + data_dict = {} + self.text = content + + self.data = self.card_data_class(**data_dict, ignore_metadata_errors=self.ignore_metadata_errors) + self._original_order = list(data_dict.keys()) + + def __str__(self): + return self.content + + def save(self, filepath: Union[Path, str]): + r"""Save a RepoCard to a file. + + Args: + filepath (`Union[Path, str]`): Filepath to the markdown file to save. + + Example: + ```python + >>> from huggingface_hub.repocard import RepoCard + >>> card = RepoCard("---\nlanguage: en\n---\n# This is a test repo card") + >>> card.save("/tmp/test.md") + + ``` + """ + filepath = Path(filepath) + filepath.parent.mkdir(parents=True, exist_ok=True) + # Preserve newlines as in the existing file. + with open(filepath, mode="w", newline="", encoding="utf-8") as f: + f.write(str(self)) + + @classmethod + def load( + cls, + repo_id_or_path: Union[str, Path], + repo_type: Optional[str] = None, + token: Optional[str] = None, + ignore_metadata_errors: bool = False, + ): + """Initialize a RepoCard from a Hugging Face Hub repo's README.md or a local filepath. + + Args: + repo_id_or_path (`Union[str, Path]`): + The repo ID associated with a Hugging Face Hub repo or a local filepath. + repo_type (`str`, *optional*): + The type of Hugging Face repo to push to. Defaults to None, which will use use "model". Other options + are "dataset" and "space". Not used when loading from a local filepath. If this is called from a child + class, the default value will be the child class's `repo_type`. + token (`str`, *optional*): + Authentication token, obtained with `huggingface_hub.HfApi.login` method. Will default to the stored token. + ignore_metadata_errors (`str`): + If True, errors while parsing the metadata section will be ignored. Some information might be lost during + the process. Use it at your own risk. + + Returns: + [`huggingface_hub.repocard.RepoCard`]: The RepoCard (or subclass) initialized from the repo's + README.md file or filepath. + + Example: + ```python + >>> from huggingface_hub.repocard import RepoCard + >>> card = RepoCard.load("nateraw/food") + >>> assert card.data.tags == ["generated_from_trainer", "image-classification", "pytorch"] + + ``` + """ + + if Path(repo_id_or_path).is_file(): + card_path = Path(repo_id_or_path) + elif isinstance(repo_id_or_path, str): + card_path = Path( + hf_hub_download( + repo_id_or_path, + constants.REPOCARD_NAME, + repo_type=repo_type or cls.repo_type, + token=token, + ) + ) + else: + raise ValueError(f"Cannot load RepoCard: path not found on disk ({repo_id_or_path}).") + + # Preserve newlines in the existing file. + with card_path.open(mode="r", newline="", encoding="utf-8") as f: + return cls(f.read(), ignore_metadata_errors=ignore_metadata_errors) + + def validate(self, repo_type: Optional[str] = None): + """Validates card against Hugging Face Hub's card validation logic. + Using this function requires access to the internet, so it is only called + internally by [`huggingface_hub.repocard.RepoCard.push_to_hub`]. + + Args: + repo_type (`str`, *optional*, defaults to "model"): + The type of Hugging Face repo to push to. Options are "model", "dataset", and "space". + If this function is called from a child class, the default will be the child class's `repo_type`. + + + Raises the following errors: + + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if the card fails validation checks. + - [`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError) + if the request to the Hub API fails for any other reason. + + + """ + + # If repo type is provided, otherwise, use the repo type of the card. + repo_type = repo_type or self.repo_type + + body = { + "repoType": repo_type, + "content": str(self), + } + headers = {"Accept": "text/plain"} + + try: + r = get_session().post("https://huggingface.co/api/validate-yaml", body, headers=headers) + r.raise_for_status() + except requests.exceptions.HTTPError as exc: + if r.status_code == 400: + raise ValueError(r.text) + else: + raise exc + + def push_to_hub( + self, + repo_id: str, + token: Optional[str] = None, + repo_type: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + revision: Optional[str] = None, + create_pr: Optional[bool] = None, + parent_commit: Optional[str] = None, + ): + """Push a RepoCard to a Hugging Face Hub repo. + + Args: + repo_id (`str`): + The repo ID of the Hugging Face Hub repo to push to. Example: "nateraw/food". + token (`str`, *optional*): + Authentication token, obtained with `huggingface_hub.HfApi.login` method. Will default to + the stored token. + repo_type (`str`, *optional*, defaults to "model"): + The type of Hugging Face repo to push to. Options are "model", "dataset", and "space". If this + function is called by a child class, it will default to the child class's `repo_type`. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit. + commit_description (`str`, *optional*) + The description of the generated commit. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the `"main"` branch. + create_pr (`bool`, *optional*): + Whether or not to create a Pull Request with this commit. Defaults to `False`. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + Returns: + `str`: URL of the commit which updated the card metadata. + """ + + # If repo type is provided, otherwise, use the repo type of the card. + repo_type = repo_type or self.repo_type + + # Validate card before pushing to hub + self.validate(repo_type=repo_type) + + with SoftTemporaryDirectory() as tmpdir: + tmp_path = Path(tmpdir) / constants.REPOCARD_NAME + tmp_path.write_text(str(self)) + url = upload_file( + path_or_fileobj=str(tmp_path), + path_in_repo=constants.REPOCARD_NAME, + repo_id=repo_id, + token=token, + repo_type=repo_type, + commit_message=commit_message, + commit_description=commit_description, + create_pr=create_pr, + revision=revision, + parent_commit=parent_commit, + ) + return url + + @classmethod + def from_template( + cls, + card_data: CardData, + template_path: Optional[str] = None, + template_str: Optional[str] = None, + **template_kwargs, + ): + """Initialize a RepoCard from a template. By default, it uses the default template. + + Templates are Jinja2 templates that can be customized by passing keyword arguments. + + Args: + card_data (`huggingface_hub.CardData`): + A huggingface_hub.CardData instance containing the metadata you want to include in the YAML + header of the repo card on the Hugging Face Hub. + template_path (`str`, *optional*): + A path to a markdown file with optional Jinja template variables that can be filled + in with `template_kwargs`. Defaults to the default template. + + Returns: + [`huggingface_hub.repocard.RepoCard`]: A RepoCard instance with the specified card data and content from the + template. + """ + if is_jinja_available(): + import jinja2 + else: + raise ImportError( + "Using RepoCard.from_template requires Jinja2 to be installed. Please" + " install it with `pip install Jinja2`." + ) + + kwargs = card_data.to_dict().copy() + kwargs.update(template_kwargs) # Template_kwargs have priority + + if template_path is not None: + template_str = Path(template_path).read_text() + if template_str is None: + template_str = Path(cls.default_template_path).read_text() + template = jinja2.Template(template_str) + content = template.render(card_data=card_data.to_yaml(), **kwargs) + return cls(content) + + +class ModelCard(RepoCard): + card_data_class = ModelCardData + default_template_path = TEMPLATE_MODELCARD_PATH + repo_type = "model" + + @classmethod + def from_template( # type: ignore # violates Liskov property but easier to use + cls, + card_data: ModelCardData, + template_path: Optional[str] = None, + template_str: Optional[str] = None, + **template_kwargs, + ): + """Initialize a ModelCard from a template. By default, it uses the default template, which can be found here: + https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/templates/modelcard_template.md + + Templates are Jinja2 templates that can be customized by passing keyword arguments. + + Args: + card_data (`huggingface_hub.ModelCardData`): + A huggingface_hub.ModelCardData instance containing the metadata you want to include in the YAML + header of the model card on the Hugging Face Hub. + template_path (`str`, *optional*): + A path to a markdown file with optional Jinja template variables that can be filled + in with `template_kwargs`. Defaults to the default template. + + Returns: + [`huggingface_hub.ModelCard`]: A ModelCard instance with the specified card data and content from the + template. + + Example: + ```python + >>> from huggingface_hub import ModelCard, ModelCardData, EvalResult + + >>> # Using the Default Template + >>> card_data = ModelCardData( + ... language='en', + ... license='mit', + ... library_name='timm', + ... tags=['image-classification', 'resnet'], + ... datasets=['beans'], + ... metrics=['accuracy'], + ... ) + >>> card = ModelCard.from_template( + ... card_data, + ... model_description='This model does x + y...' + ... ) + + >>> # Including Evaluation Results + >>> card_data = ModelCardData( + ... language='en', + ... tags=['image-classification', 'resnet'], + ... eval_results=[ + ... EvalResult( + ... task_type='image-classification', + ... dataset_type='beans', + ... dataset_name='Beans', + ... metric_type='accuracy', + ... metric_value=0.9, + ... ), + ... ], + ... model_name='my-cool-model', + ... ) + >>> card = ModelCard.from_template(card_data) + + >>> # Using a Custom Template + >>> card_data = ModelCardData( + ... language='en', + ... tags=['image-classification', 'resnet'] + ... ) + >>> card = ModelCard.from_template( + ... card_data=card_data, + ... template_path='./src/huggingface_hub/templates/modelcard_template.md', + ... custom_template_var='custom value', # will be replaced in template if it exists + ... ) + + ``` + """ + return super().from_template(card_data, template_path, template_str, **template_kwargs) + + +class DatasetCard(RepoCard): + card_data_class = DatasetCardData + default_template_path = TEMPLATE_DATASETCARD_PATH + repo_type = "dataset" + + @classmethod + def from_template( # type: ignore # violates Liskov property but easier to use + cls, + card_data: DatasetCardData, + template_path: Optional[str] = None, + template_str: Optional[str] = None, + **template_kwargs, + ): + """Initialize a DatasetCard from a template. By default, it uses the default template, which can be found here: + https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/templates/datasetcard_template.md + + Templates are Jinja2 templates that can be customized by passing keyword arguments. + + Args: + card_data (`huggingface_hub.DatasetCardData`): + A huggingface_hub.DatasetCardData instance containing the metadata you want to include in the YAML + header of the dataset card on the Hugging Face Hub. + template_path (`str`, *optional*): + A path to a markdown file with optional Jinja template variables that can be filled + in with `template_kwargs`. Defaults to the default template. + + Returns: + [`huggingface_hub.DatasetCard`]: A DatasetCard instance with the specified card data and content from the + template. + + Example: + ```python + >>> from huggingface_hub import DatasetCard, DatasetCardData + + >>> # Using the Default Template + >>> card_data = DatasetCardData( + ... language='en', + ... license='mit', + ... annotations_creators='crowdsourced', + ... task_categories=['text-classification'], + ... task_ids=['sentiment-classification', 'text-scoring'], + ... multilinguality='monolingual', + ... pretty_name='My Text Classification Dataset', + ... ) + >>> card = DatasetCard.from_template( + ... card_data, + ... pretty_name=card_data.pretty_name, + ... ) + + >>> # Using a Custom Template + >>> card_data = DatasetCardData( + ... language='en', + ... license='mit', + ... ) + >>> card = DatasetCard.from_template( + ... card_data=card_data, + ... template_path='./src/huggingface_hub/templates/datasetcard_template.md', + ... custom_template_var='custom value', # will be replaced in template if it exists + ... ) + + ``` + """ + return super().from_template(card_data, template_path, template_str, **template_kwargs) + + +class SpaceCard(RepoCard): + card_data_class = SpaceCardData + default_template_path = TEMPLATE_MODELCARD_PATH + repo_type = "space" + + +def _detect_line_ending(content: str) -> Literal["\r", "\n", "\r\n", None]: # noqa: F722 + """Detect the line ending of a string. Used by RepoCard to avoid making huge diff on newlines. + + Uses same implementation as in Hub server, keep it in sync. + + Returns: + str: The detected line ending of the string. + """ + cr = content.count("\r") + lf = content.count("\n") + crlf = content.count("\r\n") + if cr + lf == 0: + return None + if crlf == cr and crlf == lf: + return "\r\n" + if cr > lf: + return "\r" + else: + return "\n" + + +def metadata_load(local_path: Union[str, Path]) -> Optional[Dict]: + content = Path(local_path).read_text() + match = REGEX_YAML_BLOCK.search(content) + if match: + yaml_block = match.group(2) + data = yaml.safe_load(yaml_block) + if data is None or isinstance(data, dict): + return data + raise ValueError("repo card metadata block should be a dict") + else: + return None + + +def metadata_save(local_path: Union[str, Path], data: Dict) -> None: + """ + Save the metadata dict in the upper YAML part Trying to preserve newlines as + in the existing file. Docs about open() with newline="" parameter: + https://docs.python.org/3/library/functions.html?highlight=open#open Does + not work with "^M" linebreaks, which are replaced by \n + """ + line_break = "\n" + content = "" + # try to detect existing newline character + if os.path.exists(local_path): + with open(local_path, "r", newline="", encoding="utf8") as readme: + content = readme.read() + if isinstance(readme.newlines, tuple): + line_break = readme.newlines[0] + elif isinstance(readme.newlines, str): + line_break = readme.newlines + + # creates a new file if it not + with open(local_path, "w", newline="", encoding="utf8") as readme: + data_yaml = yaml_dump(data, sort_keys=False, line_break=line_break) + # sort_keys: keep dict order + match = REGEX_YAML_BLOCK.search(content) + if match: + output = content[: match.start()] + f"---{line_break}{data_yaml}---{line_break}" + content[match.end() :] + else: + output = f"---{line_break}{data_yaml}---{line_break}{content}" + + readme.write(output) + readme.close() + + +def metadata_eval_result( + *, + model_pretty_name: str, + task_pretty_name: str, + task_id: str, + metrics_pretty_name: str, + metrics_id: str, + metrics_value: Any, + dataset_pretty_name: str, + dataset_id: str, + metrics_config: Optional[str] = None, + metrics_verified: bool = False, + dataset_config: Optional[str] = None, + dataset_split: Optional[str] = None, + dataset_revision: Optional[str] = None, + metrics_verification_token: Optional[str] = None, +) -> Dict: + """ + Creates a metadata dict with the result from a model evaluated on a dataset. + + Args: + model_pretty_name (`str`): + The name of the model in natural language. + task_pretty_name (`str`): + The name of a task in natural language. + task_id (`str`): + Example: automatic-speech-recognition. A task id. + metrics_pretty_name (`str`): + A name for the metric in natural language. Example: Test WER. + metrics_id (`str`): + Example: wer. A metric id from https://hf.co/metrics. + metrics_value (`Any`): + The value from the metric. Example: 20.0 or "20.0 ± 1.2". + dataset_pretty_name (`str`): + The name of the dataset in natural language. + dataset_id (`str`): + Example: common_voice. A dataset id from https://hf.co/datasets. + metrics_config (`str`, *optional*): + The name of the metric configuration used in `load_metric()`. + Example: bleurt-large-512 in `load_metric("bleurt", "bleurt-large-512")`. + metrics_verified (`bool`, *optional*, defaults to `False`): + Indicates whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. Automatically computed by Hugging Face, do not set. + dataset_config (`str`, *optional*): + Example: fr. The name of the dataset configuration used in `load_dataset()`. + dataset_split (`str`, *optional*): + Example: test. The name of the dataset split used in `load_dataset()`. + dataset_revision (`str`, *optional*): + Example: 5503434ddd753f426f4b38109466949a1217c2bb. The name of the dataset dataset revision + used in `load_dataset()`. + metrics_verification_token (`bool`, *optional*): + A JSON Web Token that is used to verify whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. + + Returns: + `dict`: a metadata dict with the result from a model evaluated on a dataset. + + Example: + ```python + >>> from huggingface_hub import metadata_eval_result + >>> results = metadata_eval_result( + ... model_pretty_name="RoBERTa fine-tuned on ReactionGIF", + ... task_pretty_name="Text Classification", + ... task_id="text-classification", + ... metrics_pretty_name="Accuracy", + ... metrics_id="accuracy", + ... metrics_value=0.2662102282047272, + ... dataset_pretty_name="ReactionJPEG", + ... dataset_id="julien-c/reactionjpeg", + ... dataset_config="default", + ... dataset_split="test", + ... ) + >>> results == { + ... 'model-index': [ + ... { + ... 'name': 'RoBERTa fine-tuned on ReactionGIF', + ... 'results': [ + ... { + ... 'task': { + ... 'type': 'text-classification', + ... 'name': 'Text Classification' + ... }, + ... 'dataset': { + ... 'name': 'ReactionJPEG', + ... 'type': 'julien-c/reactionjpeg', + ... 'config': 'default', + ... 'split': 'test' + ... }, + ... 'metrics': [ + ... { + ... 'type': 'accuracy', + ... 'value': 0.2662102282047272, + ... 'name': 'Accuracy', + ... 'verified': False + ... } + ... ] + ... } + ... ] + ... } + ... ] + ... } + True + + ``` + """ + + return { + "model-index": eval_results_to_model_index( + model_name=model_pretty_name, + eval_results=[ + EvalResult( + task_name=task_pretty_name, + task_type=task_id, + metric_name=metrics_pretty_name, + metric_type=metrics_id, + metric_value=metrics_value, + dataset_name=dataset_pretty_name, + dataset_type=dataset_id, + metric_config=metrics_config, + verified=metrics_verified, + verify_token=metrics_verification_token, + dataset_config=dataset_config, + dataset_split=dataset_split, + dataset_revision=dataset_revision, + ) + ], + ) + } + + +@validate_hf_hub_args +def metadata_update( + repo_id: str, + metadata: Dict, + *, + repo_type: Optional[str] = None, + overwrite: bool = False, + token: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + revision: Optional[str] = None, + create_pr: bool = False, + parent_commit: Optional[str] = None, +) -> str: + """ + Updates the metadata in the README.md of a repository on the Hugging Face Hub. + If the README.md file doesn't exist yet, a new one is created with metadata and an + the default ModelCard or DatasetCard template. For `space` repo, an error is thrown + as a Space cannot exist without a `README.md` file. + + Args: + repo_id (`str`): + The name of the repository. + metadata (`dict`): + A dictionary containing the metadata to be updated. + repo_type (`str`, *optional*): + Set to `"dataset"` or `"space"` if updating to a dataset or space, + `None` or `"model"` if updating to a model. Default is `None`. + overwrite (`bool`, *optional*, defaults to `False`): + If set to `True` an existing field can be overwritten, otherwise + attempting to overwrite an existing field will cause an error. + token (`str`, *optional*): + The Hugging Face authentication token. + commit_message (`str`, *optional*): + The summary / title / first line of the generated commit. Defaults to + `f"Update metadata with huggingface_hub"` + commit_description (`str` *optional*) + The description of the generated commit + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the + `"main"` branch. + create_pr (`boolean`, *optional*): + Whether or not to create a Pull Request from `revision` with that commit. + Defaults to `False`. + parent_commit (`str`, *optional*): + The OID / SHA of the parent commit, as a hexadecimal string. Shorthands (7 first characters) are also supported. + If specified and `create_pr` is `False`, the commit will fail if `revision` does not point to `parent_commit`. + If specified and `create_pr` is `True`, the pull request will be created from `parent_commit`. + Specifying `parent_commit` ensures the repo has not changed before committing the changes, and can be + especially useful if the repo is updated / committed to concurrently. + Returns: + `str`: URL of the commit which updated the card metadata. + + Example: + ```python + >>> from huggingface_hub import metadata_update + >>> metadata = {'model-index': [{'name': 'RoBERTa fine-tuned on ReactionGIF', + ... 'results': [{'dataset': {'name': 'ReactionGIF', + ... 'type': 'julien-c/reactiongif'}, + ... 'metrics': [{'name': 'Recall', + ... 'type': 'recall', + ... 'value': 0.7762102282047272}], + ... 'task': {'name': 'Text Classification', + ... 'type': 'text-classification'}}]}]} + >>> url = metadata_update("hf-internal-testing/reactiongif-roberta-card", metadata) + + ``` + """ + commit_message = commit_message if commit_message is not None else "Update metadata with huggingface_hub" + + # Card class given repo_type + card_class: Type[RepoCard] + if repo_type is None or repo_type == "model": + card_class = ModelCard + elif repo_type == "dataset": + card_class = DatasetCard + elif repo_type == "space": + card_class = RepoCard + else: + raise ValueError(f"Unknown repo_type: {repo_type}") + + # Either load repo_card from the Hub or create an empty one. + # NOTE: Will not create the repo if it doesn't exist. + try: + card = card_class.load(repo_id, token=token, repo_type=repo_type) + except EntryNotFoundError: + if repo_type == "space": + raise ValueError("Cannot update metadata on a Space that doesn't contain a `README.md` file.") + + # Initialize a ModelCard or DatasetCard from default template and no data. + card = card_class.from_template(CardData()) + + for key, value in metadata.items(): + if key == "model-index": + # if the new metadata doesn't include a name, either use existing one or repo name + if "name" not in value[0]: + value[0]["name"] = getattr(card, "model_name", repo_id) + model_name, new_results = model_index_to_eval_results(value) + if card.data.eval_results is None: + card.data.eval_results = new_results + card.data.model_name = model_name + else: + existing_results = card.data.eval_results + + # Iterate over new results + # Iterate over existing results + # If both results describe the same metric but value is different: + # If overwrite=True: overwrite the metric value + # Else: raise ValueError + # Else: append new result to existing ones. + for new_result in new_results: + result_found = False + for existing_result in existing_results: + if new_result.is_equal_except_value(existing_result): + if new_result != existing_result and not overwrite: + raise ValueError( + "You passed a new value for the existing metric" + f" 'name: {new_result.metric_name}, type: " + f"{new_result.metric_type}'. Set `overwrite=True`" + " to overwrite existing metrics." + ) + result_found = True + existing_result.metric_value = new_result.metric_value + if existing_result.verified is True: + existing_result.verify_token = new_result.verify_token + if not result_found: + card.data.eval_results.append(new_result) + else: + # Any metadata that is not a result metric + if card.data.get(key) is not None and not overwrite and card.data.get(key) != value: + raise ValueError( + f"You passed a new value for the existing meta data field '{key}'." + " Set `overwrite=True` to overwrite existing metadata." + ) + else: + card.data[key] = value + + return card.push_to_hub( + repo_id, + token=token, + repo_type=repo_type, + commit_message=commit_message, + commit_description=commit_description, + create_pr=create_pr, + revision=revision, + parent_commit=parent_commit, + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/repocard_data.py b/lib/python3.12/site-packages/huggingface_hub/repocard_data.py new file mode 100644 index 0000000000000000000000000000000000000000..62215f2274e482d4ed69a1d6deeafdf34fc5a6a4 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/repocard_data.py @@ -0,0 +1,770 @@ +import copy +from collections import defaultdict +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +from huggingface_hub.utils import logging, yaml_dump + + +logger = logging.get_logger(__name__) + + +@dataclass +class EvalResult: + """ + Flattened representation of individual evaluation results found in model-index of Model Cards. + + For more information on the model-index spec, see https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1. + + Args: + task_type (`str`): + The task identifier. Example: "image-classification". + dataset_type (`str`): + The dataset identifier. Example: "common_voice". Use dataset id from https://hf.co/datasets. + dataset_name (`str`): + A pretty name for the dataset. Example: "Common Voice (French)". + metric_type (`str`): + The metric identifier. Example: "wer". Use metric id from https://hf.co/metrics. + metric_value (`Any`): + The metric value. Example: 0.9 or "20.0 ± 1.2". + task_name (`str`, *optional*): + A pretty name for the task. Example: "Speech Recognition". + dataset_config (`str`, *optional*): + The name of the dataset configuration used in `load_dataset()`. + Example: fr in `load_dataset("common_voice", "fr")`. See the `datasets` docs for more info: + https://hf.co/docs/datasets/package_reference/loading_methods#datasets.load_dataset.name + dataset_split (`str`, *optional*): + The split used in `load_dataset()`. Example: "test". + dataset_revision (`str`, *optional*): + The revision (AKA Git Sha) of the dataset used in `load_dataset()`. + Example: 5503434ddd753f426f4b38109466949a1217c2bb + dataset_args (`Dict[str, Any]`, *optional*): + The arguments passed during `Metric.compute()`. Example for `bleu`: `{"max_order": 4}` + metric_name (`str`, *optional*): + A pretty name for the metric. Example: "Test WER". + metric_config (`str`, *optional*): + The name of the metric configuration used in `load_metric()`. + Example: bleurt-large-512 in `load_metric("bleurt", "bleurt-large-512")`. + See the `datasets` docs for more info: https://huggingface.co/docs/datasets/v2.1.0/en/loading#load-configurations + metric_args (`Dict[str, Any]`, *optional*): + The arguments passed during `Metric.compute()`. Example for `bleu`: max_order: 4 + verified (`bool`, *optional*): + Indicates whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. Automatically computed by Hugging Face, do not set. + verify_token (`str`, *optional*): + A JSON Web Token that is used to verify whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. + source_name (`str`, *optional*): + The name of the source of the evaluation result. Example: "Open LLM Leaderboard". + source_url (`str`, *optional*): + The URL of the source of the evaluation result. Example: "https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard". + """ + + # Required + + # The task identifier + # Example: automatic-speech-recognition + task_type: str + + # The dataset identifier + # Example: common_voice. Use dataset id from https://hf.co/datasets + dataset_type: str + + # A pretty name for the dataset. + # Example: Common Voice (French) + dataset_name: str + + # The metric identifier + # Example: wer. Use metric id from https://hf.co/metrics + metric_type: str + + # Value of the metric. + # Example: 20.0 or "20.0 ± 1.2" + metric_value: Any + + # Optional + + # A pretty name for the task. + # Example: Speech Recognition + task_name: Optional[str] = None + + # The name of the dataset configuration used in `load_dataset()`. + # Example: fr in `load_dataset("common_voice", "fr")`. + # See the `datasets` docs for more info: + # https://huggingface.co/docs/datasets/package_reference/loading_methods#datasets.load_dataset.name + dataset_config: Optional[str] = None + + # The split used in `load_dataset()`. + # Example: test + dataset_split: Optional[str] = None + + # The revision (AKA Git Sha) of the dataset used in `load_dataset()`. + # Example: 5503434ddd753f426f4b38109466949a1217c2bb + dataset_revision: Optional[str] = None + + # The arguments passed during `Metric.compute()`. + # Example for `bleu`: max_order: 4 + dataset_args: Optional[Dict[str, Any]] = None + + # A pretty name for the metric. + # Example: Test WER + metric_name: Optional[str] = None + + # The name of the metric configuration used in `load_metric()`. + # Example: bleurt-large-512 in `load_metric("bleurt", "bleurt-large-512")`. + # See the `datasets` docs for more info: https://huggingface.co/docs/datasets/v2.1.0/en/loading#load-configurations + metric_config: Optional[str] = None + + # The arguments passed during `Metric.compute()`. + # Example for `bleu`: max_order: 4 + metric_args: Optional[Dict[str, Any]] = None + + # Indicates whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. Automatically computed by Hugging Face, do not set. + verified: Optional[bool] = None + + # A JSON Web Token that is used to verify whether the metrics originate from Hugging Face's [evaluation service](https://huggingface.co/spaces/autoevaluate/model-evaluator) or not. + verify_token: Optional[str] = None + + # The name of the source of the evaluation result. + # Example: Open LLM Leaderboard + source_name: Optional[str] = None + + # The URL of the source of the evaluation result. + # Example: https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard + source_url: Optional[str] = None + + @property + def unique_identifier(self) -> tuple: + """Returns a tuple that uniquely identifies this evaluation.""" + return ( + self.task_type, + self.dataset_type, + self.dataset_config, + self.dataset_split, + self.dataset_revision, + ) + + def is_equal_except_value(self, other: "EvalResult") -> bool: + """ + Return True if `self` and `other` describe exactly the same metric but with a + different value. + """ + for key, _ in self.__dict__.items(): + if key == "metric_value": + continue + # For metrics computed by Hugging Face's evaluation service, `verify_token` is derived from `metric_value`, + # so we exclude it here in the comparison. + if key != "verify_token" and getattr(self, key) != getattr(other, key): + return False + return True + + def __post_init__(self) -> None: + if self.source_name is not None and self.source_url is None: + raise ValueError("If `source_name` is provided, `source_url` must also be provided.") + + +@dataclass +class CardData: + """Structure containing metadata from a RepoCard. + + [`CardData`] is the parent class of [`ModelCardData`] and [`DatasetCardData`]. + + Metadata can be exported as a dictionary or YAML. Export can be customized to alter the representation of the data + (example: flatten evaluation results). `CardData` behaves as a dictionary (can get, pop, set values) but do not + inherit from `dict` to allow this export step. + """ + + def __init__(self, ignore_metadata_errors: bool = False, **kwargs): + self.__dict__.update(kwargs) + + def to_dict(self): + """Converts CardData to a dict. + + Returns: + `dict`: CardData represented as a dictionary ready to be dumped to a YAML + block for inclusion in a README.md file. + """ + + data_dict = copy.deepcopy(self.__dict__) + self._to_dict(data_dict) + return {key: value for key, value in data_dict.items() if value is not None} + + def _to_dict(self, data_dict): + """Use this method in child classes to alter the dict representation of the data. Alter the dict in-place. + + Args: + data_dict (`dict`): The raw dict representation of the card data. + """ + pass + + def to_yaml(self, line_break=None, original_order: Optional[List[str]] = None) -> str: + """Dumps CardData to a YAML block for inclusion in a README.md file. + + Args: + line_break (str, *optional*): + The line break to use when dumping to yaml. + + Returns: + `str`: CardData represented as a YAML block. + """ + if original_order: + self.__dict__ = { + k: self.__dict__[k] + for k in original_order + list(set(self.__dict__.keys()) - set(original_order)) + if k in self.__dict__ + } + return yaml_dump(self.to_dict(), sort_keys=False, line_break=line_break).strip() + + def __repr__(self): + return repr(self.__dict__) + + def __str__(self): + return self.to_yaml() + + def get(self, key: str, default: Any = None) -> Any: + """Get value for a given metadata key.""" + value = self.__dict__.get(key) + return default if value is None else value + + def pop(self, key: str, default: Any = None) -> Any: + """Pop value for a given metadata key.""" + return self.__dict__.pop(key, default) + + def __getitem__(self, key: str) -> Any: + """Get value for a given metadata key.""" + return self.__dict__[key] + + def __setitem__(self, key: str, value: Any) -> None: + """Set value for a given metadata key.""" + self.__dict__[key] = value + + def __contains__(self, key: str) -> bool: + """Check if a given metadata key is set.""" + return key in self.__dict__ + + def __len__(self) -> int: + """Return the number of metadata keys set.""" + return len(self.__dict__) + + +def _validate_eval_results( + eval_results: Optional[Union[EvalResult, List[EvalResult]]], + model_name: Optional[str], +) -> List[EvalResult]: + if eval_results is None: + return [] + if isinstance(eval_results, EvalResult): + eval_results = [eval_results] + if not isinstance(eval_results, list) or not all(isinstance(r, EvalResult) for r in eval_results): + raise ValueError( + f"`eval_results` should be of type `EvalResult` or a list of `EvalResult`, got {type(eval_results)}." + ) + if model_name is None: + raise ValueError("Passing `eval_results` requires `model_name` to be set.") + return eval_results + + +class ModelCardData(CardData): + """Model Card Metadata that is used by Hugging Face Hub when included at the top of your README.md + + Args: + base_model (`str` or `List[str]`, *optional*): + The identifier of the base model from which the model derives. This is applicable for example if your model is a + fine-tune or adapter of an existing model. The value must be the ID of a model on the Hub (or a list of IDs + if your model derives from multiple models). Defaults to None. + datasets (`Union[str, List[str]]`, *optional*): + Dataset or list of datasets that were used to train this model. Should be a dataset ID + found on https://hf.co/datasets. Defaults to None. + eval_results (`Union[List[EvalResult], EvalResult]`, *optional*): + List of `huggingface_hub.EvalResult` that define evaluation results of the model. If provided, + `model_name` is used to as a name on PapersWithCode's leaderboards. Defaults to `None`. + language (`Union[str, List[str]]`, *optional*): + Language of model's training data or metadata. It must be an ISO 639-1, 639-2 or + 639-3 code (two/three letters), or a special value like "code", "multilingual". Defaults to `None`. + library_name (`str`, *optional*): + Name of library used by this model. Example: keras or any library from + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/model-libraries.ts. + Defaults to None. + license (`str`, *optional*): + License of this model. Example: apache-2.0 or any license from + https://huggingface.co/docs/hub/repositories-licenses. Defaults to None. + license_name (`str`, *optional*): + Name of the license of this model. Defaults to None. To be used in conjunction with `license_link`. + Common licenses (Apache-2.0, MIT, CC-BY-SA-4.0) do not need a name. In that case, use `license` instead. + license_link (`str`, *optional*): + Link to the license of this model. Defaults to None. To be used in conjunction with `license_name`. + Common licenses (Apache-2.0, MIT, CC-BY-SA-4.0) do not need a link. In that case, use `license` instead. + metrics (`List[str]`, *optional*): + List of metrics used to evaluate this model. Should be a metric name that can be found + at https://hf.co/metrics. Example: 'accuracy'. Defaults to None. + model_name (`str`, *optional*): + A name for this model. It is used along with + `eval_results` to construct the `model-index` within the card's metadata. The name + you supply here is what will be used on PapersWithCode's leaderboards. If None is provided + then the repo name is used as a default. Defaults to None. + pipeline_tag (`str`, *optional*): + The pipeline tag associated with the model. Example: "text-classification". + tags (`List[str]`, *optional*): + List of tags to add to your model that can be used when filtering on the Hugging + Face Hub. Defaults to None. + ignore_metadata_errors (`str`): + If True, errors while parsing the metadata section will be ignored. Some information might be lost during + the process. Use it at your own risk. + kwargs (`dict`, *optional*): + Additional metadata that will be added to the model card. Defaults to None. + + Example: + ```python + >>> from huggingface_hub import ModelCardData + >>> card_data = ModelCardData( + ... language="en", + ... license="mit", + ... library_name="timm", + ... tags=['image-classification', 'resnet'], + ... ) + >>> card_data.to_dict() + {'language': 'en', 'license': 'mit', 'library_name': 'timm', 'tags': ['image-classification', 'resnet']} + + ``` + """ + + def __init__( + self, + *, + base_model: Optional[Union[str, List[str]]] = None, + datasets: Optional[Union[str, List[str]]] = None, + eval_results: Optional[List[EvalResult]] = None, + language: Optional[Union[str, List[str]]] = None, + library_name: Optional[str] = None, + license: Optional[str] = None, + license_name: Optional[str] = None, + license_link: Optional[str] = None, + metrics: Optional[List[str]] = None, + model_name: Optional[str] = None, + pipeline_tag: Optional[str] = None, + tags: Optional[List[str]] = None, + ignore_metadata_errors: bool = False, + **kwargs, + ): + self.base_model = base_model + self.datasets = datasets + self.eval_results = eval_results + self.language = language + self.library_name = library_name + self.license = license + self.license_name = license_name + self.license_link = license_link + self.metrics = metrics + self.model_name = model_name + self.pipeline_tag = pipeline_tag + self.tags = _to_unique_list(tags) + + model_index = kwargs.pop("model-index", None) + if model_index: + try: + model_name, eval_results = model_index_to_eval_results(model_index) + self.model_name = model_name + self.eval_results = eval_results + except (KeyError, TypeError) as error: + if ignore_metadata_errors: + logger.warning("Invalid model-index. Not loading eval results into CardData.") + else: + raise ValueError( + f"Invalid `model_index` in metadata cannot be parsed: {error.__class__} {error}. Pass" + " `ignore_metadata_errors=True` to ignore this error while loading a Model Card. Warning:" + " some information will be lost. Use it at your own risk." + ) + + super().__init__(**kwargs) + + if self.eval_results: + try: + self.eval_results = _validate_eval_results(self.eval_results, self.model_name) + except Exception as e: + if ignore_metadata_errors: + logger.warning(f"Failed to validate eval_results: {e}. Not loading eval results into CardData.") + else: + raise ValueError(f"Failed to validate eval_results: {e}") from e + + def _to_dict(self, data_dict): + """Format the internal data dict. In this case, we convert eval results to a valid model index""" + if self.eval_results is not None: + data_dict["model-index"] = eval_results_to_model_index(self.model_name, self.eval_results) + del data_dict["eval_results"], data_dict["model_name"] + + +class DatasetCardData(CardData): + """Dataset Card Metadata that is used by Hugging Face Hub when included at the top of your README.md + + Args: + language (`List[str]`, *optional*): + Language of dataset's data or metadata. It must be an ISO 639-1, 639-2 or + 639-3 code (two/three letters), or a special value like "code", "multilingual". + license (`Union[str, List[str]]`, *optional*): + License(s) of this dataset. Example: apache-2.0 or any license from + https://huggingface.co/docs/hub/repositories-licenses. + annotations_creators (`Union[str, List[str]]`, *optional*): + How the annotations for the dataset were created. + Options are: 'found', 'crowdsourced', 'expert-generated', 'machine-generated', 'no-annotation', 'other'. + language_creators (`Union[str, List[str]]`, *optional*): + How the text-based data in the dataset was created. + Options are: 'found', 'crowdsourced', 'expert-generated', 'machine-generated', 'other' + multilinguality (`Union[str, List[str]]`, *optional*): + Whether the dataset is multilingual. + Options are: 'monolingual', 'multilingual', 'translation', 'other'. + size_categories (`Union[str, List[str]]`, *optional*): + The number of examples in the dataset. Options are: 'n<1K', '1K1T', and 'other'. + source_datasets (`List[str]]`, *optional*): + Indicates whether the dataset is an original dataset or extended from another existing dataset. + Options are: 'original' and 'extended'. + task_categories (`Union[str, List[str]]`, *optional*): + What categories of task does the dataset support? + task_ids (`Union[str, List[str]]`, *optional*): + What specific tasks does the dataset support? + paperswithcode_id (`str`, *optional*): + ID of the dataset on PapersWithCode. + pretty_name (`str`, *optional*): + A more human-readable name for the dataset. (ex. "Cats vs. Dogs") + train_eval_index (`Dict`, *optional*): + A dictionary that describes the necessary spec for doing evaluation on the Hub. + If not provided, it will be gathered from the 'train-eval-index' key of the kwargs. + config_names (`Union[str, List[str]]`, *optional*): + A list of the available dataset configs for the dataset. + """ + + def __init__( + self, + *, + language: Optional[Union[str, List[str]]] = None, + license: Optional[Union[str, List[str]]] = None, + annotations_creators: Optional[Union[str, List[str]]] = None, + language_creators: Optional[Union[str, List[str]]] = None, + multilinguality: Optional[Union[str, List[str]]] = None, + size_categories: Optional[Union[str, List[str]]] = None, + source_datasets: Optional[List[str]] = None, + task_categories: Optional[Union[str, List[str]]] = None, + task_ids: Optional[Union[str, List[str]]] = None, + paperswithcode_id: Optional[str] = None, + pretty_name: Optional[str] = None, + train_eval_index: Optional[Dict] = None, + config_names: Optional[Union[str, List[str]]] = None, + ignore_metadata_errors: bool = False, + **kwargs, + ): + self.annotations_creators = annotations_creators + self.language_creators = language_creators + self.language = language + self.license = license + self.multilinguality = multilinguality + self.size_categories = size_categories + self.source_datasets = source_datasets + self.task_categories = task_categories + self.task_ids = task_ids + self.paperswithcode_id = paperswithcode_id + self.pretty_name = pretty_name + self.config_names = config_names + + # TODO - maybe handle this similarly to EvalResult? + self.train_eval_index = train_eval_index or kwargs.pop("train-eval-index", None) + super().__init__(**kwargs) + + def _to_dict(self, data_dict): + data_dict["train-eval-index"] = data_dict.pop("train_eval_index") + + +class SpaceCardData(CardData): + """Space Card Metadata that is used by Hugging Face Hub when included at the top of your README.md + + To get an exhaustive reference of Spaces configuration, please visit https://huggingface.co/docs/hub/spaces-config-reference#spaces-configuration-reference. + + Args: + title (`str`, *optional*) + Title of the Space. + sdk (`str`, *optional*) + SDK of the Space (one of `gradio`, `streamlit`, `docker`, or `static`). + sdk_version (`str`, *optional*) + Version of the used SDK (if Gradio/Streamlit sdk). + python_version (`str`, *optional*) + Python version used in the Space (if Gradio/Streamlit sdk). + app_file (`str`, *optional*) + Path to your main application file (which contains either gradio or streamlit Python code, or static html code). + Path is relative to the root of the repository. + app_port (`str`, *optional*) + Port on which your application is running. Used only if sdk is `docker`. + license (`str`, *optional*) + License of this model. Example: apache-2.0 or any license from + https://huggingface.co/docs/hub/repositories-licenses. + duplicated_from (`str`, *optional*) + ID of the original Space if this is a duplicated Space. + models (List[`str`], *optional*) + List of models related to this Space. Should be a dataset ID found on https://hf.co/models. + datasets (`List[str]`, *optional*) + List of datasets related to this Space. Should be a dataset ID found on https://hf.co/datasets. + tags (`List[str]`, *optional*) + List of tags to add to your Space that can be used when filtering on the Hub. + ignore_metadata_errors (`str`): + If True, errors while parsing the metadata section will be ignored. Some information might be lost during + the process. Use it at your own risk. + kwargs (`dict`, *optional*): + Additional metadata that will be added to the space card. + + Example: + ```python + >>> from huggingface_hub import SpaceCardData + >>> card_data = SpaceCardData( + ... title="Dreambooth Training", + ... license="mit", + ... sdk="gradio", + ... duplicated_from="multimodalart/dreambooth-training" + ... ) + >>> card_data.to_dict() + {'title': 'Dreambooth Training', 'sdk': 'gradio', 'license': 'mit', 'duplicated_from': 'multimodalart/dreambooth-training'} + ``` + """ + + def __init__( + self, + *, + title: Optional[str] = None, + sdk: Optional[str] = None, + sdk_version: Optional[str] = None, + python_version: Optional[str] = None, + app_file: Optional[str] = None, + app_port: Optional[int] = None, + license: Optional[str] = None, + duplicated_from: Optional[str] = None, + models: Optional[List[str]] = None, + datasets: Optional[List[str]] = None, + tags: Optional[List[str]] = None, + ignore_metadata_errors: bool = False, + **kwargs, + ): + self.title = title + self.sdk = sdk + self.sdk_version = sdk_version + self.python_version = python_version + self.app_file = app_file + self.app_port = app_port + self.license = license + self.duplicated_from = duplicated_from + self.models = models + self.datasets = datasets + self.tags = _to_unique_list(tags) + super().__init__(**kwargs) + + +def model_index_to_eval_results(model_index: List[Dict[str, Any]]) -> Tuple[str, List[EvalResult]]: + """Takes in a model index and returns the model name and a list of `huggingface_hub.EvalResult` objects. + + A detailed spec of the model index can be found here: + https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1 + + Args: + model_index (`List[Dict[str, Any]]`): + A model index data structure, likely coming from a README.md file on the + Hugging Face Hub. + + Returns: + model_name (`str`): + The name of the model as found in the model index. This is used as the + identifier for the model on leaderboards like PapersWithCode. + eval_results (`List[EvalResult]`): + A list of `huggingface_hub.EvalResult` objects containing the metrics + reported in the provided model_index. + + Example: + ```python + >>> from huggingface_hub.repocard_data import model_index_to_eval_results + >>> # Define a minimal model index + >>> model_index = [ + ... { + ... "name": "my-cool-model", + ... "results": [ + ... { + ... "task": { + ... "type": "image-classification" + ... }, + ... "dataset": { + ... "type": "beans", + ... "name": "Beans" + ... }, + ... "metrics": [ + ... { + ... "type": "accuracy", + ... "value": 0.9 + ... } + ... ] + ... } + ... ] + ... } + ... ] + >>> model_name, eval_results = model_index_to_eval_results(model_index) + >>> model_name + 'my-cool-model' + >>> eval_results[0].task_type + 'image-classification' + >>> eval_results[0].metric_type + 'accuracy' + + ``` + """ + + eval_results = [] + for elem in model_index: + name = elem["name"] + results = elem["results"] + for result in results: + task_type = result["task"]["type"] + task_name = result["task"].get("name") + dataset_type = result["dataset"]["type"] + dataset_name = result["dataset"]["name"] + dataset_config = result["dataset"].get("config") + dataset_split = result["dataset"].get("split") + dataset_revision = result["dataset"].get("revision") + dataset_args = result["dataset"].get("args") + source_name = result.get("source", {}).get("name") + source_url = result.get("source", {}).get("url") + + for metric in result["metrics"]: + metric_type = metric["type"] + metric_value = metric["value"] + metric_name = metric.get("name") + metric_args = metric.get("args") + metric_config = metric.get("config") + verified = metric.get("verified") + verify_token = metric.get("verifyToken") + + eval_result = EvalResult( + task_type=task_type, # Required + dataset_type=dataset_type, # Required + dataset_name=dataset_name, # Required + metric_type=metric_type, # Required + metric_value=metric_value, # Required + task_name=task_name, + dataset_config=dataset_config, + dataset_split=dataset_split, + dataset_revision=dataset_revision, + dataset_args=dataset_args, + metric_name=metric_name, + metric_args=metric_args, + metric_config=metric_config, + verified=verified, + verify_token=verify_token, + source_name=source_name, + source_url=source_url, + ) + eval_results.append(eval_result) + return name, eval_results + + +def _remove_none(obj): + """ + Recursively remove `None` values from a dict. Borrowed from: https://stackoverflow.com/a/20558778 + """ + if isinstance(obj, (list, tuple, set)): + return type(obj)(_remove_none(x) for x in obj if x is not None) + elif isinstance(obj, dict): + return type(obj)((_remove_none(k), _remove_none(v)) for k, v in obj.items() if k is not None and v is not None) + else: + return obj + + +def eval_results_to_model_index(model_name: str, eval_results: List[EvalResult]) -> List[Dict[str, Any]]: + """Takes in given model name and list of `huggingface_hub.EvalResult` and returns a + valid model-index that will be compatible with the format expected by the + Hugging Face Hub. + + Args: + model_name (`str`): + Name of the model (ex. "my-cool-model"). This is used as the identifier + for the model on leaderboards like PapersWithCode. + eval_results (`List[EvalResult]`): + List of `huggingface_hub.EvalResult` objects containing the metrics to be + reported in the model-index. + + Returns: + model_index (`List[Dict[str, Any]]`): The eval_results converted to a model-index. + + Example: + ```python + >>> from huggingface_hub.repocard_data import eval_results_to_model_index, EvalResult + >>> # Define minimal eval_results + >>> eval_results = [ + ... EvalResult( + ... task_type="image-classification", # Required + ... dataset_type="beans", # Required + ... dataset_name="Beans", # Required + ... metric_type="accuracy", # Required + ... metric_value=0.9, # Required + ... ) + ... ] + >>> eval_results_to_model_index("my-cool-model", eval_results) + [{'name': 'my-cool-model', 'results': [{'task': {'type': 'image-classification'}, 'dataset': {'name': 'Beans', 'type': 'beans'}, 'metrics': [{'type': 'accuracy', 'value': 0.9}]}]}] + + ``` + """ + + # Metrics are reported on a unique task-and-dataset basis. + # Here, we make a map of those pairs and the associated EvalResults. + task_and_ds_types_map: Dict[Any, List[EvalResult]] = defaultdict(list) + for eval_result in eval_results: + task_and_ds_types_map[eval_result.unique_identifier].append(eval_result) + + # Use the map from above to generate the model index data. + model_index_data = [] + for results in task_and_ds_types_map.values(): + # All items from `results` share same metadata + sample_result = results[0] + data = { + "task": { + "type": sample_result.task_type, + "name": sample_result.task_name, + }, + "dataset": { + "name": sample_result.dataset_name, + "type": sample_result.dataset_type, + "config": sample_result.dataset_config, + "split": sample_result.dataset_split, + "revision": sample_result.dataset_revision, + "args": sample_result.dataset_args, + }, + "metrics": [ + { + "type": result.metric_type, + "value": result.metric_value, + "name": result.metric_name, + "config": result.metric_config, + "args": result.metric_args, + "verified": result.verified, + "verifyToken": result.verify_token, + } + for result in results + ], + } + if sample_result.source_url is not None: + source = { + "url": sample_result.source_url, + } + if sample_result.source_name is not None: + source["name"] = sample_result.source_name + data["source"] = source + model_index_data.append(data) + + # TODO - Check if there cases where this list is longer than one? + # Finally, the model index itself is list of dicts. + model_index = [ + { + "name": model_name, + "results": model_index_data, + } + ] + return _remove_none(model_index) + + +def _to_unique_list(tags: Optional[List[str]]) -> Optional[List[str]]: + if tags is None: + return tags + unique_tags = [] # make tags unique + keep order explicitly + for tag in tags: + if tag not in unique_tags: + unique_tags.append(tag) + return unique_tags diff --git a/lib/python3.12/site-packages/huggingface_hub/repository.py b/lib/python3.12/site-packages/huggingface_hub/repository.py new file mode 100644 index 0000000000000000000000000000000000000000..af1ab72fb458340f3fc211f0c5ef577b6471fda1 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/repository.py @@ -0,0 +1,1477 @@ +import atexit +import os +import re +import subprocess +import threading +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Dict, Iterator, List, Optional, Tuple, TypedDict, Union +from urllib.parse import urlparse + +from huggingface_hub import constants +from huggingface_hub.repocard import metadata_load, metadata_save + +from .hf_api import HfApi, repo_type_and_id_from_hf_id +from .lfs import LFS_MULTIPART_UPLOAD_COMMAND +from .utils import ( + SoftTemporaryDirectory, + get_token, + logging, + run_subprocess, + tqdm, + validate_hf_hub_args, +) +from .utils._deprecation import _deprecate_method + + +logger = logging.get_logger(__name__) + + +class CommandInProgress: + """ + Utility to follow commands launched asynchronously. + """ + + def __init__( + self, + title: str, + is_done_method: Callable, + status_method: Callable, + process: subprocess.Popen, + post_method: Optional[Callable] = None, + ): + self.title = title + self._is_done = is_done_method + self._status = status_method + self._process = process + self._stderr = "" + self._stdout = "" + self._post_method = post_method + + @property + def is_done(self) -> bool: + """ + Whether the process is done. + """ + result = self._is_done() + + if result and self._post_method is not None: + self._post_method() + self._post_method = None + + return result + + @property + def status(self) -> int: + """ + The exit code/status of the current action. Will return `0` if the + command has completed successfully, and a number between 1 and 255 if + the process errored-out. + + Will return -1 if the command is still ongoing. + """ + return self._status() + + @property + def failed(self) -> bool: + """ + Whether the process errored-out. + """ + return self.status > 0 + + @property + def stderr(self) -> str: + """ + The current output message on the standard error. + """ + if self._process.stderr is not None: + self._stderr += self._process.stderr.read() + return self._stderr + + @property + def stdout(self) -> str: + """ + The current output message on the standard output. + """ + if self._process.stdout is not None: + self._stdout += self._process.stdout.read() + return self._stdout + + def __repr__(self): + status = self.status + + if status == -1: + status = "running" + + return ( + f"[{self.title} command, status code: {status}," + f" {'in progress.' if not self.is_done else 'finished.'} PID:" + f" {self._process.pid}]" + ) + + +def is_git_repo(folder: Union[str, Path]) -> bool: + """ + Check if the folder is the root or part of a git repository + + Args: + folder (`str`): + The folder in which to run the command. + + Returns: + `bool`: `True` if the repository is part of a repository, `False` + otherwise. + """ + folder_exists = os.path.exists(os.path.join(folder, ".git")) + git_branch = subprocess.run("git branch".split(), cwd=folder, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + return folder_exists and git_branch.returncode == 0 + + +def is_local_clone(folder: Union[str, Path], remote_url: str) -> bool: + """ + Check if the folder is a local clone of the remote_url + + Args: + folder (`str` or `Path`): + The folder in which to run the command. + remote_url (`str`): + The url of a git repository. + + Returns: + `bool`: `True` if the repository is a local clone of the remote + repository specified, `False` otherwise. + """ + if not is_git_repo(folder): + return False + + remotes = run_subprocess("git remote -v", folder).stdout + + # Remove token for the test with remotes. + remote_url = re.sub(r"https://.*@", "https://", remote_url) + remotes = [re.sub(r"https://.*@", "https://", remote) for remote in remotes.split()] + return remote_url in remotes + + +def is_tracked_with_lfs(filename: Union[str, Path]) -> bool: + """ + Check if the file passed is tracked with git-lfs. + + Args: + filename (`str` or `Path`): + The filename to check. + + Returns: + `bool`: `True` if the file passed is tracked with git-lfs, `False` + otherwise. + """ + folder = Path(filename).parent + filename = Path(filename).name + + try: + p = run_subprocess("git check-attr -a".split() + [filename], folder) + attributes = p.stdout.strip() + except subprocess.CalledProcessError as exc: + if not is_git_repo(folder): + return False + else: + raise OSError(exc.stderr) + + if len(attributes) == 0: + return False + + found_lfs_tag = {"diff": False, "merge": False, "filter": False} + + for attribute in attributes.split("\n"): + for tag in found_lfs_tag.keys(): + if tag in attribute and "lfs" in attribute: + found_lfs_tag[tag] = True + + return all(found_lfs_tag.values()) + + +def is_git_ignored(filename: Union[str, Path]) -> bool: + """ + Check if file is git-ignored. Supports nested .gitignore files. + + Args: + filename (`str` or `Path`): + The filename to check. + + Returns: + `bool`: `True` if the file passed is ignored by `git`, `False` + otherwise. + """ + folder = Path(filename).parent + filename = Path(filename).name + + try: + p = run_subprocess("git check-ignore".split() + [filename], folder, check=False) + # Will return exit code 1 if not gitignored + is_ignored = not bool(p.returncode) + except subprocess.CalledProcessError as exc: + raise OSError(exc.stderr) + + return is_ignored + + +def is_binary_file(filename: Union[str, Path]) -> bool: + """ + Check if file is a binary file. + + Args: + filename (`str` or `Path`): + The filename to check. + + Returns: + `bool`: `True` if the file passed is a binary file, `False` otherwise. + """ + try: + with open(filename, "rb") as f: + content = f.read(10 * (1024**2)) # Read a maximum of 10MB + + # Code sample taken from the following stack overflow thread + # https://stackoverflow.com/questions/898669/how-can-i-detect-if-a-file-is-binary-non-text-in-python/7392391#7392391 + text_chars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7F}) + return bool(content.translate(None, text_chars)) + except UnicodeDecodeError: + return True + + +def files_to_be_staged(pattern: str = ".", folder: Union[str, Path, None] = None) -> List[str]: + """ + Returns a list of filenames that are to be staged. + + Args: + pattern (`str` or `Path`): + The pattern of filenames to check. Put `.` to get all files. + folder (`str` or `Path`): + The folder in which to run the command. + + Returns: + `List[str]`: List of files that are to be staged. + """ + try: + p = run_subprocess("git ls-files --exclude-standard -mo".split() + [pattern], folder) + if len(p.stdout.strip()): + files = p.stdout.strip().split("\n") + else: + files = [] + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + return files + + +def is_tracked_upstream(folder: Union[str, Path]) -> bool: + """ + Check if the current checked-out branch is tracked upstream. + + Args: + folder (`str` or `Path`): + The folder in which to run the command. + + Returns: + `bool`: `True` if the current checked-out branch is tracked upstream, + `False` otherwise. + """ + try: + run_subprocess("git rev-parse --symbolic-full-name --abbrev-ref @{u}", folder) + return True + except subprocess.CalledProcessError as exc: + if "HEAD" in exc.stderr: + raise OSError("No branch checked out") + + return False + + +def commits_to_push(folder: Union[str, Path], upstream: Optional[str] = None) -> int: + """ + Check the number of commits that would be pushed upstream + + Args: + folder (`str` or `Path`): + The folder in which to run the command. + upstream (`str`, *optional*): + The name of the upstream repository with which the comparison should be + made. + + Returns: + `int`: Number of commits that would be pushed upstream were a `git + push` to proceed. + """ + try: + result = run_subprocess(f"git cherry -v {upstream or ''}", folder) + return len(result.stdout.split("\n")) - 1 + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + +class PbarT(TypedDict): + # Used to store an opened progress bar in `_lfs_log_progress` + bar: tqdm + past_bytes: int + + +@contextmanager +def _lfs_log_progress(): + """ + This is a context manager that will log the Git LFS progress of cleaning, + smudging, pulling and pushing. + """ + + if logger.getEffectiveLevel() >= logging.ERROR: + try: + yield + except Exception: + pass + return + + def output_progress(stopping_event: threading.Event): + """ + To be launched as a separate thread with an event meaning it should stop + the tail. + """ + # Key is tuple(state, filename), value is a dict(tqdm bar and a previous value) + pbars: Dict[Tuple[str, str], PbarT] = {} + + def close_pbars(): + for pbar in pbars.values(): + pbar["bar"].update(pbar["bar"].total - pbar["past_bytes"]) + pbar["bar"].refresh() + pbar["bar"].close() + + def tail_file(filename) -> Iterator[str]: + """ + Creates a generator to be iterated through, which will return each + line one by one. Will stop tailing the file if the stopping_event is + set. + """ + with open(filename, "r") as file: + current_line = "" + while True: + if stopping_event.is_set(): + close_pbars() + break + + line_bit = file.readline() + if line_bit is not None and not len(line_bit.strip()) == 0: + current_line += line_bit + if current_line.endswith("\n"): + yield current_line + current_line = "" + else: + time.sleep(1) + + # If the file isn't created yet, wait for a few seconds before trying again. + # Can be interrupted with the stopping_event. + while not os.path.exists(os.environ["GIT_LFS_PROGRESS"]): + if stopping_event.is_set(): + close_pbars() + return + + time.sleep(2) + + for line in tail_file(os.environ["GIT_LFS_PROGRESS"]): + try: + state, file_progress, byte_progress, filename = line.split() + except ValueError as error: + # Try/except to ease debugging. See https://github.com/huggingface/huggingface_hub/issues/1373. + raise ValueError(f"Cannot unpack LFS progress line:\n{line}") from error + description = f"{state.capitalize()} file {filename}" + + current_bytes, total_bytes = byte_progress.split("/") + current_bytes_int = int(current_bytes) + total_bytes_int = int(total_bytes) + + pbar = pbars.get((state, filename)) + if pbar is None: + # Initialize progress bar + pbars[(state, filename)] = { + "bar": tqdm( + desc=description, + initial=current_bytes_int, + total=total_bytes_int, + unit="B", + unit_scale=True, + unit_divisor=1024, + name="huggingface_hub.lfs_upload", + ), + "past_bytes": int(current_bytes), + } + else: + # Update progress bar + pbar["bar"].update(current_bytes_int - pbar["past_bytes"]) + pbar["past_bytes"] = current_bytes_int + + current_lfs_progress_value = os.environ.get("GIT_LFS_PROGRESS", "") + + with SoftTemporaryDirectory() as tmpdir: + os.environ["GIT_LFS_PROGRESS"] = os.path.join(tmpdir, "lfs_progress") + logger.debug(f"Following progress in {os.environ['GIT_LFS_PROGRESS']}") + + exit_event = threading.Event() + x = threading.Thread(target=output_progress, args=(exit_event,), daemon=True) + x.start() + + try: + yield + finally: + exit_event.set() + x.join() + + os.environ["GIT_LFS_PROGRESS"] = current_lfs_progress_value + + +class Repository: + """ + Helper class to wrap the git and git-lfs commands. + + The aim is to facilitate interacting with huggingface.co hosted model or + dataset repos, though not a lot here (if any) is actually specific to + huggingface.co. + + + + [`Repository`] is deprecated in favor of the http-based alternatives implemented in + [`HfApi`]. Given its large adoption in legacy code, the complete removal of + [`Repository`] will only happen in release `v1.0`. For more details, please read + https://huggingface.co/docs/huggingface_hub/concepts/git_vs_http. + + + """ + + command_queue: List[CommandInProgress] + + @validate_hf_hub_args + @_deprecate_method( + version="1.0", + message=( + "Please prefer the http-based alternatives instead. Given its large adoption in legacy code, the complete" + " removal is only planned on next major release.\nFor more details, please read" + " https://huggingface.co/docs/huggingface_hub/concepts/git_vs_http." + ), + ) + def __init__( + self, + local_dir: Union[str, Path], + clone_from: Optional[str] = None, + repo_type: Optional[str] = None, + token: Union[bool, str] = True, + git_user: Optional[str] = None, + git_email: Optional[str] = None, + revision: Optional[str] = None, + skip_lfs_files: bool = False, + client: Optional[HfApi] = None, + ): + """ + Instantiate a local clone of a git repo. + + If `clone_from` is set, the repo will be cloned from an existing remote repository. + If the remote repo does not exist, a `EnvironmentError` exception will be thrown. + Please create the remote repo first using [`create_repo`]. + + `Repository` uses the local git credentials by default. If explicitly set, the `token` + or the `git_user`/`git_email` pair will be used instead. + + Args: + local_dir (`str` or `Path`): + path (e.g. `'my_trained_model/'`) to the local directory, where + the `Repository` will be initialized. + clone_from (`str`, *optional*): + Either a repository url or `repo_id`. + Example: + - `"https://huggingface.co/philschmid/playground-tests"` + - `"philschmid/playground-tests"` + repo_type (`str`, *optional*): + To set when cloning a repo from a repo_id. Default is model. + token (`bool` or `str`, *optional*): + A valid authentication token (see https://huggingface.co/settings/token). + If `None` or `True` and machine is logged in (through `huggingface-cli login` + or [`~huggingface_hub.login`]), token will be retrieved from the cache. + If `False`, token is not sent in the request header. + git_user (`str`, *optional*): + will override the `git config user.name` for committing and + pushing files to the hub. + git_email (`str`, *optional*): + will override the `git config user.email` for committing and + pushing files to the hub. + revision (`str`, *optional*): + Revision to checkout after initializing the repository. If the + revision doesn't exist, a branch will be created with that + revision name from the default branch's current HEAD. + skip_lfs_files (`bool`, *optional*, defaults to `False`): + whether to skip git-LFS files or not. + client (`HfApi`, *optional*): + Instance of [`HfApi`] to use when calling the HF Hub API. A new + instance will be created if this is left to `None`. + + Raises: + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If the remote repository set in `clone_from` does not exist. + """ + if isinstance(local_dir, Path): + local_dir = str(local_dir) + os.makedirs(local_dir, exist_ok=True) + self.local_dir = os.path.join(os.getcwd(), local_dir) + self._repo_type = repo_type + self.command_queue = [] + self.skip_lfs_files = skip_lfs_files + self.client = client if client is not None else HfApi() + + self.check_git_versions() + + if isinstance(token, str): + self.huggingface_token: Optional[str] = token + elif token is False: + self.huggingface_token = None + else: + # if `True` -> explicit use of the cached token + # if `None` -> implicit use of the cached token + self.huggingface_token = get_token() + + if clone_from is not None: + self.clone_from(repo_url=clone_from) + else: + if is_git_repo(self.local_dir): + logger.debug("[Repository] is a valid git repo") + else: + raise ValueError("If not specifying `clone_from`, you need to pass Repository a valid git clone.") + + if self.huggingface_token is not None and (git_email is None or git_user is None): + user = self.client.whoami(self.huggingface_token) + + if git_email is None: + git_email = user.get("email") + + if git_user is None: + git_user = user.get("fullname") + + if git_user is not None or git_email is not None: + self.git_config_username_and_email(git_user, git_email) + + self.lfs_enable_largefiles() + self.git_credential_helper_store() + + if revision is not None: + self.git_checkout(revision, create_branch_ok=True) + + # This ensures that all commands exit before exiting the Python runtime. + # This will ensure all pushes register on the hub, even if other errors happen in subsequent operations. + atexit.register(self.wait_for_commands) + + @property + def current_branch(self) -> str: + """ + Returns the current checked out branch. + + Returns: + `str`: Current checked out branch. + """ + try: + result = run_subprocess("git rev-parse --abbrev-ref HEAD", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + return result + + def check_git_versions(self): + """ + Checks that `git` and `git-lfs` can be run. + + Raises: + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `git` or `git-lfs` are not installed. + """ + try: + git_version = run_subprocess("git --version", self.local_dir).stdout.strip() + except FileNotFoundError: + raise EnvironmentError("Looks like you do not have git installed, please install.") + + try: + lfs_version = run_subprocess("git-lfs --version", self.local_dir).stdout.strip() + except FileNotFoundError: + raise EnvironmentError( + "Looks like you do not have git-lfs installed, please install." + " You can install from https://git-lfs.github.com/." + " Then run `git lfs install` (you only have to do this once)." + ) + logger.info(git_version + "\n" + lfs_version) + + @validate_hf_hub_args + def clone_from(self, repo_url: str, token: Union[bool, str, None] = None): + """ + Clone from a remote. If the folder already exists, will try to clone the + repository within it. + + If this folder is a git repository with linked history, will try to + update the repository. + + Args: + repo_url (`str`): + The URL from which to clone the repository + token (`Union[str, bool]`, *optional*): + Whether to use the authentication token. It can be: + - a string which is the token itself + - `False`, which would not use the authentication token + - `True`, which would fetch the authentication token from the + local folder and use it (you should be logged in for this to + work). + - `None`, which would retrieve the value of + `self.huggingface_token`. + + + + Raises the following error: + + - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + if an organization token (starts with "api_org") is passed. Use must use + your own personal access token (see https://hf.co/settings/tokens). + + - [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + if you are trying to clone the repository in a non-empty folder, or if the + `git` operations raise errors. + + + """ + token = ( + token # str -> use it + if isinstance(token, str) + else ( + None # `False` -> explicit no token + if token is False + else self.huggingface_token # `None` or `True` -> use default + ) + ) + if token is not None and token.startswith("api_org"): + raise ValueError( + "You must use your personal access token, not an Organization token" + " (see https://hf.co/settings/tokens)." + ) + + hub_url = self.client.endpoint + if hub_url in repo_url or ("http" not in repo_url and len(repo_url.split("/")) <= 2): + repo_type, namespace, repo_name = repo_type_and_id_from_hf_id(repo_url, hub_url=hub_url) + repo_id = f"{namespace}/{repo_name}" if namespace is not None else repo_name + + if repo_type is not None: + self._repo_type = repo_type + + repo_url = hub_url + "/" + + if self._repo_type in constants.REPO_TYPES_URL_PREFIXES: + repo_url += constants.REPO_TYPES_URL_PREFIXES[self._repo_type] + + if token is not None: + # Add token in git url when provided + scheme = urlparse(repo_url).scheme + repo_url = repo_url.replace(f"{scheme}://", f"{scheme}://user:{token}@") + + repo_url += repo_id + + # For error messages, it's cleaner to show the repo url without the token. + clean_repo_url = re.sub(r"(https?)://.*@", r"\1://", repo_url) + try: + run_subprocess("git lfs install", self.local_dir) + + # checks if repository is initialized in a empty repository or in one with files + if len(os.listdir(self.local_dir)) == 0: + logger.warning(f"Cloning {clean_repo_url} into local empty directory.") + + with _lfs_log_progress(): + env = os.environ.copy() + + if self.skip_lfs_files: + env.update({"GIT_LFS_SKIP_SMUDGE": "1"}) + + run_subprocess( + # 'git lfs clone' is deprecated (will display a warning in the terminal) + # but we still use it as it provides a nicer UX when downloading large + # files (shows progress). + f"{'git clone' if self.skip_lfs_files else 'git lfs clone'} {repo_url} .", + self.local_dir, + env=env, + ) + else: + # Check if the folder is the root of a git repository + if not is_git_repo(self.local_dir): + raise EnvironmentError( + "Tried to clone a repository in a non-empty folder that isn't" + f" a git repository ('{self.local_dir}'). If you really want to" + f" do this, do it manually:\n cd {self.local_dir} && git init" + " && git remote add origin && git pull origin main\n or clone" + " repo to a new folder and move your existing files there" + " afterwards." + ) + + if is_local_clone(self.local_dir, repo_url): + logger.warning( + f"{self.local_dir} is already a clone of {clean_repo_url}." + " Make sure you pull the latest changes with" + " `repo.git_pull()`." + ) + else: + output = run_subprocess("git remote get-url origin", self.local_dir, check=False) + + error_msg = ( + f"Tried to clone {clean_repo_url} in an unrelated git" + " repository.\nIf you believe this is an error, please add" + f" a remote with the following URL: {clean_repo_url}." + ) + if output.returncode == 0: + clean_local_remote_url = re.sub(r"https://.*@", "https://", output.stdout) + error_msg += f"\nLocal path has its origin defined as: {clean_local_remote_url}" + raise EnvironmentError(error_msg) + + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_config_username_and_email(self, git_user: Optional[str] = None, git_email: Optional[str] = None): + """ + Sets git username and email (only in the current repo). + + Args: + git_user (`str`, *optional*): + The username to register through `git`. + git_email (`str`, *optional*): + The email to register through `git`. + """ + try: + if git_user is not None: + run_subprocess("git config user.name".split() + [git_user], self.local_dir) + + if git_email is not None: + run_subprocess(f"git config user.email {git_email}".split(), self.local_dir) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_credential_helper_store(self): + """ + Sets the git credential helper to `store` + """ + try: + run_subprocess("git config credential.helper store", self.local_dir) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_head_hash(self) -> str: + """ + Get commit sha on top of HEAD. + + Returns: + `str`: The current checked out commit SHA. + """ + try: + p = run_subprocess("git rev-parse HEAD", self.local_dir) + return p.stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_remote_url(self) -> str: + """ + Get URL to origin remote. + + Returns: + `str`: The URL of the `origin` remote. + """ + try: + p = run_subprocess("git config --get remote.origin.url", self.local_dir) + url = p.stdout.strip() + # Strip basic auth info. + return re.sub(r"https://.*@", "https://", url) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_head_commit_url(self) -> str: + """ + Get URL to last commit on HEAD. We assume it's been pushed, and the url + scheme is the same one as for GitHub or HuggingFace. + + Returns: + `str`: The URL to the current checked-out commit. + """ + sha = self.git_head_hash() + url = self.git_remote_url() + if url.endswith("/"): + url = url[:-1] + return f"{url}/commit/{sha}" + + def list_deleted_files(self) -> List[str]: + """ + Returns a list of the files that are deleted in the working directory or + index. + + Returns: + `List[str]`: A list of files that have been deleted in the working + directory or index. + """ + try: + git_status = run_subprocess("git status -s", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + if len(git_status) == 0: + return [] + + # Receives a status like the following + # D .gitignore + # D new_file.json + # AD new_file1.json + # ?? new_file2.json + # ?? new_file4.json + + # Strip each line of whitespaces + modified_files_statuses = [status.strip() for status in git_status.split("\n")] + + # Only keep files that are deleted using the D prefix + deleted_files_statuses = [status for status in modified_files_statuses if "D" in status.split()[0]] + + # Remove the D prefix and strip to keep only the relevant filename + deleted_files = [status.split()[-1].strip() for status in deleted_files_statuses] + + return deleted_files + + def lfs_track(self, patterns: Union[str, List[str]], filename: bool = False): + """ + Tell git-lfs to track files according to a pattern. + + Setting the `filename` argument to `True` will treat the arguments as + literal filenames, not as patterns. Any special glob characters in the + filename will be escaped when writing to the `.gitattributes` file. + + Args: + patterns (`Union[str, List[str]]`): + The pattern, or list of patterns, to track with git-lfs. + filename (`bool`, *optional*, defaults to `False`): + Whether to use the patterns as literal filenames. + """ + if isinstance(patterns, str): + patterns = [patterns] + try: + for pattern in patterns: + run_subprocess( + f"git lfs track {'--filename' if filename else ''} {pattern}", + self.local_dir, + ) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def lfs_untrack(self, patterns: Union[str, List[str]]): + """ + Tell git-lfs to untrack those files. + + Args: + patterns (`Union[str, List[str]]`): + The pattern, or list of patterns, to untrack with git-lfs. + """ + if isinstance(patterns, str): + patterns = [patterns] + try: + for pattern in patterns: + run_subprocess("git lfs untrack".split() + [pattern], self.local_dir) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def lfs_enable_largefiles(self): + """ + HF-specific. This enables upload support of files >5GB. + """ + try: + lfs_config = "git config lfs.customtransfer.multipart" + run_subprocess(f"{lfs_config}.path huggingface-cli", self.local_dir) + run_subprocess( + f"{lfs_config}.args {LFS_MULTIPART_UPLOAD_COMMAND}", + self.local_dir, + ) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def auto_track_binary_files(self, pattern: str = ".") -> List[str]: + """ + Automatically track binary files with git-lfs. + + Args: + pattern (`str`, *optional*, defaults to "."): + The pattern with which to track files that are binary. + + Returns: + `List[str]`: List of filenames that are now tracked due to being + binary files + """ + files_to_be_tracked_with_lfs = [] + + deleted_files = self.list_deleted_files() + + for filename in files_to_be_staged(pattern, folder=self.local_dir): + if filename in deleted_files: + continue + + path_to_file = os.path.join(os.getcwd(), self.local_dir, filename) + + if not (is_tracked_with_lfs(path_to_file) or is_git_ignored(path_to_file)): + size_in_mb = os.path.getsize(path_to_file) / (1024 * 1024) + + if size_in_mb >= 10: + logger.warning( + "Parsing a large file to check if binary or not. Tracking large" + " files using `repository.auto_track_large_files` is" + " recommended so as to not load the full file in memory." + ) + + is_binary = is_binary_file(path_to_file) + + if is_binary: + self.lfs_track(filename) + files_to_be_tracked_with_lfs.append(filename) + + # Cleanup the .gitattributes if files were deleted + self.lfs_untrack(deleted_files) + + return files_to_be_tracked_with_lfs + + def auto_track_large_files(self, pattern: str = ".") -> List[str]: + """ + Automatically track large files (files that weigh more than 10MBs) with + git-lfs. + + Args: + pattern (`str`, *optional*, defaults to "."): + The pattern with which to track files that are above 10MBs. + + Returns: + `List[str]`: List of filenames that are now tracked due to their + size. + """ + files_to_be_tracked_with_lfs = [] + + deleted_files = self.list_deleted_files() + + for filename in files_to_be_staged(pattern, folder=self.local_dir): + if filename in deleted_files: + continue + + path_to_file = os.path.join(os.getcwd(), self.local_dir, filename) + size_in_mb = os.path.getsize(path_to_file) / (1024 * 1024) + + if size_in_mb >= 10 and not is_tracked_with_lfs(path_to_file) and not is_git_ignored(path_to_file): + self.lfs_track(filename) + files_to_be_tracked_with_lfs.append(filename) + + # Cleanup the .gitattributes if files were deleted + self.lfs_untrack(deleted_files) + + return files_to_be_tracked_with_lfs + + def lfs_prune(self, recent=False): + """ + git lfs prune + + Args: + recent (`bool`, *optional*, defaults to `False`): + Whether to prune files even if they were referenced by recent + commits. See the following + [link](https://github.com/git-lfs/git-lfs/blob/f3d43f0428a84fc4f1e5405b76b5a73ec2437e65/docs/man/git-lfs-prune.1.ronn#recent-files) + for more information. + """ + try: + with _lfs_log_progress(): + result = run_subprocess(f"git lfs prune {'--recent' if recent else ''}", self.local_dir) + logger.info(result.stdout) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_pull(self, rebase: bool = False, lfs: bool = False): + """ + git pull + + Args: + rebase (`bool`, *optional*, defaults to `False`): + Whether to rebase the current branch on top of the upstream + branch after fetching. + lfs (`bool`, *optional*, defaults to `False`): + Whether to fetch the LFS files too. This option only changes the + behavior when a repository was cloned without fetching the LFS + files; calling `repo.git_pull(lfs=True)` will then fetch the LFS + file from the remote repository. + """ + command = "git pull" if not lfs else "git lfs pull" + if rebase: + command += " --rebase" + try: + with _lfs_log_progress(): + result = run_subprocess(command, self.local_dir) + logger.info(result.stdout) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_add(self, pattern: str = ".", auto_lfs_track: bool = False): + """ + git add + + Setting the `auto_lfs_track` parameter to `True` will automatically + track files that are larger than 10MB with `git-lfs`. + + Args: + pattern (`str`, *optional*, defaults to "."): + The pattern with which to add files to staging. + auto_lfs_track (`bool`, *optional*, defaults to `False`): + Whether to automatically track large and binary files with + git-lfs. Any file over 10MB in size, or in binary format, will + be automatically tracked. + """ + if auto_lfs_track: + # Track files according to their size (>=10MB) + tracked_files = self.auto_track_large_files(pattern) + + # Read the remaining files and track them if they're binary + tracked_files.extend(self.auto_track_binary_files(pattern)) + + if tracked_files: + logger.warning( + f"Adding files tracked by Git LFS: {tracked_files}. This may take a" + " bit of time if the files are large." + ) + + try: + result = run_subprocess("git add -v".split() + [pattern], self.local_dir) + logger.info(f"Adding to index:\n{result.stdout}\n") + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def git_commit(self, commit_message: str = "commit files to HF hub"): + """ + git commit + + Args: + commit_message (`str`, *optional*, defaults to "commit files to HF hub"): + The message attributed to the commit. + """ + try: + result = run_subprocess("git commit -v -m".split() + [commit_message], self.local_dir) + logger.info(f"Committed:\n{result.stdout}\n") + except subprocess.CalledProcessError as exc: + if len(exc.stderr) > 0: + raise EnvironmentError(exc.stderr) + else: + raise EnvironmentError(exc.stdout) + + def git_push( + self, + upstream: Optional[str] = None, + blocking: bool = True, + auto_lfs_prune: bool = False, + ) -> Union[str, Tuple[str, CommandInProgress]]: + """ + git push + + If used without setting `blocking`, will return url to commit on remote + repo. If used with `blocking=True`, will return a tuple containing the + url to commit and the command object to follow for information about the + process. + + Args: + upstream (`str`, *optional*): + Upstream to which this should push. If not specified, will push + to the lastly defined upstream or to the default one (`origin + main`). + blocking (`bool`, *optional*, defaults to `True`): + Whether the function should return only when the push has + finished. Setting this to `False` will return an + `CommandInProgress` object which has an `is_done` property. This + property will be set to `True` when the push is finished. + auto_lfs_prune (`bool`, *optional*, defaults to `False`): + Whether to automatically prune files once they have been pushed + to the remote. + """ + command = "git push" + + if upstream: + command += f" --set-upstream {upstream}" + + number_of_commits = commits_to_push(self.local_dir, upstream) + + if number_of_commits > 1: + logger.warning(f"Several commits ({number_of_commits}) will be pushed upstream.") + if blocking: + logger.warning("The progress bars may be unreliable.") + + try: + with _lfs_log_progress(): + process = subprocess.Popen( + command.split(), + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + encoding="utf-8", + cwd=self.local_dir, + ) + + if blocking: + stdout, stderr = process.communicate() + return_code = process.poll() + process.kill() + + if len(stderr): + logger.warning(stderr) + + if return_code: + raise subprocess.CalledProcessError(return_code, process.args, output=stdout, stderr=stderr) + + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + if not blocking: + + def status_method(): + status = process.poll() + if status is None: + return -1 + else: + return status + + command_in_progress = CommandInProgress( + "push", + is_done_method=lambda: process.poll() is not None, + status_method=status_method, + process=process, + post_method=self.lfs_prune if auto_lfs_prune else None, + ) + + self.command_queue.append(command_in_progress) + + return self.git_head_commit_url(), command_in_progress + + if auto_lfs_prune: + self.lfs_prune() + + return self.git_head_commit_url() + + def git_checkout(self, revision: str, create_branch_ok: bool = False): + """ + git checkout a given revision + + Specifying `create_branch_ok` to `True` will create the branch to the + given revision if that revision doesn't exist. + + Args: + revision (`str`): + The revision to checkout. + create_branch_ok (`str`, *optional*, defaults to `False`): + Whether creating a branch named with the `revision` passed at + the current checked-out reference if `revision` isn't an + existing revision is allowed. + """ + try: + result = run_subprocess(f"git checkout {revision}", self.local_dir) + logger.warning(f"Checked out {revision} from {self.current_branch}.") + logger.warning(result.stdout) + except subprocess.CalledProcessError as exc: + if not create_branch_ok: + raise EnvironmentError(exc.stderr) + else: + try: + result = run_subprocess(f"git checkout -b {revision}", self.local_dir) + logger.warning( + f"Revision `{revision}` does not exist. Created and checked out branch `{revision}`." + ) + logger.warning(result.stdout) + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def tag_exists(self, tag_name: str, remote: Optional[str] = None) -> bool: + """ + Check if a tag exists or not. + + Args: + tag_name (`str`): + The name of the tag to check. + remote (`str`, *optional*): + Whether to check if the tag exists on a remote. This parameter + should be the identifier of the remote. + + Returns: + `bool`: Whether the tag exists. + """ + if remote: + try: + result = run_subprocess(f"git ls-remote origin refs/tags/{tag_name}", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + return len(result) != 0 + else: + try: + git_tags = run_subprocess("git tag", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + git_tags = git_tags.split("\n") + return tag_name in git_tags + + def delete_tag(self, tag_name: str, remote: Optional[str] = None) -> bool: + """ + Delete a tag, both local and remote, if it exists + + Args: + tag_name (`str`): + The tag name to delete. + remote (`str`, *optional*): + The remote on which to delete the tag. + + Returns: + `bool`: `True` if deleted, `False` if the tag didn't exist. + If remote is not passed, will just be updated locally + """ + delete_locally = True + delete_remotely = True + + if not self.tag_exists(tag_name): + delete_locally = False + + if not self.tag_exists(tag_name, remote=remote): + delete_remotely = False + + if delete_locally: + try: + run_subprocess(["git", "tag", "-d", tag_name], self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + if remote and delete_remotely: + try: + run_subprocess(f"git push {remote} --delete {tag_name}", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + return True + + def add_tag(self, tag_name: str, message: Optional[str] = None, remote: Optional[str] = None): + """ + Add a tag at the current head and push it + + If remote is None, will just be updated locally + + If no message is provided, the tag will be lightweight. if a message is + provided, the tag will be annotated. + + Args: + tag_name (`str`): + The name of the tag to be added. + message (`str`, *optional*): + The message that accompanies the tag. The tag will turn into an + annotated tag if a message is passed. + remote (`str`, *optional*): + The remote on which to add the tag. + """ + if message: + tag_args = ["git", "tag", "-a", tag_name, "-m", message] + else: + tag_args = ["git", "tag", tag_name] + + try: + run_subprocess(tag_args, self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + if remote: + try: + run_subprocess(f"git push {remote} {tag_name}", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + def is_repo_clean(self) -> bool: + """ + Return whether or not the git status is clean or not + + Returns: + `bool`: `True` if the git status is clean, `False` otherwise. + """ + try: + git_status = run_subprocess("git status --porcelain", self.local_dir).stdout.strip() + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + return len(git_status) == 0 + + def push_to_hub( + self, + commit_message: str = "commit files to HF hub", + blocking: bool = True, + clean_ok: bool = True, + auto_lfs_prune: bool = False, + ) -> Union[None, str, Tuple[str, CommandInProgress]]: + """ + Helper to add, commit, and push files to remote repository on the + HuggingFace Hub. Will automatically track large files (>10MB). + + Args: + commit_message (`str`): + Message to use for the commit. + blocking (`bool`, *optional*, defaults to `True`): + Whether the function should return only when the `git push` has + finished. + clean_ok (`bool`, *optional*, defaults to `True`): + If True, this function will return None if the repo is + untouched. Default behavior is to fail because the git command + fails. + auto_lfs_prune (`bool`, *optional*, defaults to `False`): + Whether to automatically prune files once they have been pushed + to the remote. + """ + if clean_ok and self.is_repo_clean(): + logger.info("Repo currently clean. Ignoring push_to_hub") + return None + self.git_add(auto_lfs_track=True) + self.git_commit(commit_message) + return self.git_push( + upstream=f"origin {self.current_branch}", + blocking=blocking, + auto_lfs_prune=auto_lfs_prune, + ) + + @contextmanager + def commit( + self, + commit_message: str, + branch: Optional[str] = None, + track_large_files: bool = True, + blocking: bool = True, + auto_lfs_prune: bool = False, + ): + """ + Context manager utility to handle committing to a repository. This + automatically tracks large files (>10Mb) with git-lfs. Set the + `track_large_files` argument to `False` if you wish to ignore that + behavior. + + Args: + commit_message (`str`): + Message to use for the commit. + branch (`str`, *optional*): + The branch on which the commit will appear. This branch will be + checked-out before any operation. + track_large_files (`bool`, *optional*, defaults to `True`): + Whether to automatically track large files or not. Will do so by + default. + blocking (`bool`, *optional*, defaults to `True`): + Whether the function should return only when the `git push` has + finished. + auto_lfs_prune (`bool`, defaults to `True`): + Whether to automatically prune files once they have been pushed + to the remote. + + Examples: + + ```python + >>> with Repository( + ... "text-files", + ... clone_from="/text-files", + ... token=True, + >>> ).commit("My first file :)"): + ... with open("file.txt", "w+") as f: + ... f.write(json.dumps({"hey": 8})) + + >>> import torch + + >>> model = torch.nn.Transformer() + >>> with Repository( + ... "torch-model", + ... clone_from="/torch-model", + ... token=True, + >>> ).commit("My cool model :)"): + ... torch.save(model.state_dict(), "model.pt") + ``` + + """ + + files_to_stage = files_to_be_staged(".", folder=self.local_dir) + + if len(files_to_stage): + files_in_msg = str(files_to_stage[:5])[:-1] + ", ...]" if len(files_to_stage) > 5 else str(files_to_stage) + logger.error( + "There exists some updated files in the local repository that are not" + f" committed: {files_in_msg}. This may lead to errors if checking out" + " a branch. These files and their modifications will be added to the" + " current commit." + ) + + if branch is not None: + self.git_checkout(branch, create_branch_ok=True) + + if is_tracked_upstream(self.local_dir): + logger.warning("Pulling changes ...") + self.git_pull(rebase=True) + else: + logger.warning(f"The current branch has no upstream branch. Will push to 'origin {self.current_branch}'") + + current_working_directory = os.getcwd() + os.chdir(os.path.join(current_working_directory, self.local_dir)) + + try: + yield self + finally: + self.git_add(auto_lfs_track=track_large_files) + + try: + self.git_commit(commit_message) + except OSError as e: + # If no changes are detected, there is nothing to commit. + if "nothing to commit" not in str(e): + raise e + + try: + self.git_push( + upstream=f"origin {self.current_branch}", + blocking=blocking, + auto_lfs_prune=auto_lfs_prune, + ) + except OSError as e: + # If no changes are detected, there is nothing to commit. + if "could not read Username" in str(e): + raise OSError("Couldn't authenticate user for push. Did you set `token` to `True`?") from e + else: + raise e + + os.chdir(current_working_directory) + + def repocard_metadata_load(self) -> Optional[Dict]: + filepath = os.path.join(self.local_dir, constants.REPOCARD_NAME) + if os.path.isfile(filepath): + return metadata_load(filepath) + return None + + def repocard_metadata_save(self, data: Dict) -> None: + return metadata_save(os.path.join(self.local_dir, constants.REPOCARD_NAME), data) + + @property + def commands_failed(self): + """ + Returns the asynchronous commands that failed. + """ + return [c for c in self.command_queue if c.status > 0] + + @property + def commands_in_progress(self): + """ + Returns the asynchronous commands that are currently in progress. + """ + return [c for c in self.command_queue if not c.is_done] + + def wait_for_commands(self): + """ + Blocking method: blocks all subsequent execution until all commands have + been processed. + """ + index = 0 + for command_failed in self.commands_failed: + logger.error(f"The {command_failed.title} command with PID {command_failed._process.pid} failed.") + logger.error(command_failed.stderr) + + while self.commands_in_progress: + if index % 10 == 0: + logger.warning( + f"Waiting for the following commands to finish before shutting down: {self.commands_in_progress}." + ) + + index += 1 + + time.sleep(1) diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__init__.py b/lib/python3.12/site-packages/huggingface_hub/serialization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8949a22a5f65ab29b7df65aa6a9df9bce0544b7e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/serialization/__init__.py @@ -0,0 +1,27 @@ +# 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. +# ruff: noqa: F401 +"""Contains helpers to serialize tensors.""" + +from ._base import StateDictSplit, split_state_dict_into_shards_factory +from ._tensorflow import get_tf_storage_size, split_tf_state_dict_into_shards +from ._torch import ( + get_torch_storage_id, + get_torch_storage_size, + load_state_dict_from_file, + load_torch_model, + save_torch_model, + save_torch_state_dict, + split_torch_state_dict_into_shards, +) diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ed54e907d8f633dd7ca60503b706fd699b2238a Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_base.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf420bcae8bbb3455de44efe2183ccfbf2dd1b73 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_dduf.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_dduf.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..89c49f21ad9d1e776db4385b9ff99fbcd8b2355d Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_dduf.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_tensorflow.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_tensorflow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03438f0e84cf351384b0d123712e68f28d0f3ad0 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_tensorflow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_torch.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb91716a787f8eddbd5f1f58795a296f2a6caaa8 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/serialization/__pycache__/_torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/_base.py b/lib/python3.12/site-packages/huggingface_hub/serialization/_base.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b6454a90e1942854dd0a095a59c92794323279 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/serialization/_base.py @@ -0,0 +1,210 @@ +# 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. +"""Contains helpers to split tensors into shards.""" + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, TypeVar, Union + +from .. import logging + + +TensorT = TypeVar("TensorT") +TensorSizeFn_T = Callable[[TensorT], int] +StorageIDFn_T = Callable[[TensorT], Optional[Any]] + +MAX_SHARD_SIZE = "5GB" +SIZE_UNITS = { + "TB": 10**12, + "GB": 10**9, + "MB": 10**6, + "KB": 10**3, +} + + +logger = logging.get_logger(__file__) + + +@dataclass +class StateDictSplit: + is_sharded: bool = field(init=False) + metadata: Dict[str, Any] + filename_to_tensors: Dict[str, List[str]] + tensor_to_filename: Dict[str, str] + + def __post_init__(self): + self.is_sharded = len(self.filename_to_tensors) > 1 + + +def split_state_dict_into_shards_factory( + state_dict: Dict[str, TensorT], + *, + get_storage_size: TensorSizeFn_T, + filename_pattern: str, + get_storage_id: StorageIDFn_T = lambda tensor: None, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, +) -> StateDictSplit: + """ + Split a model state dictionary in shards so that each shard is smaller than a given size. + + The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization + made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we + have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not + [6+2+2GB], [6+2GB], [6GB]. + + + + If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a + size greater than `max_shard_size`. + + + + Args: + state_dict (`Dict[str, Tensor]`): + The state dictionary to save. + get_storage_size (`Callable[[Tensor], int]`): + A function that returns the size of a tensor when saved on disk in bytes. + get_storage_id (`Callable[[Tensor], Optional[Any]]`, *optional*): + A function that returns a unique identifier to a tensor storage. Multiple different tensors can share the + same underlying storage. This identifier is guaranteed to be unique and constant for this tensor's storage + during its lifetime. Two tensor storages with non-overlapping lifetimes may have the same id. + filename_pattern (`str`, *optional*): + The pattern to generate the files names in which the model will be saved. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + max_shard_size (`int` or `str`, *optional*): + The maximum size of each shard, in bytes. Defaults to 5GB. + + Returns: + [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them. + """ + storage_id_to_tensors: Dict[Any, List[str]] = {} + + shard_list: List[Dict[str, TensorT]] = [] + current_shard: Dict[str, TensorT] = {} + current_shard_size = 0 + total_size = 0 + + if isinstance(max_shard_size, str): + max_shard_size = parse_size_to_int(max_shard_size) + + for key, tensor in state_dict.items(): + # when bnb serialization is used the weights in the state dict can be strings + # check: https://github.com/huggingface/transformers/pull/24416 for more details + if isinstance(tensor, str): + logger.info("Skipping tensor %s as it is a string (bnb serialization)", key) + continue + + # If a `tensor` shares the same underlying storage as another tensor, we put `tensor` in the same `block` + storage_id = get_storage_id(tensor) + if storage_id is not None: + if storage_id in storage_id_to_tensors: + # We skip this tensor for now and will reassign to correct shard later + storage_id_to_tensors[storage_id].append(key) + continue + else: + # This is the first tensor with this storage_id, we create a new entry + # in the storage_id_to_tensors dict => we will assign the shard id later + storage_id_to_tensors[storage_id] = [key] + + # Compute tensor size + tensor_size = get_storage_size(tensor) + + # If this tensor is bigger than the maximal size, we put it in its own shard + if tensor_size > max_shard_size: + total_size += tensor_size + shard_list.append({key: tensor}) + continue + + # If this tensor is going to tip up over the maximal size, we split. + # Current shard already has some tensors, we add it to the list of shards and create a new one. + if current_shard_size + tensor_size > max_shard_size: + shard_list.append(current_shard) + current_shard = {} + current_shard_size = 0 + + # Add the tensor to the current shard + current_shard[key] = tensor + current_shard_size += tensor_size + total_size += tensor_size + + # Add the last shard + if len(current_shard) > 0: + shard_list.append(current_shard) + nb_shards = len(shard_list) + + # Loop over the tensors that share the same storage and assign them together + for storage_id, keys in storage_id_to_tensors.items(): + # Let's try to find the shard where the first tensor of this storage is and put all tensors in the same shard + for shard in shard_list: + if keys[0] in shard: + for key in keys: + shard[key] = state_dict[key] + break + + # If we only have one shard, we return it => no need to build the index + if nb_shards == 1: + filename = filename_pattern.format(suffix="") + return StateDictSplit( + metadata={"total_size": total_size}, + filename_to_tensors={filename: list(state_dict.keys())}, + tensor_to_filename={key: filename for key in state_dict.keys()}, + ) + + # Now that each tensor is assigned to a shard, let's assign a filename to each shard + tensor_name_to_filename = {} + filename_to_tensors = {} + for idx, shard in enumerate(shard_list): + filename = filename_pattern.format(suffix=f"-{idx + 1:05d}-of-{nb_shards:05d}") + for key in shard: + tensor_name_to_filename[key] = filename + filename_to_tensors[filename] = list(shard.keys()) + + # Build the index and return + return StateDictSplit( + metadata={"total_size": total_size}, + filename_to_tensors=filename_to_tensors, + tensor_to_filename=tensor_name_to_filename, + ) + + +def parse_size_to_int(size_as_str: str) -> int: + """ + Parse a size expressed as a string with digits and unit (like `"5MB"`) to an integer (in bytes). + + Supported units are "TB", "GB", "MB", "KB". + + Args: + size_as_str (`str`): The size to convert. Will be directly returned if an `int`. + + Example: + + ```py + >>> parse_size_to_int("5MB") + 5000000 + ``` + """ + size_as_str = size_as_str.strip() + + # Parse unit + unit = size_as_str[-2:].upper() + if unit not in SIZE_UNITS: + raise ValueError(f"Unit '{unit}' not supported. Supported units are TB, GB, MB, KB. Got '{size_as_str}'.") + multiplier = SIZE_UNITS[unit] + + # Parse value + try: + value = float(size_as_str[:-2].strip()) + except ValueError as e: + raise ValueError(f"Could not parse the size value from '{size_as_str}': {e}") from e + + return int(value * multiplier) diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/_dduf.py b/lib/python3.12/site-packages/huggingface_hub/serialization/_dduf.py new file mode 100644 index 0000000000000000000000000000000000000000..a1debadb3ac8a45716f0359b932dc065f09edb84 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/serialization/_dduf.py @@ -0,0 +1,387 @@ +import json +import logging +import mmap +import os +import shutil +import zipfile +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, Generator, Iterable, Tuple, Union + +from ..errors import DDUFCorruptedFileError, DDUFExportError, DDUFInvalidEntryNameError + + +logger = logging.getLogger(__name__) + +DDUF_ALLOWED_ENTRIES = { + # Allowed file extensions in a DDUF file + ".json", + ".model", + ".safetensors", + ".txt", +} + +DDUF_FOLDER_REQUIRED_ENTRIES = { + # Each folder must contain at least one of these entries + "config.json", + "tokenizer_config.json", + "preprocessor_config.json", + "scheduler_config.json", +} + + +@dataclass +class DDUFEntry: + """Object representing a file entry in a DDUF file. + + See [`read_dduf_file`] for how to read a DDUF file. + + Attributes: + filename (str): + The name of the file in the DDUF archive. + offset (int): + The offset of the file in the DDUF archive. + length (int): + The length of the file in the DDUF archive. + dduf_path (str): + The path to the DDUF archive (for internal use). + """ + + filename: str + length: int + offset: int + + dduf_path: Path = field(repr=False) + + @contextmanager + def as_mmap(self) -> Generator[bytes, None, None]: + """Open the file as a memory-mapped file. + + Useful to load safetensors directly from the file. + + Example: + ```py + >>> import safetensors.torch + >>> with entry.as_mmap() as mm: + ... tensors = safetensors.torch.load(mm) + ``` + """ + with self.dduf_path.open("rb") as f: + with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mm: + yield mm[self.offset : self.offset + self.length] + + def read_text(self, encoding: str = "utf-8") -> str: + """Read the file as text. + + Useful for '.txt' and '.json' entries. + + Example: + ```py + >>> import json + >>> index = json.loads(entry.read_text()) + ``` + """ + with self.dduf_path.open("rb") as f: + f.seek(self.offset) + return f.read(self.length).decode(encoding=encoding) + + +def read_dduf_file(dduf_path: Union[os.PathLike, str]) -> Dict[str, DDUFEntry]: + """ + Read a DDUF file and return a dictionary of entries. + + Only the metadata is read, the data is not loaded in memory. + + Args: + dduf_path (`str` or `os.PathLike`): + The path to the DDUF file to read. + + Returns: + `Dict[str, DDUFEntry]`: + A dictionary of [`DDUFEntry`] indexed by filename. + + Raises: + - [`DDUFCorruptedFileError`]: If the DDUF file is corrupted (i.e. doesn't follow the DDUF format). + + Example: + ```python + >>> import json + >>> import safetensors.torch + >>> from huggingface_hub import read_dduf_file + + # Read DDUF metadata + >>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf") + + # Returns a mapping filename <> DDUFEntry + >>> dduf_entries["model_index.json"] + DDUFEntry(filename='model_index.json', offset=66, length=587) + + # Load model index as JSON + >>> json.loads(dduf_entries["model_index.json"].read_text()) + {'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', ... + + # Load VAE weights using safetensors + >>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm: + ... state_dict = safetensors.torch.load(mm) + ``` + """ + entries = {} + dduf_path = Path(dduf_path) + logger.info(f"Reading DDUF file {dduf_path}") + with zipfile.ZipFile(str(dduf_path), "r") as zf: + for info in zf.infolist(): + logger.debug(f"Reading entry {info.filename}") + if info.compress_type != zipfile.ZIP_STORED: + raise DDUFCorruptedFileError("Data must not be compressed in DDUF file.") + + try: + _validate_dduf_entry_name(info.filename) + except DDUFInvalidEntryNameError as e: + raise DDUFCorruptedFileError(f"Invalid entry name in DDUF file: {info.filename}") from e + + offset = _get_data_offset(zf, info) + + entries[info.filename] = DDUFEntry( + filename=info.filename, offset=offset, length=info.file_size, dduf_path=dduf_path + ) + + # Consistency checks on the DDUF file + if "model_index.json" not in entries: + raise DDUFCorruptedFileError("Missing required 'model_index.json' entry in DDUF file.") + index = json.loads(entries["model_index.json"].read_text()) + _validate_dduf_structure(index, entries.keys()) + + logger.info(f"Done reading DDUF file {dduf_path}. Found {len(entries)} entries") + return entries + + +def export_entries_as_dduf( + dduf_path: Union[str, os.PathLike], entries: Iterable[Tuple[str, Union[str, Path, bytes]]] +) -> None: + """Write a DDUF file from an iterable of entries. + + This is a lower-level helper than [`export_folder_as_dduf`] that allows more flexibility when serializing data. + In particular, you don't need to save the data on disk before exporting it in the DDUF file. + + Args: + dduf_path (`str` or `os.PathLike`): + The path to the DDUF file to write. + entries (`Iterable[Tuple[str, Union[str, Path, bytes]]]`): + An iterable of entries to write in the DDUF file. Each entry is a tuple with the filename and the content. + The filename should be the path to the file in the DDUF archive. + The content can be a string or a pathlib.Path representing a path to a file on the local disk or directly the content as bytes. + + Raises: + - [`DDUFExportError`]: If anything goes wrong during the export (e.g. invalid entry name, missing 'model_index.json', etc.). + + Example: + ```python + # Export specific files from the local disk. + >>> from huggingface_hub import export_entries_as_dduf + >>> export_entries_as_dduf( + ... dduf_path="stable-diffusion-v1-4-FP16.dduf", + ... entries=[ # List entries to add to the DDUF file (here, only FP16 weights) + ... ("model_index.json", "path/to/model_index.json"), + ... ("vae/config.json", "path/to/vae/config.json"), + ... ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"), + ... ("text_encoder/config.json", "path/to/text_encoder/config.json"), + ... ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"), + ... # ... add more entries here + ... ] + ... ) + ``` + + ```python + # Export state_dicts one by one from a loaded pipeline + >>> from diffusers import DiffusionPipeline + >>> from typing import Generator, Tuple + >>> import safetensors.torch + >>> from huggingface_hub import export_entries_as_dduf + >>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4") + ... # ... do some work with the pipeline + + >>> def as_entries(pipe: DiffusionPipeline) -> Generator[Tuple[str, bytes], None, None]: + ... # Build an generator that yields the entries to add to the DDUF file. + ... # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file. + ... # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time) + ... yield "vae/config.json", pipe.vae.to_json_string().encode() + ... yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict()) + ... yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode() + ... yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict()) + ... # ... add more entries here + + >>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe)) + ``` + """ + logger.info(f"Exporting DDUF file '{dduf_path}'") + filenames = set() + index = None + with zipfile.ZipFile(str(dduf_path), "w", zipfile.ZIP_STORED) as archive: + for filename, content in entries: + if filename in filenames: + raise DDUFExportError(f"Can't add duplicate entry: {filename}") + filenames.add(filename) + + if filename == "model_index.json": + try: + index = json.loads(_load_content(content).decode()) + except json.JSONDecodeError as e: + raise DDUFExportError("Failed to parse 'model_index.json'.") from e + + try: + filename = _validate_dduf_entry_name(filename) + except DDUFInvalidEntryNameError as e: + raise DDUFExportError(f"Invalid entry name: {filename}") from e + logger.debug(f"Adding entry '{filename}' to DDUF file") + _dump_content_in_archive(archive, filename, content) + + # Consistency checks on the DDUF file + if index is None: + raise DDUFExportError("Missing required 'model_index.json' entry in DDUF file.") + try: + _validate_dduf_structure(index, filenames) + except DDUFCorruptedFileError as e: + raise DDUFExportError("Invalid DDUF file structure.") from e + + logger.info(f"Done writing DDUF file {dduf_path}") + + +def export_folder_as_dduf(dduf_path: Union[str, os.PathLike], folder_path: Union[str, os.PathLike]) -> None: + """ + Export a folder as a DDUF file. + + AUses [`export_entries_as_dduf`] under the hood. + + Args: + dduf_path (`str` or `os.PathLike`): + The path to the DDUF file to write. + folder_path (`str` or `os.PathLike`): + The path to the folder containing the diffusion model. + + Example: + ```python + >>> from huggingface_hub import export_folder_as_dduf + >>> export_folder_as_dduf(dduf_path="FLUX.1-dev.dduf", folder_path="path/to/FLUX.1-dev") + ``` + """ + folder_path = Path(folder_path) + + def _iterate_over_folder() -> Iterable[Tuple[str, Path]]: + for path in Path(folder_path).glob("**/*"): + if not path.is_file(): + continue + if path.suffix not in DDUF_ALLOWED_ENTRIES: + logger.debug(f"Skipping file '{path}' (file type not allowed)") + continue + path_in_archive = path.relative_to(folder_path) + if len(path_in_archive.parts) >= 3: + logger.debug(f"Skipping file '{path}' (nested directories not allowed)") + continue + yield path_in_archive.as_posix(), path + + export_entries_as_dduf(dduf_path, _iterate_over_folder()) + + +def _dump_content_in_archive(archive: zipfile.ZipFile, filename: str, content: Union[str, os.PathLike, bytes]) -> None: + with archive.open(filename, "w", force_zip64=True) as archive_fh: + if isinstance(content, (str, Path)): + content_path = Path(content) + with content_path.open("rb") as content_fh: + shutil.copyfileobj(content_fh, archive_fh, 1024 * 1024 * 8) # type: ignore[misc] + elif isinstance(content, bytes): + archive_fh.write(content) + else: + raise DDUFExportError(f"Invalid content type for {filename}. Must be str, Path or bytes.") + + +def _load_content(content: Union[str, Path, bytes]) -> bytes: + """Load the content of an entry as bytes. + + Used only for small checks (not to dump content into archive). + """ + if isinstance(content, (str, Path)): + return Path(content).read_bytes() + elif isinstance(content, bytes): + return content + else: + raise DDUFExportError(f"Invalid content type. Must be str, Path or bytes. Got {type(content)}.") + + +def _validate_dduf_entry_name(entry_name: str) -> str: + if "." + entry_name.split(".")[-1] not in DDUF_ALLOWED_ENTRIES: + raise DDUFInvalidEntryNameError(f"File type not allowed: {entry_name}") + if "\\" in entry_name: + raise DDUFInvalidEntryNameError(f"Entry names must use UNIX separators ('/'). Got {entry_name}.") + entry_name = entry_name.strip("/") + if entry_name.count("/") > 1: + raise DDUFInvalidEntryNameError(f"DDUF only supports 1 level of directory. Got {entry_name}.") + return entry_name + + +def _validate_dduf_structure(index: Any, entry_names: Iterable[str]) -> None: + """ + Consistency checks on the DDUF file structure. + + Rules: + - The 'model_index.json' entry is required and must contain a dictionary. + - Each folder name must correspond to an entry in 'model_index.json'. + - Each folder must contain at least a config file ('config.json', 'tokenizer_config.json', 'preprocessor_config.json', 'scheduler_config.json'). + + Args: + index (Any): + The content of the 'model_index.json' entry. + entry_names (Iterable[str]): + The list of entry names in the DDUF file. + + Raises: + - [`DDUFCorruptedFileError`]: If the DDUF file is corrupted (i.e. doesn't follow the DDUF format). + """ + if not isinstance(index, dict): + raise DDUFCorruptedFileError(f"Invalid 'model_index.json' content. Must be a dictionary. Got {type(index)}.") + + dduf_folders = {entry.split("/")[0] for entry in entry_names if "/" in entry} + for folder in dduf_folders: + if folder not in index: + raise DDUFCorruptedFileError(f"Missing required entry '{folder}' in 'model_index.json'.") + if not any(f"{folder}/{required_entry}" in entry_names for required_entry in DDUF_FOLDER_REQUIRED_ENTRIES): + raise DDUFCorruptedFileError( + f"Missing required file in folder '{folder}'. Must contains at least one of {DDUF_FOLDER_REQUIRED_ENTRIES}." + ) + + +def _get_data_offset(zf: zipfile.ZipFile, info: zipfile.ZipInfo) -> int: + """ + Calculate the data offset for a file in a ZIP archive. + + Args: + zf (`zipfile.ZipFile`): + The opened ZIP file. Must be opened in read mode. + info (`zipfile.ZipInfo`): + The file info. + + Returns: + int: The offset of the file data in the ZIP archive. + """ + if zf.fp is None: + raise DDUFCorruptedFileError("ZipFile object must be opened in read mode.") + + # Step 1: Get the local file header offset + header_offset = info.header_offset + + # Step 2: Read the local file header + zf.fp.seek(header_offset) + local_file_header = zf.fp.read(30) # Fixed-size part of the local header + + if len(local_file_header) < 30: + raise DDUFCorruptedFileError("Incomplete local file header.") + + # Step 3: Parse the header fields to calculate the start of file data + # Local file header: https://en.wikipedia.org/wiki/ZIP_(file_format)#File_headers + filename_len = int.from_bytes(local_file_header[26:28], "little") + extra_field_len = int.from_bytes(local_file_header[28:30], "little") + + # Data offset is after the fixed header, filename, and extra fields + data_offset = header_offset + 30 + filename_len + extra_field_len + + return data_offset diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/_tensorflow.py b/lib/python3.12/site-packages/huggingface_hub/serialization/_tensorflow.py new file mode 100644 index 0000000000000000000000000000000000000000..59ed8110b28f4891d67e754fdfbfa47a26f85be1 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/serialization/_tensorflow.py @@ -0,0 +1,95 @@ +# 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. +"""Contains tensorflow-specific helpers.""" + +import math +import re +from typing import TYPE_CHECKING, Dict, Union + +from .. import constants +from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory + + +if TYPE_CHECKING: + import tensorflow as tf + + +def split_tf_state_dict_into_shards( + state_dict: Dict[str, "tf.Tensor"], + *, + filename_pattern: str = constants.TF2_WEIGHTS_FILE_PATTERN, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, +) -> StateDictSplit: + """ + Split a model state dictionary in shards so that each shard is smaller than a given size. + + The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization + made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we + have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not + [6+2+2GB], [6+2GB], [6GB]. + + + + If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a + size greater than `max_shard_size`. + + + + Args: + state_dict (`Dict[str, Tensor]`): + The state dictionary to save. + filename_pattern (`str`, *optional*): + The pattern to generate the files names in which the model will be saved. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"tf_model{suffix}.h5"`. + max_shard_size (`int` or `str`, *optional*): + The maximum size of each shard, in bytes. Defaults to 5GB. + + Returns: + [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them. + """ + return split_state_dict_into_shards_factory( + state_dict, + max_shard_size=max_shard_size, + filename_pattern=filename_pattern, + get_storage_size=get_tf_storage_size, + ) + + +def get_tf_storage_size(tensor: "tf.Tensor") -> int: + # Return `math.ceil` since dtype byte size can be a float (e.g., 0.125 for tf.bool). + # Better to overestimate than underestimate. + return math.ceil(tensor.numpy().size * _dtype_byte_size_tf(tensor.dtype)) + + +def _dtype_byte_size_tf(dtype) -> float: + """ + Returns the size (in bytes) occupied by one parameter of type `dtype`. + Taken from https://github.com/huggingface/transformers/blob/74d9d0cebb0263a3f8ab9c280569170cc74651d0/src/transformers/modeling_tf_utils.py#L608. + NOTE: why not `tensor.numpy().nbytes`? + Example: + ```py + >>> _dtype_byte_size(tf.float32) + 4 + ``` + """ + import tensorflow as tf + + if dtype == tf.bool: + return 1 / 8 + bit_search = re.search(r"[^\d](\d+)$", dtype.name) + if bit_search is None: + raise ValueError(f"`dtype` is not a valid dtype: {dtype}.") + bit_size = int(bit_search.groups()[0]) + return bit_size // 8 diff --git a/lib/python3.12/site-packages/huggingface_hub/serialization/_torch.py b/lib/python3.12/site-packages/huggingface_hub/serialization/_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..daa4154b45b4b17cbe5c8d841e29cbbc9aedd05d --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/serialization/_torch.py @@ -0,0 +1,1033 @@ +# 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. +"""Contains pytorch-specific helpers.""" + +import importlib +import json +import os +import re +from collections import defaultdict, namedtuple +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Set, Tuple, Union + +from packaging import version + +from .. import constants, logging +from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory + + +logger = logging.get_logger(__file__) + +if TYPE_CHECKING: + import torch + +# SAVING + + +def save_torch_model( + model: "torch.nn.Module", + save_directory: Union[str, Path], + *, + filename_pattern: Optional[str] = None, + force_contiguous: bool = True, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, + metadata: Optional[Dict[str, str]] = None, + safe_serialization: bool = True, + is_main_process: bool = True, + shared_tensors_to_discard: Optional[List[str]] = None, +): + """ + Saves a given torch model to disk, handling sharding and shared tensors issues. + + See also [`save_torch_state_dict`] to save a state dict with more flexibility. + + For more information about tensor sharing, check out [this guide](https://huggingface.co/docs/safetensors/torch_shared_tensors). + + The model state dictionary is split into shards so that each shard is smaller than a given size. The shards are + saved in the `save_directory` with the given `filename_pattern`. If the model is too big to fit in a single shard, + an index file is saved in the `save_directory` to indicate where each tensor is saved. This helper uses + [`split_torch_state_dict_into_shards`] under the hood. If `safe_serialization` is `True`, the shards are saved as + safetensors (the default). Otherwise, the shards are saved as pickle. + + Before saving the model, the `save_directory` is cleaned from any previous shard files. + + + + If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a + size greater than `max_shard_size`. + + + + + + If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving. + + + + Args: + model (`torch.nn.Module`): + The model to save on disk. + save_directory (`str` or `Path`): + The directory in which the model will be saved. + filename_pattern (`str`, *optional*): + The pattern to generate the files names in which the model will be saved. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"model{suffix}.safetensors"` or `pytorch_model{suffix}.bin` depending on `safe_serialization` + parameter. + force_contiguous (`boolean`, *optional*): + Forcing the state_dict to be saved as contiguous tensors. This has no effect on the correctness of the + model, but it could potentially change performance if the layout of the tensor was chosen specifically for + that reason. Defaults to `True`. + max_shard_size (`int` or `str`, *optional*): + The maximum size of each shard, in bytes. Defaults to 5GB. + metadata (`Dict[str, str]`, *optional*): + Extra information to save along with the model. Some metadata will be added for each dropped tensors. + This information will not be enough to recover the entire shared structure but might help understanding + things. + safe_serialization (`bool`, *optional*): + Whether to save as safetensors, which is the default behavior. If `False`, the shards are saved as pickle. + Safe serialization is recommended for security reasons. Saving as pickle is deprecated and will be removed + in a future version. + is_main_process (`bool`, *optional*): + Whether the process calling this is the main process or not. Useful when in distributed training like + TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on + the main process to avoid race conditions. Defaults to True. + shared_tensors_to_discard (`List[str]`, *optional*): + List of tensor names to drop when saving shared tensors. If not provided and shared tensors are + detected, it will drop the first name alphabetically. + + Example: + + ```py + >>> from huggingface_hub import save_torch_model + >>> model = ... # A PyTorch model + + # Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors. + >>> save_torch_model(model, "path/to/folder") + + # Load model back + >>> from huggingface_hub import load_torch_model # TODO + >>> load_torch_model(model, "path/to/folder") + >>> + ``` + """ + save_torch_state_dict( + state_dict=model.state_dict(), + filename_pattern=filename_pattern, + force_contiguous=force_contiguous, + max_shard_size=max_shard_size, + metadata=metadata, + safe_serialization=safe_serialization, + save_directory=save_directory, + is_main_process=is_main_process, + shared_tensors_to_discard=shared_tensors_to_discard, + ) + + +def save_torch_state_dict( + state_dict: Dict[str, "torch.Tensor"], + save_directory: Union[str, Path], + *, + filename_pattern: Optional[str] = None, + force_contiguous: bool = True, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, + metadata: Optional[Dict[str, str]] = None, + safe_serialization: bool = True, + is_main_process: bool = True, + shared_tensors_to_discard: Optional[List[str]] = None, +) -> None: + """ + Save a model state dictionary to the disk, handling sharding and shared tensors issues. + + See also [`save_torch_model`] to directly save a PyTorch model. + + For more information about tensor sharing, check out [this guide](https://huggingface.co/docs/safetensors/torch_shared_tensors). + + The model state dictionary is split into shards so that each shard is smaller than a given size. The shards are + saved in the `save_directory` with the given `filename_pattern`. If the model is too big to fit in a single shard, + an index file is saved in the `save_directory` to indicate where each tensor is saved. This helper uses + [`split_torch_state_dict_into_shards`] under the hood. If `safe_serialization` is `True`, the shards are saved as + safetensors (the default). Otherwise, the shards are saved as pickle. + + Before saving the model, the `save_directory` is cleaned from any previous shard files. + + + + If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a + size greater than `max_shard_size`. + + + + + + If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving. + + + + Args: + state_dict (`Dict[str, torch.Tensor]`): + The state dictionary to save. + save_directory (`str` or `Path`): + The directory in which the model will be saved. + filename_pattern (`str`, *optional*): + The pattern to generate the files names in which the model will be saved. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"model{suffix}.safetensors"` or `pytorch_model{suffix}.bin` depending on `safe_serialization` + parameter. + force_contiguous (`boolean`, *optional*): + Forcing the state_dict to be saved as contiguous tensors. This has no effect on the correctness of the + model, but it could potentially change performance if the layout of the tensor was chosen specifically for + that reason. Defaults to `True`. + max_shard_size (`int` or `str`, *optional*): + The maximum size of each shard, in bytes. Defaults to 5GB. + metadata (`Dict[str, str]`, *optional*): + Extra information to save along with the model. Some metadata will be added for each dropped tensors. + This information will not be enough to recover the entire shared structure but might help understanding + things. + safe_serialization (`bool`, *optional*): + Whether to save as safetensors, which is the default behavior. If `False`, the shards are saved as pickle. + Safe serialization is recommended for security reasons. Saving as pickle is deprecated and will be removed + in a future version. + is_main_process (`bool`, *optional*): + Whether the process calling this is the main process or not. Useful when in distributed training like + TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on + the main process to avoid race conditions. Defaults to True. + shared_tensors_to_discard (`List[str]`, *optional*): + List of tensor names to drop when saving shared tensors. If not provided and shared tensors are + detected, it will drop the first name alphabetically. + + Example: + + ```py + >>> from huggingface_hub import save_torch_state_dict + >>> model = ... # A PyTorch model + + # Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors. + >>> state_dict = model_to_save.state_dict() + >>> save_torch_state_dict(state_dict, "path/to/folder") + ``` + """ + save_directory = str(save_directory) + + if filename_pattern is None: + filename_pattern = ( + constants.SAFETENSORS_WEIGHTS_FILE_PATTERN + if safe_serialization + else constants.PYTORCH_WEIGHTS_FILE_PATTERN + ) + + if metadata is None: + metadata = {} + if safe_serialization: + try: + from safetensors.torch import save_file as save_file_fn + except ImportError as e: + raise ImportError( + "Please install `safetensors` to use safe serialization. " + "You can install it with `pip install safetensors`." + ) from e + # Clean state dict for safetensors + state_dict = _clean_state_dict_for_safetensors( + state_dict, + metadata, + force_contiguous=force_contiguous, + shared_tensors_to_discard=shared_tensors_to_discard, + ) + else: + from torch import save as save_file_fn # type: ignore[assignment] + + logger.warning( + "You are using unsafe serialization. Due to security reasons, it is recommended not to load " + "pickled models from untrusted sources. If you intend to share your model, we strongly recommend " + "using safe serialization by installing `safetensors` with `pip install safetensors`." + ) + # Split dict + state_dict_split = split_torch_state_dict_into_shards( + state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + ) + + # Only main process should clean up existing files to avoid race conditions in distributed environment + if is_main_process: + existing_files_regex = re.compile(filename_pattern.format(suffix=r"(-\d{5}-of-\d{5})?") + r"(\.index\.json)?") + for filename in os.listdir(save_directory): + if existing_files_regex.match(filename): + try: + logger.debug(f"Removing existing file '{filename}' from folder.") + os.remove(os.path.join(save_directory, filename)) + except Exception as e: + logger.warning( + f"Error when trying to remove existing '{filename}' from folder: {e}. Continuing..." + ) + + # Save each shard + per_file_metadata = {"format": "pt"} + if not state_dict_split.is_sharded: + per_file_metadata.update(metadata) + safe_file_kwargs = {"metadata": per_file_metadata} if safe_serialization else {} + for filename, tensors in state_dict_split.filename_to_tensors.items(): + shard = {tensor: state_dict[tensor] for tensor in tensors} + save_file_fn(shard, os.path.join(save_directory, filename), **safe_file_kwargs) + logger.debug(f"Shard saved to {filename}") + + # Save the index (if any) + if state_dict_split.is_sharded: + index_path = filename_pattern.format(suffix="") + ".index.json" + index = { + "metadata": {**state_dict_split.metadata, **metadata}, + "weight_map": state_dict_split.tensor_to_filename, + } + with open(os.path.join(save_directory, index_path), "w") as f: + json.dump(index, f, indent=2) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}). " + f"Model weighs have been saved in {len(state_dict_split.filename_to_tensors)} checkpoint shards. " + f"You can find where each parameters has been saved in the index located at {index_path}." + ) + + logger.info(f"Model weights successfully saved to {save_directory}!") + + +def split_torch_state_dict_into_shards( + state_dict: Dict[str, "torch.Tensor"], + *, + filename_pattern: str = constants.SAFETENSORS_WEIGHTS_FILE_PATTERN, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, +) -> StateDictSplit: + """ + Split a model state dictionary in shards so that each shard is smaller than a given size. + + The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization + made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we + have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not + [6+2+2GB], [6+2GB], [6GB]. + + + + + To save a model state dictionary to the disk, see [`save_torch_state_dict`]. This helper uses + `split_torch_state_dict_into_shards` under the hood. + + + + + + If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a + size greater than `max_shard_size`. + + + + Args: + state_dict (`Dict[str, torch.Tensor]`): + The state dictionary to save. + filename_pattern (`str`, *optional*): + The pattern to generate the files names in which the model will be saved. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"model{suffix}.safetensors"`. + max_shard_size (`int` or `str`, *optional*): + The maximum size of each shard, in bytes. Defaults to 5GB. + + Returns: + [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them. + + Example: + ```py + >>> import json + >>> import os + >>> from safetensors.torch import save_file as safe_save_file + >>> from huggingface_hub import split_torch_state_dict_into_shards + + >>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str): + ... state_dict_split = split_torch_state_dict_into_shards(state_dict) + ... for filename, tensors in state_dict_split.filename_to_tensors.items(): + ... shard = {tensor: state_dict[tensor] for tensor in tensors} + ... safe_save_file( + ... shard, + ... os.path.join(save_directory, filename), + ... metadata={"format": "pt"}, + ... ) + ... if state_dict_split.is_sharded: + ... index = { + ... "metadata": state_dict_split.metadata, + ... "weight_map": state_dict_split.tensor_to_filename, + ... } + ... with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f: + ... f.write(json.dumps(index, indent=2)) + ``` + """ + return split_state_dict_into_shards_factory( + state_dict, + max_shard_size=max_shard_size, + filename_pattern=filename_pattern, + get_storage_size=get_torch_storage_size, + get_storage_id=get_torch_storage_id, + ) + + +# LOADING + + +def load_torch_model( + model: "torch.nn.Module", + checkpoint_path: Union[str, os.PathLike], + *, + strict: bool = False, + safe: bool = True, + weights_only: bool = False, + map_location: Optional[Union[str, "torch.device"]] = None, + mmap: bool = False, + filename_pattern: Optional[str] = None, +) -> NamedTuple: + """ + Load a checkpoint into a model, handling both sharded and non-sharded checkpoints. + + Args: + model (`torch.nn.Module`): + The model in which to load the checkpoint. + checkpoint_path (`str` or `os.PathLike`): + Path to either the checkpoint file or directory containing the checkpoint(s). + strict (`bool`, *optional*, defaults to `False`): + Whether to strictly enforce that the keys in the model state dict match the keys in the checkpoint. + safe (`bool`, *optional*, defaults to `True`): + If `safe` is True, the safetensors files will be loaded. If `safe` is False, the function + will first attempt to load safetensors files if they are available, otherwise it will fall back to loading + pickle files. `filename_pattern` parameter takes precedence over `safe` parameter. + weights_only (`bool`, *optional*, defaults to `False`): + If True, only loads the model weights without optimizer states and other metadata. + Only supported in PyTorch >= 1.13. + map_location (`str` or `torch.device`, *optional*): + A `torch.device` object, string or a dict specifying how to remap storage locations. It + indicates the location where all tensors should be loaded. + mmap (`bool`, *optional*, defaults to `False`): + Whether to use memory-mapped file loading. Memory mapping can improve loading performance + for large models in PyTorch >= 2.1.0 with zipfile-based checkpoints. + filename_pattern (`str`, *optional*): + The pattern to look for the index file. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"model{suffix}.safetensors"`. + Returns: + `NamedTuple`: A named tuple with `missing_keys` and `unexpected_keys` fields. + - `missing_keys` is a list of str containing the missing keys, i.e. keys that are in the model but not in the checkpoint. + - `unexpected_keys` is a list of str containing the unexpected keys, i.e. keys that are in the checkpoint but not in the model. + + Raises: + [`FileNotFoundError`](https://docs.python.org/3/library/exceptions.html#FileNotFoundError) + If the checkpoint file or directory does not exist. + [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + If safetensors or torch is not installed when trying to load a .safetensors file or a PyTorch checkpoint respectively. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the checkpoint path is invalid or if the checkpoint format cannot be determined. + + Example: + ```python + >>> from huggingface_hub import load_torch_model + >>> model = ... # A PyTorch model + >>> load_torch_model(model, "path/to/checkpoint") + ``` + """ + checkpoint_path = Path(checkpoint_path) + + if not checkpoint_path.exists(): + raise ValueError(f"Checkpoint path {checkpoint_path} does not exist") + # 1. Check if checkpoint is a single file + if checkpoint_path.is_file(): + state_dict = load_state_dict_from_file( + checkpoint_file=checkpoint_path, + map_location=map_location, + weights_only=weights_only, + ) + return model.load_state_dict(state_dict, strict=strict) + + # 2. If not, checkpoint_path is a directory + if filename_pattern is None: + filename_pattern = constants.SAFETENSORS_WEIGHTS_FILE_PATTERN + index_path = checkpoint_path / (filename_pattern.format(suffix="") + ".index.json") + # Only fallback to pickle format if safetensors index is not found and safe is False. + if not index_path.is_file() and not safe: + filename_pattern = constants.PYTORCH_WEIGHTS_FILE_PATTERN + + index_path = checkpoint_path / (filename_pattern.format(suffix="") + ".index.json") + + if index_path.is_file(): + return _load_sharded_checkpoint( + model=model, + save_directory=checkpoint_path, + strict=strict, + weights_only=weights_only, + filename_pattern=filename_pattern, + ) + + # Look for single model file + model_files = list(checkpoint_path.glob("*.safetensors" if safe else "*.bin")) + if len(model_files) == 1: + state_dict = load_state_dict_from_file( + checkpoint_file=model_files[0], + map_location=map_location, + weights_only=weights_only, + mmap=mmap, + ) + return model.load_state_dict(state_dict, strict=strict) + + raise ValueError( + f"Directory '{checkpoint_path}' does not contain a valid checkpoint. " + "Expected either a sharded checkpoint with an index file, or a single model file." + ) + + +def _load_sharded_checkpoint( + model: "torch.nn.Module", + save_directory: os.PathLike, + *, + strict: bool = False, + weights_only: bool = False, + filename_pattern: str = constants.SAFETENSORS_WEIGHTS_FILE_PATTERN, +) -> NamedTuple: + """ + Loads a sharded checkpoint into a model. This is the same as + [`torch.nn.Module.load_state_dict`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict) + but for a sharded checkpoint. Each shard is loaded one by one and removed from memory after being loaded into the model. + + Args: + model (`torch.nn.Module`): + The model in which to load the checkpoint. + save_directory (`str` or `os.PathLike`): + A path to a folder containing the sharded checkpoint. + strict (`bool`, *optional*, defaults to `False`): + Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint. + weights_only (`bool`, *optional*, defaults to `False`): + If True, only loads the model weights without optimizer states and other metadata. + Only supported in PyTorch >= 1.13. + filename_pattern (`str`, *optional*, defaults to `"model{suffix}.safetensors"`): + The pattern to look for the index file. Pattern must be a string that + can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix` + Defaults to `"model{suffix}.safetensors"`. + + Returns: + `NamedTuple`: A named tuple with `missing_keys` and `unexpected_keys` fields, + - `missing_keys` is a list of str containing the missing keys + - `unexpected_keys` is a list of str containing the unexpected keys + """ + + # 1. Load and validate index file + # The index file contains mapping of parameter names to shard files + index_path = filename_pattern.format(suffix="") + ".index.json" + index_file = os.path.join(save_directory, index_path) + with open(index_file, "r", encoding="utf-8") as f: + index = json.load(f) + + # 2. Validate keys if in strict mode + # This is done before loading any shards to fail fast + if strict: + _validate_keys_for_strict_loading(model, index["weight_map"].keys()) + + # 3. Load each shard using `load_state_dict` + # Get unique shard files (multiple parameters can be in same shard) + shard_files = list(set(index["weight_map"].values())) + for shard_file in shard_files: + # Load shard into memory + shard_path = os.path.join(save_directory, shard_file) + state_dict = load_state_dict_from_file( + shard_path, + map_location="cpu", + weights_only=weights_only, + ) + # Update model with parameters from this shard + model.load_state_dict(state_dict, strict=strict) + # Explicitly remove the state dict from memory + del state_dict + + # 4. Return compatibility info + loaded_keys = set(index["weight_map"].keys()) + model_keys = set(model.state_dict().keys()) + return _IncompatibleKeys( + missing_keys=list(model_keys - loaded_keys), unexpected_keys=list(loaded_keys - model_keys) + ) + + +def load_state_dict_from_file( + checkpoint_file: Union[str, os.PathLike], + map_location: Optional[Union[str, "torch.device"]] = None, + weights_only: bool = False, + mmap: bool = False, +) -> Union[Dict[str, "torch.Tensor"], Any]: + """ + Loads a checkpoint file, handling both safetensors and pickle checkpoint formats. + + Args: + checkpoint_file (`str` or `os.PathLike`): + Path to the checkpoint file to load. Can be either a safetensors or pickle (`.bin`) checkpoint. + map_location (`str` or `torch.device`, *optional*): + A `torch.device` object, string or a dict specifying how to remap storage locations. It + indicates the location where all tensors should be loaded. + weights_only (`bool`, *optional*, defaults to `False`): + If True, only loads the model weights without optimizer states and other metadata. + Only supported for pickle (`.bin`) checkpoints with PyTorch >= 1.13. Has no effect when + loading safetensors files. + mmap (`bool`, *optional*, defaults to `False`): + Whether to use memory-mapped file loading. Memory mapping can improve loading performance + for large models in PyTorch >= 2.1.0 with zipfile-based checkpoints. Has no effect when + loading safetensors files, as the `safetensors` library uses memory mapping by default. + + Returns: + `Union[Dict[str, "torch.Tensor"], Any]`: The loaded checkpoint. + - For safetensors files: always returns a dictionary mapping parameter names to tensors. + - For pickle files: returns any Python object that was pickled (commonly a state dict, but could be + an entire model, optimizer state, or any other Python object). + + Raises: + [`FileNotFoundError`](https://docs.python.org/3/library/exceptions.html#FileNotFoundError) + If the checkpoint file does not exist. + [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) + If safetensors or torch is not installed when trying to load a .safetensors file or a PyTorch checkpoint respectively. + [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) + If the checkpoint file format is invalid or if git-lfs files are not properly downloaded. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the checkpoint file path is empty or invalid. + + Example: + ```python + >>> from huggingface_hub import load_state_dict_from_file + + # Load a PyTorch checkpoint + >>> state_dict = load_state_dict_from_file("path/to/model.bin", map_location="cpu") + >>> model.load_state_dict(state_dict) + + # Load a safetensors checkpoint + >>> state_dict = load_state_dict_from_file("path/to/model.safetensors") + >>> model.load_state_dict(state_dict) + ``` + """ + checkpoint_path = Path(checkpoint_file) + + # Check if file exists and is a regular file (not a directory) + if not checkpoint_path.is_file(): + raise FileNotFoundError( + f"No checkpoint file found at '{checkpoint_path}'. Please verify the path is correct and " + "the file has been properly downloaded." + ) + + # Load safetensors checkpoint + if checkpoint_path.suffix == ".safetensors": + try: + from safetensors import safe_open + from safetensors.torch import load_file + except ImportError as e: + raise ImportError( + "Please install `safetensors` to load safetensors checkpoint. " + "You can install it with `pip install safetensors`." + ) from e + + # Check format of the archive + with safe_open(checkpoint_file, framework="pt") as f: # type: ignore[attr-defined] + metadata = f.metadata() + # see comment: https://github.com/huggingface/transformers/blob/3d213b57fe74302e5902d68ed9478c3ad1aaa713/src/transformers/modeling_utils.py#L3966 + if metadata is not None and metadata.get("format") not in ["pt", "mlx"]: + raise OSError( + f"The safetensors archive passed at {checkpoint_file} does not contain the valid metadata. Make sure " + "you save your model with the `save_torch_model` method." + ) + device = str(map_location.type) if map_location is not None and hasattr(map_location, "type") else map_location + # meta device is not supported with safetensors, falling back to CPU + if device == "meta": + logger.warning("Meta device is not supported with safetensors. Falling back to CPU device.") + device = "cpu" + return load_file(checkpoint_file, device=device) # type: ignore[arg-type] + # Otherwise, load from pickle + try: + import torch + from torch import load + except ImportError as e: + raise ImportError( + "Please install `torch` to load torch tensors. You can install it with `pip install torch`." + ) from e + # Add additional kwargs, mmap is only supported in torch >= 2.1.0 + additional_kwargs = {} + if version.parse(torch.__version__) >= version.parse("2.1.0"): + additional_kwargs["mmap"] = mmap + + # weights_only is only supported in torch >= 1.13.0 + if version.parse(torch.__version__) >= version.parse("1.13.0"): + additional_kwargs["weights_only"] = weights_only + + return load( + checkpoint_file, + map_location=map_location, + **additional_kwargs, + ) + + +# HELPERS + + +def _validate_keys_for_strict_loading( + model: "torch.nn.Module", + loaded_keys: Iterable[str], +) -> None: + """ + Validate that model keys match loaded keys when strict loading is enabled. + + Args: + model: The PyTorch model being loaded + loaded_keys: The keys present in the checkpoint + + Raises: + RuntimeError: If there are missing or unexpected keys in strict mode + """ + loaded_keys_set = set(loaded_keys) + model_keys = set(model.state_dict().keys()) + missing_keys = model_keys - loaded_keys_set # Keys in model but not in checkpoint + unexpected_keys = loaded_keys_set - model_keys # Keys in checkpoint but not in model + + if missing_keys or unexpected_keys: + error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}" + if missing_keys: + str_missing_keys = ",".join([f'"{k}"' for k in sorted(missing_keys)]) + error_message += f"\nMissing key(s): {str_missing_keys}." + if unexpected_keys: + str_unexpected_keys = ",".join([f'"{k}"' for k in sorted(unexpected_keys)]) + error_message += f"\nUnexpected key(s): {str_unexpected_keys}." + raise RuntimeError(error_message) + + +def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]: + """Returns a unique id for plain tensor + or a (potentially nested) Tuple of unique id for the flattened Tensor + if the input is a wrapper tensor subclass Tensor + """ + + try: + from torch.distributed.tensor import DTensor + + if isinstance(tensor, DTensor): + local_tensor = tensor.to_local() + return local_tensor.storage().data_ptr() + except ImportError: + pass + + try: + # for torch 2.1 and above we can also handle tensor subclasses + from torch.utils._python_dispatch import is_traceable_wrapper_subclass + + if is_traceable_wrapper_subclass(tensor): + attrs, _ = tensor.__tensor_flatten__() # type: ignore[attr-defined] + return tuple(_get_unique_id(getattr(tensor, attr)) for attr in attrs) + + except ImportError: + # for torch version less than 2.1, we can fallback to original implementation + pass + + if tensor.device.type == "xla" and is_torch_tpu_available(): + # NOTE: xla tensors dont have storage + # use some other unique id to distinguish. + # this is a XLA tensor, it must be created using torch_xla's + # device. So the following import is safe: + import torch_xla # type: ignore[import] + + unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor) + else: + unique_id = storage_ptr(tensor) + + return unique_id + + +def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[Tuple["torch.device", Union[int, Tuple[Any, ...]], int]]: + """ + Return unique identifier to a tensor storage. + + Multiple different tensors can share the same underlying storage. This identifier is + guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with + non-overlapping lifetimes may have the same id. + In the case of meta tensors, we return None since we can't tell if they share the same storage. + + Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278. + """ + if tensor.device.type == "meta": + return None + else: + return tensor.device, _get_unique_id(tensor), get_torch_storage_size(tensor) + + +def get_torch_storage_size(tensor: "torch.Tensor") -> int: + """ + Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L31C1-L41C59 + """ + try: + from torch.distributed.tensor import DTensor + + if isinstance(tensor, DTensor): + # this returns the size of the FULL tensor in bytes + return tensor.nbytes + except ImportError: + pass + + try: + # for torch 2.1 and above we can also handle tensor subclasses + from torch.utils._python_dispatch import is_traceable_wrapper_subclass + + if is_traceable_wrapper_subclass(tensor): + attrs, _ = tensor.__tensor_flatten__() # type: ignore[attr-defined] + return sum(get_torch_storage_size(getattr(tensor, attr)) for attr in attrs) + except ImportError: + # for torch version less than 2.1, we can fallback to original implementation + pass + + try: + return tensor.untyped_storage().nbytes() + except AttributeError: + # Fallback for torch==1.10 + try: + return tensor.storage().size() * _get_dtype_size(tensor.dtype) + except NotImplementedError: + # Fallback for meta storage + # On torch >=2.0 this is the tensor size + return tensor.nelement() * _get_dtype_size(tensor.dtype) + + +@lru_cache() +def is_torch_tpu_available(check_device=True): + """ + Checks if `torch_xla` is installed and potentially if a TPU is in the environment + + Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/utils/import_utils.py#L463. + """ + if importlib.util.find_spec("torch_xla") is not None: + if check_device: + # We need to check if `xla_device` can be found, will raise a RuntimeError if not + try: + import torch_xla.core.xla_model as xm # type: ignore[import] + + _ = xm.xla_device() + return True + except RuntimeError: + return False + return True + return False + + +def storage_ptr(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L11. + """ + try: + # for torch 2.1 and above we can also handle tensor subclasses + from torch.utils._python_dispatch import is_traceable_wrapper_subclass + + if is_traceable_wrapper_subclass(tensor): + return _get_unique_id(tensor) # type: ignore + except ImportError: + # for torch version less than 2.1, we can fallback to original implementation + pass + + try: + return tensor.untyped_storage().data_ptr() + except Exception: + # Fallback for torch==1.10 + try: + return tensor.storage().data_ptr() + except NotImplementedError: + # Fallback for meta storage + return 0 + + +def _clean_state_dict_for_safetensors( + state_dict: Dict[str, "torch.Tensor"], + metadata: Dict[str, str], + force_contiguous: bool = True, + shared_tensors_to_discard: Optional[List[str]] = None, +): + """Remove shared tensors from state_dict and update metadata accordingly (for reloading). + + Warning: `state_dict` and `metadata` are mutated in-place! + + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L155. + """ + to_removes = _remove_duplicate_names(state_dict, discard_names=shared_tensors_to_discard) + for kept_name, to_remove_group in to_removes.items(): + for to_remove in to_remove_group: + if metadata is None: + metadata = {} + + if to_remove not in metadata: + # Do not override user data + metadata[to_remove] = kept_name + del state_dict[to_remove] + if force_contiguous: + state_dict = {k: v.contiguous() for k, v in state_dict.items()} + return state_dict + + +def _end_ptr(tensor: "torch.Tensor") -> int: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L23. + """ + if tensor.nelement(): + stop = tensor.view(-1)[-1].data_ptr() + _get_dtype_size(tensor.dtype) + else: + stop = tensor.data_ptr() + return stop + + +def _filter_shared_not_shared(tensors: List[Set[str]], state_dict: Dict[str, "torch.Tensor"]) -> List[Set[str]]: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L44 + """ + filtered_tensors = [] + for shared in tensors: + if len(shared) < 2: + filtered_tensors.append(shared) + continue + + areas = [] + for name in shared: + tensor = state_dict[name] + areas.append((tensor.data_ptr(), _end_ptr(tensor), name)) + areas.sort() + + _, last_stop, last_name = areas[0] + filtered_tensors.append({last_name}) + for start, stop, name in areas[1:]: + if start >= last_stop: + filtered_tensors.append({name}) + else: + filtered_tensors[-1].add(name) + last_stop = stop + + return filtered_tensors + + +def _find_shared_tensors(state_dict: Dict[str, "torch.Tensor"]) -> List[Set[str]]: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L69. + """ + import torch + + tensors_dict = defaultdict(set) + for k, v in state_dict.items(): + if v.device != torch.device("meta") and storage_ptr(v) != 0 and get_torch_storage_size(v) != 0: + # Need to add device as key because of multiple GPU. + tensors_dict[(v.device, storage_ptr(v), get_torch_storage_size(v))].add(k) + tensors = list(sorted(tensors_dict.values())) + tensors = _filter_shared_not_shared(tensors, state_dict) + return tensors + + +def _is_complete(tensor: "torch.Tensor") -> bool: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L80 + """ + try: + # for torch 2.1 and above we can also handle tensor subclasses + from torch.utils._python_dispatch import is_traceable_wrapper_subclass + + if is_traceable_wrapper_subclass(tensor): + attrs, _ = tensor.__tensor_flatten__() # type: ignore[attr-defined] + return all(_is_complete(getattr(tensor, attr)) for attr in attrs) + except ImportError: + # for torch version less than 2.1, we can fallback to original implementation + pass + + return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _get_dtype_size( + tensor.dtype + ) == get_torch_storage_size(tensor) + + +def _remove_duplicate_names( + state_dict: Dict[str, "torch.Tensor"], + *, + preferred_names: Optional[List[str]] = None, + discard_names: Optional[List[str]] = None, +) -> Dict[str, List[str]]: + """ + Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L80 + """ + if preferred_names is None: + preferred_names = [] + unique_preferred_names = set(preferred_names) + if discard_names is None: + discard_names = [] + unique_discard_names = set(discard_names) + + shareds = _find_shared_tensors(state_dict) + to_remove = defaultdict(list) + for shared in shareds: + complete_names = set([name for name in shared if _is_complete(state_dict[name])]) + if not complete_names: + raise RuntimeError( + "Error while trying to find names to remove to save state dict, but found no suitable name to keep" + f" for saving amongst: {shared}. None is covering the entire storage. Refusing to save/load the model" + " since you could be storing much more memory than needed. Please refer to" + " https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an" + " issue." + ) + + keep_name = sorted(list(complete_names))[0] + + # Mechanism to preferentially select keys to keep + # coming from the on-disk file to allow + # loading models saved with a different choice + # of keep_name + preferred = complete_names.difference(unique_discard_names) + if preferred: + keep_name = sorted(list(preferred))[0] + + if unique_preferred_names: + preferred = unique_preferred_names.intersection(complete_names) + if preferred: + keep_name = sorted(list(preferred))[0] + for name in sorted(shared): + if name != keep_name: + to_remove[keep_name].append(name) + return to_remove + + +@lru_cache() +def _get_dtype_size(dtype: "torch.dtype") -> int: + """ + Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L344 + """ + import torch + + # torch.float8 formats require 2.1; we do not support these dtypes on earlier versions + _float8_e4m3fn = getattr(torch, "float8_e4m3fn", None) + _float8_e5m2 = getattr(torch, "float8_e5m2", None) + _SIZE = { + torch.int64: 8, + torch.float32: 4, + torch.int32: 4, + torch.bfloat16: 2, + torch.float16: 2, + torch.int16: 2, + torch.uint8: 1, + torch.int8: 1, + torch.bool: 1, + torch.float64: 8, + _float8_e4m3fn: 1, + _float8_e5m2: 1, + } + return _SIZE[dtype] + + +class _IncompatibleKeys(namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])): + """ + This is used to report missing and unexpected keys in the state dict. + Taken from https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py#L52. + + """ + + def __repr__(self) -> str: + if not self.missing_keys and not self.unexpected_keys: + return "" + return super().__repr__() + + __str__ = __repr__ diff --git a/lib/python3.12/site-packages/huggingface_hub/templates/datasetcard_template.md b/lib/python3.12/site-packages/huggingface_hub/templates/datasetcard_template.md new file mode 100644 index 0000000000000000000000000000000000000000..9af29ebbed93653ec74a8952e314e7554323ef15 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/templates/datasetcard_template.md @@ -0,0 +1,143 @@ +--- +# For reference on dataset card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/datasetcard.md?plain=1 +# Doc / guide: https://huggingface.co/docs/hub/datasets-cards +{{ card_data }} +--- + +# Dataset Card for {{ pretty_name | default("Dataset Name", true) }} + + + +{{ dataset_summary | default("", true) }} + +## Dataset Details + +### Dataset Description + + + +{{ dataset_description | default("", true) }} + +- **Curated by:** {{ curators | default("[More Information Needed]", true)}} +- **Funded by [optional]:** {{ funded_by | default("[More Information Needed]", true)}} +- **Shared by [optional]:** {{ shared_by | default("[More Information Needed]", true)}} +- **Language(s) (NLP):** {{ language | default("[More Information Needed]", true)}} +- **License:** {{ license | default("[More Information Needed]", true)}} + +### Dataset Sources [optional] + + + +- **Repository:** {{ repo | default("[More Information Needed]", true)}} +- **Paper [optional]:** {{ paper | default("[More Information Needed]", true)}} +- **Demo [optional]:** {{ demo | default("[More Information Needed]", true)}} + +## Uses + + + +### Direct Use + + + +{{ direct_use | default("[More Information Needed]", true)}} + +### Out-of-Scope Use + + + +{{ out_of_scope_use | default("[More Information Needed]", true)}} + +## Dataset Structure + + + +{{ dataset_structure | default("[More Information Needed]", true)}} + +## Dataset Creation + +### Curation Rationale + + + +{{ curation_rationale_section | default("[More Information Needed]", true)}} + +### Source Data + + + +#### Data Collection and Processing + + + +{{ data_collection_and_processing_section | default("[More Information Needed]", true)}} + +#### Who are the source data producers? + + + +{{ source_data_producers_section | default("[More Information Needed]", true)}} + +### Annotations [optional] + + + +#### Annotation process + + + +{{ annotation_process_section | default("[More Information Needed]", true)}} + +#### Who are the annotators? + + + +{{ who_are_annotators_section | default("[More Information Needed]", true)}} + +#### Personal and Sensitive Information + + + +{{ personal_and_sensitive_information | default("[More Information Needed]", true)}} + +## Bias, Risks, and Limitations + + + +{{ bias_risks_limitations | default("[More Information Needed]", true)}} + +### Recommendations + + + +{{ bias_recommendations | default("Users should be made aware of the risks, biases and limitations of the dataset. More information needed for further recommendations.", true)}} + +## Citation [optional] + + + +**BibTeX:** + +{{ citation_bibtex | default("[More Information Needed]", true)}} + +**APA:** + +{{ citation_apa | default("[More Information Needed]", true)}} + +## Glossary [optional] + + + +{{ glossary | default("[More Information Needed]", true)}} + +## More Information [optional] + +{{ more_information | default("[More Information Needed]", true)}} + +## Dataset Card Authors [optional] + +{{ dataset_card_authors | default("[More Information Needed]", true)}} + +## Dataset Card Contact + +{{ dataset_card_contact | default("[More Information Needed]", true)}} diff --git a/lib/python3.12/site-packages/huggingface_hub/templates/modelcard_template.md b/lib/python3.12/site-packages/huggingface_hub/templates/modelcard_template.md new file mode 100644 index 0000000000000000000000000000000000000000..79ca15e4547debac763b390ef8e4b715e6f6403f --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/templates/modelcard_template.md @@ -0,0 +1,200 @@ +--- +# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1 +# Doc / guide: https://huggingface.co/docs/hub/model-cards +{{ card_data }} +--- + +# Model Card for {{ model_id | default("Model ID", true) }} + + + +{{ model_summary | default("", true) }} + +## Model Details + +### Model Description + + + +{{ model_description | default("", true) }} + +- **Developed by:** {{ developers | default("[More Information Needed]", true)}} +- **Funded by [optional]:** {{ funded_by | default("[More Information Needed]", true)}} +- **Shared by [optional]:** {{ shared_by | default("[More Information Needed]", true)}} +- **Model type:** {{ model_type | default("[More Information Needed]", true)}} +- **Language(s) (NLP):** {{ language | default("[More Information Needed]", true)}} +- **License:** {{ license | default("[More Information Needed]", true)}} +- **Finetuned from model [optional]:** {{ base_model | default("[More Information Needed]", true)}} + +### Model Sources [optional] + + + +- **Repository:** {{ repo | default("[More Information Needed]", true)}} +- **Paper [optional]:** {{ paper | default("[More Information Needed]", true)}} +- **Demo [optional]:** {{ demo | default("[More Information Needed]", true)}} + +## Uses + + + +### Direct Use + + + +{{ direct_use | default("[More Information Needed]", true)}} + +### Downstream Use [optional] + + + +{{ downstream_use | default("[More Information Needed]", true)}} + +### Out-of-Scope Use + + + +{{ out_of_scope_use | default("[More Information Needed]", true)}} + +## Bias, Risks, and Limitations + + + +{{ bias_risks_limitations | default("[More Information Needed]", true)}} + +### Recommendations + + + +{{ bias_recommendations | default("Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.", true)}} + +## How to Get Started with the Model + +Use the code below to get started with the model. + +{{ get_started_code | default("[More Information Needed]", true)}} + +## Training Details + +### Training Data + + + +{{ training_data | default("[More Information Needed]", true)}} + +### Training Procedure + + + +#### Preprocessing [optional] + +{{ preprocessing | default("[More Information Needed]", true)}} + + +#### Training Hyperparameters + +- **Training regime:** {{ training_regime | default("[More Information Needed]", true)}} + +#### Speeds, Sizes, Times [optional] + + + +{{ speeds_sizes_times | default("[More Information Needed]", true)}} + +## Evaluation + + + +### Testing Data, Factors & Metrics + +#### Testing Data + + + +{{ testing_data | default("[More Information Needed]", true)}} + +#### Factors + + + +{{ testing_factors | default("[More Information Needed]", true)}} + +#### Metrics + + + +{{ testing_metrics | default("[More Information Needed]", true)}} + +### Results + +{{ results | default("[More Information Needed]", true)}} + +#### Summary + +{{ results_summary | default("", true) }} + +## Model Examination [optional] + + + +{{ model_examination | default("[More Information Needed]", true)}} + +## Environmental Impact + + + +Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700). + +- **Hardware Type:** {{ hardware_type | default("[More Information Needed]", true)}} +- **Hours used:** {{ hours_used | default("[More Information Needed]", true)}} +- **Cloud Provider:** {{ cloud_provider | default("[More Information Needed]", true)}} +- **Compute Region:** {{ cloud_region | default("[More Information Needed]", true)}} +- **Carbon Emitted:** {{ co2_emitted | default("[More Information Needed]", true)}} + +## Technical Specifications [optional] + +### Model Architecture and Objective + +{{ model_specs | default("[More Information Needed]", true)}} + +### Compute Infrastructure + +{{ compute_infrastructure | default("[More Information Needed]", true)}} + +#### Hardware + +{{ hardware_requirements | default("[More Information Needed]", true)}} + +#### Software + +{{ software | default("[More Information Needed]", true)}} + +## Citation [optional] + + + +**BibTeX:** + +{{ citation_bibtex | default("[More Information Needed]", true)}} + +**APA:** + +{{ citation_apa | default("[More Information Needed]", true)}} + +## Glossary [optional] + + + +{{ glossary | default("[More Information Needed]", true)}} + +## More Information [optional] + +{{ more_information | default("[More Information Needed]", true)}} + +## Model Card Authors [optional] + +{{ model_card_authors | default("[More Information Needed]", true)}} + +## Model Card Contact + +{{ model_card_contact | default("[More Information Needed]", true)}} diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__init__.py b/lib/python3.12/site-packages/huggingface_hub/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..992eac104bd80de97444003172e926d5ad4522a0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/__init__.py @@ -0,0 +1,117 @@ +# coding=utf-8 +# Copyright 2021 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 + +# ruff: noqa: F401 +from huggingface_hub.errors import ( + BadRequestError, + CacheNotFound, + CorruptedCacheException, + DisabledRepoError, + EntryNotFoundError, + FileMetadataError, + GatedRepoError, + HfHubHTTPError, + HFValidationError, + LocalEntryNotFoundError, + LocalTokenNotFoundError, + NotASafetensorsRepoError, + OfflineModeIsEnabled, + RepositoryNotFoundError, + RevisionNotFoundError, + SafetensorsParsingError, +) + +from . import tqdm as _tqdm # _tqdm is the module +from ._auth import get_stored_tokens, get_token +from ._cache_assets import cached_assets_path +from ._cache_manager import ( + CachedFileInfo, + CachedRepoInfo, + CachedRevisionInfo, + DeleteCacheStrategy, + HFCacheInfo, + scan_cache_dir, +) +from ._chunk_utils import chunk_iterable +from ._datetime import parse_datetime +from ._experimental import experimental +from ._fixes import SoftTemporaryDirectory, WeakFileLock, yaml_dump +from ._git_credential import list_credential_helpers, set_git_credential, unset_git_credential +from ._headers import build_hf_headers, get_token_to_send +from ._hf_folder import HfFolder +from ._http import ( + configure_http_backend, + fix_hf_endpoint_in_url, + get_session, + hf_raise_for_status, + http_backoff, + reset_sessions, +) +from ._pagination import paginate +from ._paths import DEFAULT_IGNORE_PATTERNS, FORBIDDEN_FOLDERS, filter_repo_objects +from ._runtime import ( + dump_environment_info, + get_aiohttp_version, + get_fastai_version, + get_fastapi_version, + get_fastcore_version, + get_gradio_version, + get_graphviz_version, + get_hf_hub_version, + get_hf_transfer_version, + get_jinja_version, + get_numpy_version, + get_pillow_version, + get_pydantic_version, + get_pydot_version, + get_python_version, + get_tensorboard_version, + get_tf_version, + get_torch_version, + is_aiohttp_available, + is_colab_enterprise, + is_fastai_available, + is_fastapi_available, + is_fastcore_available, + is_google_colab, + is_gradio_available, + is_graphviz_available, + is_hf_transfer_available, + is_jinja_available, + is_notebook, + is_numpy_available, + is_package_available, + is_pillow_available, + is_pydantic_available, + is_pydot_available, + is_safetensors_available, + is_tensorboard_available, + is_tf_available, + is_torch_available, +) +from ._safetensors import SafetensorsFileMetadata, SafetensorsRepoMetadata, TensorInfo +from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess +from ._telemetry import send_telemetry +from ._typing import is_jsonable, is_simple_optional_type, unwrap_simple_optional_type +from ._validators import smoothly_deprecate_use_auth_token, validate_hf_hub_args, validate_repo_id +from ._xet import ( + XetConnectionInfo, + XetFileData, + XetTokenType, + fetch_xet_connection_info_from_repo_info, + parse_xet_file_data_from_response, + refresh_xet_connection_info, +) +from .tqdm import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm, tqdm_stream_file diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6ff449702789434fad02da94206ad2455543c8f Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_auth.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_auth.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c788e1a8c882aea55bbf46af03d3951fa4a2a391 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_auth.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_assets.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_assets.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1464f1523d5e3303c534662cf6666efc5276736 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_assets.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_manager.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbd5d4c351206cfba2fbae2cf0f3afa178fc9068 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_cache_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_chunk_utils.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_chunk_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9db11522ccecc17be87c8799b68221fb231dcd5e Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_chunk_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_datetime.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_datetime.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08d17b6947057f505e8188efdf7ba46d514750e5 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_datetime.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_deprecation.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_deprecation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f98fadbcd75bc3b0b89f0c578fb74c91d90adc3 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_deprecation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_experimental.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_experimental.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf6300c3a5d9351f8cc6c40848eb480572ac09dc Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_experimental.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_fixes.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_fixes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73c5a25d791580867c07aa280850b67d8ac08995 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_fixes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_git_credential.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_git_credential.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4966efc621a9dd721b8d4e92d2563999e48731bf Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_git_credential.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_headers.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_headers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a5a6970f3692b49d037da6be5e68cd367d27bcb Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_headers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_hf_folder.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_hf_folder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0cd960a93b661ebc5eff5747a421c22c0139e4b6 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_hf_folder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_http.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_http.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d196e4a2ecb694792cf7d0177584bca27c16163 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_http.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_lfs.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_lfs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eda5d5ea8b268af2eeb5b530111812a40463c0f1 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_lfs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_pagination.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_pagination.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0931bbeeaa91d3be26c5193956e40ac89dd74a5f Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_pagination.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_paths.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_paths.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c358d747c6bc83fd7784b7a628e1488e5edda53b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_paths.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_runtime.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_runtime.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eefbc603c644f32455e04c07e9017598c7cd614b Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_runtime.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_safetensors.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_safetensors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f40ec7435b15e447d9eaa7e2aaa509ce06883245 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_safetensors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_subprocess.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_subprocess.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eea2a73c5c50a5adbe71422c682ab907e71309b6 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_subprocess.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_telemetry.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_telemetry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90a8a9c4cef9dafc085e41de03803cf45181008 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_telemetry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_typing.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_typing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9ce0c8f88e6b40d3a47f5ffb9d0ccd46c056cbd Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_typing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_validators.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddb2b05ff3c10b9273901a505c101166faa38fd9 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_xet.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_xet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ec12bd64c3eba4170b2af67245a6edd1dab93bd Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/_xet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/endpoint_helpers.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/endpoint_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..daf6478d9b683c7cce8ac219a362e6dd51e7bb99 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/endpoint_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/insecure_hashlib.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/insecure_hashlib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..870625e65382252584d3dfe2953d6b2f88053097 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/insecure_hashlib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f50478771316fff997b170b03bf333fabd5f82f4 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/sha.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/sha.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5559dfc07dbf91f8157c605e0c734d16b94d54b0 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/sha.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/tqdm.cpython-312.pyc b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/tqdm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1922bb5c0c7f85b54b2cdd69410ff50fa7275649 Binary files /dev/null and b/lib/python3.12/site-packages/huggingface_hub/utils/__pycache__/tqdm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_auth.py b/lib/python3.12/site-packages/huggingface_hub/utils/_auth.py new file mode 100644 index 0000000000000000000000000000000000000000..c70280aec4e6ca1ca11f6452d1bf9e57e9953834 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_auth.py @@ -0,0 +1,214 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains an helper to get the token from machine (env variable, secret or config file).""" + +import configparser +import logging +import os +import warnings +from pathlib import Path +from threading import Lock +from typing import Dict, Optional + +from .. import constants +from ._runtime import is_colab_enterprise, is_google_colab + + +_IS_GOOGLE_COLAB_CHECKED = False +_GOOGLE_COLAB_SECRET_LOCK = Lock() +_GOOGLE_COLAB_SECRET: Optional[str] = None + +logger = logging.getLogger(__name__) + + +def get_token() -> Optional[str]: + """ + Get token if user is logged in. + + Note: in most cases, you should use [`huggingface_hub.utils.build_hf_headers`] instead. This method is only useful + if you want to retrieve the token for other purposes than sending an HTTP request. + + Token is retrieved in priority from the `HF_TOKEN` environment variable. Otherwise, we read the token file located + in the Hugging Face home folder. Returns None if user is not logged in. To log in, use [`login`] or + `huggingface-cli login`. + + Returns: + `str` or `None`: The token, `None` if it doesn't exist. + """ + return _get_token_from_google_colab() or _get_token_from_environment() or _get_token_from_file() + + +def _get_token_from_google_colab() -> Optional[str]: + """Get token from Google Colab secrets vault using `google.colab.userdata.get(...)`. + + Token is read from the vault only once per session and then stored in a global variable to avoid re-requesting + access to the vault. + """ + # If it's not a Google Colab or it's Colab Enterprise, fallback to environment variable or token file authentication + if not is_google_colab() or is_colab_enterprise(): + return None + + # `google.colab.userdata` is not thread-safe + # This can lead to a deadlock if multiple threads try to access it at the same time + # (typically when using `snapshot_download`) + # => use a lock + # See https://github.com/huggingface/huggingface_hub/issues/1952 for more details. + with _GOOGLE_COLAB_SECRET_LOCK: + global _GOOGLE_COLAB_SECRET + global _IS_GOOGLE_COLAB_CHECKED + + if _IS_GOOGLE_COLAB_CHECKED: # request access only once + return _GOOGLE_COLAB_SECRET + + try: + from google.colab import userdata # type: ignore + from google.colab.errors import Error as ColabError # type: ignore + except ImportError: + return None + + try: + token = userdata.get("HF_TOKEN") + _GOOGLE_COLAB_SECRET = _clean_token(token) + except userdata.NotebookAccessError: + # Means the user has a secret call `HF_TOKEN` and got a popup "please grand access to HF_TOKEN" and refused it + # => warn user but ignore error => do not re-request access to user + warnings.warn( + "\nAccess to the secret `HF_TOKEN` has not been granted on this notebook." + "\nYou will not be requested again." + "\nPlease restart the session if you want to be prompted again." + ) + _GOOGLE_COLAB_SECRET = None + except userdata.SecretNotFoundError: + # Means the user did not define a `HF_TOKEN` secret => warn + warnings.warn( + "\nThe secret `HF_TOKEN` does not exist in your Colab secrets." + "\nTo authenticate with the Hugging Face Hub, create a token in your settings tab " + "(https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session." + "\nYou will be able to reuse this secret in all of your notebooks." + "\nPlease note that authentication is recommended but still optional to access public models or datasets." + ) + _GOOGLE_COLAB_SECRET = None + except ColabError as e: + # Something happen but we don't know what => recommend to open a GitHub issue + warnings.warn( + f"\nError while fetching `HF_TOKEN` secret value from your vault: '{str(e)}'." + "\nYou are not authenticated with the Hugging Face Hub in this notebook." + "\nIf the error persists, please let us know by opening an issue on GitHub " + "(https://github.com/huggingface/huggingface_hub/issues/new)." + ) + _GOOGLE_COLAB_SECRET = None + + _IS_GOOGLE_COLAB_CHECKED = True + return _GOOGLE_COLAB_SECRET + + +def _get_token_from_environment() -> Optional[str]: + # `HF_TOKEN` has priority (keep `HUGGING_FACE_HUB_TOKEN` for backward compatibility) + return _clean_token(os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")) + + +def _get_token_from_file() -> Optional[str]: + try: + return _clean_token(Path(constants.HF_TOKEN_PATH).read_text()) + except FileNotFoundError: + return None + + +def get_stored_tokens() -> Dict[str, str]: + """ + Returns the parsed INI file containing the access tokens. + The file is located at `HF_STORED_TOKENS_PATH`, defaulting to `~/.cache/huggingface/stored_tokens`. + If the file does not exist, an empty dictionary is returned. + + Returns: `Dict[str, str]` + Key is the token name and value is the token. + """ + tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + if not tokens_path.exists(): + stored_tokens = {} + config = configparser.ConfigParser() + try: + config.read(tokens_path) + stored_tokens = {token_name: config.get(token_name, "hf_token") for token_name in config.sections()} + except configparser.Error as e: + logger.error(f"Error parsing stored tokens file: {e}") + stored_tokens = {} + return stored_tokens + + +def _save_stored_tokens(stored_tokens: Dict[str, str]) -> None: + """ + Saves the given configuration to the stored tokens file. + + Args: + stored_tokens (`Dict[str, str]`): + The stored tokens to save. Key is the token name and value is the token. + """ + stored_tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + + # Write the stored tokens into an INI file + config = configparser.ConfigParser() + for token_name in sorted(stored_tokens.keys()): + config.add_section(token_name) + config.set(token_name, "hf_token", stored_tokens[token_name]) + + stored_tokens_path.parent.mkdir(parents=True, exist_ok=True) + with stored_tokens_path.open("w") as config_file: + config.write(config_file) + + +def _get_token_by_name(token_name: str) -> Optional[str]: + """ + Get the token by name. + + Args: + token_name (`str`): + The name of the token to get. + + Returns: + `str` or `None`: The token, `None` if it doesn't exist. + + """ + stored_tokens = get_stored_tokens() + if token_name not in stored_tokens: + return None + return _clean_token(stored_tokens[token_name]) + + +def _save_token(token: str, token_name: str) -> None: + """ + Save the given token. + + If the stored tokens file does not exist, it will be created. + Args: + token (`str`): + The token to save. + token_name (`str`): + The name of the token. + """ + tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + stored_tokens = get_stored_tokens() + stored_tokens[token_name] = token + _save_stored_tokens(stored_tokens) + logger.info(f"The token `{token_name}` has been saved to {tokens_path}") + + +def _clean_token(token: Optional[str]) -> Optional[str]: + """Clean token by removing trailing and leading spaces and newlines. + + If token is an empty string, return None. + """ + if token is None: + return None + return token.replace("\r", "").replace("\n", "").strip() or None diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_cache_assets.py b/lib/python3.12/site-packages/huggingface_hub/utils/_cache_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..e5d435df9b0bb0c67c0bcb5ef65711e9aef367f6 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_cache_assets.py @@ -0,0 +1,135 @@ +# coding=utf-8 +# Copyright 2019-present, the HuggingFace Inc. team. +# +# 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 pathlib import Path +from typing import Union + +from ..constants import HF_ASSETS_CACHE + + +def cached_assets_path( + library_name: str, + namespace: str = "default", + subfolder: str = "default", + *, + assets_dir: Union[str, Path, None] = None, +): + """Return a folder path to cache arbitrary files. + + `huggingface_hub` provides a canonical folder path to store assets. This is the + recommended way to integrate cache in a downstream library as it will benefit from + the builtins tools to scan and delete the cache properly. + + The distinction is made between files cached from the Hub and assets. Files from the + Hub are cached in a git-aware manner and entirely managed by `huggingface_hub`. See + [related documentation](https://huggingface.co/docs/huggingface_hub/how-to-cache). + All other files that a downstream library caches are considered to be "assets" + (files downloaded from external sources, extracted from a .tar archive, preprocessed + for training,...). + + Once the folder path is generated, it is guaranteed to exist and to be a directory. + The path is based on 3 levels of depth: the library name, a namespace and a + subfolder. Those 3 levels grants flexibility while allowing `huggingface_hub` to + expect folders when scanning/deleting parts of the assets cache. Within a library, + it is expected that all namespaces share the same subset of subfolder names but this + is not a mandatory rule. The downstream library has then full control on which file + structure to adopt within its cache. Namespace and subfolder are optional (would + default to a `"default/"` subfolder) but library name is mandatory as we want every + downstream library to manage its own cache. + + Expected tree: + ```text + assets/ + └── datasets/ + │ ├── SQuAD/ + │ │ ├── downloaded/ + │ │ ├── extracted/ + │ │ └── processed/ + │ ├── Helsinki-NLP--tatoeba_mt/ + │ ├── downloaded/ + │ ├── extracted/ + │ └── processed/ + └── transformers/ + ├── default/ + │ ├── something/ + ├── bert-base-cased/ + │ ├── default/ + │ └── training/ + hub/ + └── models--julien-c--EsperBERTo-small/ + ├── blobs/ + │ ├── (...) + │ ├── (...) + ├── refs/ + │ └── (...) + └── [ 128] snapshots/ + ├── 2439f60ef33a0d46d85da5001d52aeda5b00ce9f/ + │ ├── (...) + └── bbc77c8132af1cc5cf678da3f1ddf2de43606d48/ + └── (...) + ``` + + + Args: + library_name (`str`): + Name of the library that will manage the cache folder. Example: `"dataset"`. + namespace (`str`, *optional*, defaults to "default"): + Namespace to which the data belongs. Example: `"SQuAD"`. + subfolder (`str`, *optional*, defaults to "default"): + Subfolder in which the data will be stored. Example: `extracted`. + assets_dir (`str`, `Path`, *optional*): + Path to the folder where assets are cached. This must not be the same folder + where Hub files are cached. Defaults to `HF_HOME / "assets"` if not provided. + Can also be set with `HF_ASSETS_CACHE` environment variable. + + Returns: + Path to the cache folder (`Path`). + + Example: + ```py + >>> from huggingface_hub import cached_assets_path + + >>> cached_assets_path(library_name="datasets", namespace="SQuAD", subfolder="download") + PosixPath('/home/wauplin/.cache/huggingface/extra/datasets/SQuAD/download') + + >>> cached_assets_path(library_name="datasets", namespace="SQuAD", subfolder="extracted") + PosixPath('/home/wauplin/.cache/huggingface/extra/datasets/SQuAD/extracted') + + >>> cached_assets_path(library_name="datasets", namespace="Helsinki-NLP/tatoeba_mt") + PosixPath('/home/wauplin/.cache/huggingface/extra/datasets/Helsinki-NLP--tatoeba_mt/default') + + >>> cached_assets_path(library_name="datasets", assets_dir="/tmp/tmp123456") + PosixPath('/tmp/tmp123456/datasets/default/default') + ``` + """ + # Resolve assets_dir + if assets_dir is None: + assets_dir = HF_ASSETS_CACHE + assets_dir = Path(assets_dir).expanduser().resolve() + + # Avoid names that could create path issues + for part in (" ", "/", "\\"): + library_name = library_name.replace(part, "--") + namespace = namespace.replace(part, "--") + subfolder = subfolder.replace(part, "--") + + # Path to subfolder is created + path = assets_dir / library_name / namespace / subfolder + try: + path.mkdir(exist_ok=True, parents=True) + except (FileExistsError, NotADirectoryError): + raise ValueError(f"Corrupted assets folder: cannot create directory because of an existing file ({path}).") + + # Return + return path diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_cache_manager.py b/lib/python3.12/site-packages/huggingface_hub/utils/_cache_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..21469c97aff138a4bd015dc537d5809ef97cf88e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_cache_manager.py @@ -0,0 +1,896 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to manage the HF cache directory.""" + +import os +import shutil +import time +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, FrozenSet, List, Literal, Optional, Set, Union + +from huggingface_hub.errors import CacheNotFound, CorruptedCacheException + +from ..commands._cli_utils import tabulate +from ..constants import HF_HUB_CACHE +from . import logging + + +logger = logging.get_logger(__name__) + +REPO_TYPE_T = Literal["model", "dataset", "space"] + +# List of OS-created helper files that need to be ignored +FILES_TO_IGNORE = [".DS_Store"] + + +@dataclass(frozen=True) +class CachedFileInfo: + """Frozen data structure holding information about a single cached file. + + Args: + file_name (`str`): + Name of the file. Example: `config.json`. + file_path (`Path`): + Path of the file in the `snapshots` directory. The file path is a symlink + referring to a blob in the `blobs` folder. + blob_path (`Path`): + Path of the blob file. This is equivalent to `file_path.resolve()`. + size_on_disk (`int`): + Size of the blob file in bytes. + blob_last_accessed (`float`): + Timestamp of the last time the blob file has been accessed (from any + revision). + blob_last_modified (`float`): + Timestamp of the last time the blob file has been modified/created. + + + + `blob_last_accessed` and `blob_last_modified` reliability can depend on the OS you + are using. See [python documentation](https://docs.python.org/3/library/os.html#os.stat_result) + for more details. + + + """ + + file_name: str + file_path: Path + blob_path: Path + size_on_disk: int + + blob_last_accessed: float + blob_last_modified: float + + @property + def blob_last_accessed_str(self) -> str: + """ + (property) Timestamp of the last time the blob file has been accessed (from any + revision), returned as a human-readable string. + + Example: "2 weeks ago". + """ + return _format_timesince(self.blob_last_accessed) + + @property + def blob_last_modified_str(self) -> str: + """ + (property) Timestamp of the last time the blob file has been modified, returned + as a human-readable string. + + Example: "2 weeks ago". + """ + return _format_timesince(self.blob_last_modified) + + @property + def size_on_disk_str(self) -> str: + """ + (property) Size of the blob file as a human-readable string. + + Example: "42.2K". + """ + return _format_size(self.size_on_disk) + + +@dataclass(frozen=True) +class CachedRevisionInfo: + """Frozen data structure holding information about a revision. + + A revision correspond to a folder in the `snapshots` folder and is populated with + the exact tree structure as the repo on the Hub but contains only symlinks. A + revision can be either referenced by 1 or more `refs` or be "detached" (no refs). + + Args: + commit_hash (`str`): + Hash of the revision (unique). + Example: `"9338f7b671827df886678df2bdd7cc7b4f36dffd"`. + snapshot_path (`Path`): + Path to the revision directory in the `snapshots` folder. It contains the + exact tree structure as the repo on the Hub. + files: (`FrozenSet[CachedFileInfo]`): + Set of [`~CachedFileInfo`] describing all files contained in the snapshot. + refs (`FrozenSet[str]`): + Set of `refs` pointing to this revision. If the revision has no `refs`, it + is considered detached. + Example: `{"main", "2.4.0"}` or `{"refs/pr/1"}`. + size_on_disk (`int`): + Sum of the blob file sizes that are symlink-ed by the revision. + last_modified (`float`): + Timestamp of the last time the revision has been created/modified. + + + + `last_accessed` cannot be determined correctly on a single revision as blob files + are shared across revisions. + + + + + + `size_on_disk` is not necessarily the sum of all file sizes because of possible + duplicated files. Besides, only blobs are taken into account, not the (negligible) + size of folders and symlinks. + + + """ + + commit_hash: str + snapshot_path: Path + size_on_disk: int + files: FrozenSet[CachedFileInfo] + refs: FrozenSet[str] + + last_modified: float + + @property + def last_modified_str(self) -> str: + """ + (property) Timestamp of the last time the revision has been modified, returned + as a human-readable string. + + Example: "2 weeks ago". + """ + return _format_timesince(self.last_modified) + + @property + def size_on_disk_str(self) -> str: + """ + (property) Sum of the blob file sizes as a human-readable string. + + Example: "42.2K". + """ + return _format_size(self.size_on_disk) + + @property + def nb_files(self) -> int: + """ + (property) Total number of files in the revision. + """ + return len(self.files) + + +@dataclass(frozen=True) +class CachedRepoInfo: + """Frozen data structure holding information about a cached repository. + + Args: + repo_id (`str`): + Repo id of the repo on the Hub. Example: `"google/fleurs"`. + repo_type (`Literal["dataset", "model", "space"]`): + Type of the cached repo. + repo_path (`Path`): + Local path to the cached repo. + size_on_disk (`int`): + Sum of the blob file sizes in the cached repo. + nb_files (`int`): + Total number of blob files in the cached repo. + revisions (`FrozenSet[CachedRevisionInfo]`): + Set of [`~CachedRevisionInfo`] describing all revisions cached in the repo. + last_accessed (`float`): + Timestamp of the last time a blob file of the repo has been accessed. + last_modified (`float`): + Timestamp of the last time a blob file of the repo has been modified/created. + + + + `size_on_disk` is not necessarily the sum of all revisions sizes because of + duplicated files. Besides, only blobs are taken into account, not the (negligible) + size of folders and symlinks. + + + + + + `last_accessed` and `last_modified` reliability can depend on the OS you are using. + See [python documentation](https://docs.python.org/3/library/os.html#os.stat_result) + for more details. + + + """ + + repo_id: str + repo_type: REPO_TYPE_T + repo_path: Path + size_on_disk: int + nb_files: int + revisions: FrozenSet[CachedRevisionInfo] + + last_accessed: float + last_modified: float + + @property + def last_accessed_str(self) -> str: + """ + (property) Last time a blob file of the repo has been accessed, returned as a + human-readable string. + + Example: "2 weeks ago". + """ + return _format_timesince(self.last_accessed) + + @property + def last_modified_str(self) -> str: + """ + (property) Last time a blob file of the repo has been modified, returned as a + human-readable string. + + Example: "2 weeks ago". + """ + return _format_timesince(self.last_modified) + + @property + def size_on_disk_str(self) -> str: + """ + (property) Sum of the blob file sizes as a human-readable string. + + Example: "42.2K". + """ + return _format_size(self.size_on_disk) + + @property + def refs(self) -> Dict[str, CachedRevisionInfo]: + """ + (property) Mapping between `refs` and revision data structures. + """ + return {ref: revision for revision in self.revisions for ref in revision.refs} + + +@dataclass(frozen=True) +class DeleteCacheStrategy: + """Frozen data structure holding the strategy to delete cached revisions. + + This object is not meant to be instantiated programmatically but to be returned by + [`~utils.HFCacheInfo.delete_revisions`]. See documentation for usage example. + + Args: + expected_freed_size (`float`): + Expected freed size once strategy is executed. + blobs (`FrozenSet[Path]`): + Set of blob file paths to be deleted. + refs (`FrozenSet[Path]`): + Set of reference file paths to be deleted. + repos (`FrozenSet[Path]`): + Set of entire repo paths to be deleted. + snapshots (`FrozenSet[Path]`): + Set of snapshots to be deleted (directory of symlinks). + """ + + expected_freed_size: int + blobs: FrozenSet[Path] + refs: FrozenSet[Path] + repos: FrozenSet[Path] + snapshots: FrozenSet[Path] + + @property + def expected_freed_size_str(self) -> str: + """ + (property) Expected size that will be freed as a human-readable string. + + Example: "42.2K". + """ + return _format_size(self.expected_freed_size) + + def execute(self) -> None: + """Execute the defined strategy. + + + + If this method is interrupted, the cache might get corrupted. Deletion order is + implemented so that references and symlinks are deleted before the actual blob + files. + + + + + + This method is irreversible. If executed, cached files are erased and must be + downloaded again. + + + """ + # Deletion order matters. Blobs are deleted in last so that the user can't end + # up in a state where a `ref`` refers to a missing snapshot or a snapshot + # symlink refers to a deleted blob. + + # Delete entire repos + for path in self.repos: + _try_delete_path(path, path_type="repo") + + # Delete snapshot directories + for path in self.snapshots: + _try_delete_path(path, path_type="snapshot") + + # Delete refs files + for path in self.refs: + _try_delete_path(path, path_type="ref") + + # Delete blob files + for path in self.blobs: + _try_delete_path(path, path_type="blob") + + logger.info(f"Cache deletion done. Saved {self.expected_freed_size_str}.") + + +@dataclass(frozen=True) +class HFCacheInfo: + """Frozen data structure holding information about the entire cache-system. + + This data structure is returned by [`scan_cache_dir`] and is immutable. + + Args: + size_on_disk (`int`): + Sum of all valid repo sizes in the cache-system. + repos (`FrozenSet[CachedRepoInfo]`): + Set of [`~CachedRepoInfo`] describing all valid cached repos found on the + cache-system while scanning. + warnings (`List[CorruptedCacheException]`): + List of [`~CorruptedCacheException`] that occurred while scanning the cache. + Those exceptions are captured so that the scan can continue. Corrupted repos + are skipped from the scan. + + + + Here `size_on_disk` is equal to the sum of all repo sizes (only blobs). However if + some cached repos are corrupted, their sizes are not taken into account. + + + """ + + size_on_disk: int + repos: FrozenSet[CachedRepoInfo] + warnings: List[CorruptedCacheException] + + @property + def size_on_disk_str(self) -> str: + """ + (property) Sum of all valid repo sizes in the cache-system as a human-readable + string. + + Example: "42.2K". + """ + return _format_size(self.size_on_disk) + + def delete_revisions(self, *revisions: str) -> DeleteCacheStrategy: + """Prepare the strategy to delete one or more revisions cached locally. + + Input revisions can be any revision hash. If a revision hash is not found in the + local cache, a warning is thrown but no error is raised. Revisions can be from + different cached repos since hashes are unique across repos, + + Examples: + ```py + >>> from huggingface_hub import scan_cache_dir + >>> cache_info = scan_cache_dir() + >>> delete_strategy = cache_info.delete_revisions( + ... "81fd1d6e7847c99f5862c9fb81387956d99ec7aa" + ... ) + >>> print(f"Will free {delete_strategy.expected_freed_size_str}.") + Will free 7.9K. + >>> delete_strategy.execute() + Cache deletion done. Saved 7.9K. + ``` + + ```py + >>> from huggingface_hub import scan_cache_dir + >>> scan_cache_dir().delete_revisions( + ... "81fd1d6e7847c99f5862c9fb81387956d99ec7aa", + ... "e2983b237dccf3ab4937c97fa717319a9ca1a96d", + ... "6c0e6080953db56375760c0471a8c5f2929baf11", + ... ).execute() + Cache deletion done. Saved 8.6G. + ``` + + + + `delete_revisions` returns a [`~utils.DeleteCacheStrategy`] object that needs to + be executed. The [`~utils.DeleteCacheStrategy`] is not meant to be modified but + allows having a dry run before actually executing the deletion. + + + """ + hashes_to_delete: Set[str] = set(revisions) + + repos_with_revisions: Dict[CachedRepoInfo, Set[CachedRevisionInfo]] = defaultdict(set) + + for repo in self.repos: + for revision in repo.revisions: + if revision.commit_hash in hashes_to_delete: + repos_with_revisions[repo].add(revision) + hashes_to_delete.remove(revision.commit_hash) + + if len(hashes_to_delete) > 0: + logger.warning(f"Revision(s) not found - cannot delete them: {', '.join(hashes_to_delete)}") + + delete_strategy_blobs: Set[Path] = set() + delete_strategy_refs: Set[Path] = set() + delete_strategy_repos: Set[Path] = set() + delete_strategy_snapshots: Set[Path] = set() + delete_strategy_expected_freed_size = 0 + + for affected_repo, revisions_to_delete in repos_with_revisions.items(): + other_revisions = affected_repo.revisions - revisions_to_delete + + # If no other revisions, it means all revisions are deleted + # -> delete the entire cached repo + if len(other_revisions) == 0: + delete_strategy_repos.add(affected_repo.repo_path) + delete_strategy_expected_freed_size += affected_repo.size_on_disk + continue + + # Some revisions of the repo will be deleted but not all. We need to filter + # which blob files will not be linked anymore. + for revision_to_delete in revisions_to_delete: + # Snapshot dir + delete_strategy_snapshots.add(revision_to_delete.snapshot_path) + + # Refs dir + for ref in revision_to_delete.refs: + delete_strategy_refs.add(affected_repo.repo_path / "refs" / ref) + + # Blobs dir + for file in revision_to_delete.files: + if file.blob_path not in delete_strategy_blobs: + is_file_alone = True + for revision in other_revisions: + for rev_file in revision.files: + if file.blob_path == rev_file.blob_path: + is_file_alone = False + break + if not is_file_alone: + break + + # Blob file not referenced by remaining revisions -> delete + if is_file_alone: + delete_strategy_blobs.add(file.blob_path) + delete_strategy_expected_freed_size += file.size_on_disk + + # Return the strategy instead of executing it. + return DeleteCacheStrategy( + blobs=frozenset(delete_strategy_blobs), + refs=frozenset(delete_strategy_refs), + repos=frozenset(delete_strategy_repos), + snapshots=frozenset(delete_strategy_snapshots), + expected_freed_size=delete_strategy_expected_freed_size, + ) + + def export_as_table(self, *, verbosity: int = 0) -> str: + """Generate a table from the [`HFCacheInfo`] object. + + Pass `verbosity=0` to get a table with a single row per repo, with columns + "repo_id", "repo_type", "size_on_disk", "nb_files", "last_accessed", "last_modified", "refs", "local_path". + + Pass `verbosity=1` to get a table with a row per repo and revision (thus multiple rows can appear for a single repo), with columns + "repo_id", "repo_type", "revision", "size_on_disk", "nb_files", "last_modified", "refs", "local_path". + + Example: + ```py + >>> from huggingface_hub.utils import scan_cache_dir + + >>> hf_cache_info = scan_cache_dir() + HFCacheInfo(...) + + >>> print(hf_cache_info.export_as_table()) + REPO ID REPO TYPE SIZE ON DISK NB FILES LAST_ACCESSED LAST_MODIFIED REFS LOCAL PATH + --------------------------------------------------- --------- ------------ -------- ------------- ------------- ---- -------------------------------------------------------------------------------------------------- + roberta-base model 2.7M 5 1 day ago 1 week ago main ~/.cache/huggingface/hub/models--roberta-base + suno/bark model 8.8K 1 1 week ago 1 week ago main ~/.cache/huggingface/hub/models--suno--bark + t5-base model 893.8M 4 4 days ago 7 months ago main ~/.cache/huggingface/hub/models--t5-base + t5-large model 3.0G 4 5 weeks ago 5 months ago main ~/.cache/huggingface/hub/models--t5-large + + >>> print(hf_cache_info.export_as_table(verbosity=1)) + REPO ID REPO TYPE REVISION SIZE ON DISK NB FILES LAST_MODIFIED REFS LOCAL PATH + --------------------------------------------------- --------- ---------------------------------------- ------------ -------- ------------- ---- ----------------------------------------------------------------------------------------------------------------------------------------------------- + roberta-base model e2da8e2f811d1448a5b465c236feacd80ffbac7b 2.7M 5 1 week ago main ~/.cache/huggingface/hub/models--roberta-base/snapshots/e2da8e2f811d1448a5b465c236feacd80ffbac7b + suno/bark model 70a8a7d34168586dc5d028fa9666aceade177992 8.8K 1 1 week ago main ~/.cache/huggingface/hub/models--suno--bark/snapshots/70a8a7d34168586dc5d028fa9666aceade177992 + t5-base model a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1 893.8M 4 7 months ago main ~/.cache/huggingface/hub/models--t5-base/snapshots/a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1 + t5-large model 150ebc2c4b72291e770f58e6057481c8d2ed331a 3.0G 4 5 months ago main ~/.cache/huggingface/hub/models--t5-large/snapshots/150ebc2c4b72291e770f58e6057481c8d2ed331a + ``` + + Args: + verbosity (`int`, *optional*): + The verbosity level. Defaults to 0. + + Returns: + `str`: The table as a string. + """ + if verbosity == 0: + return tabulate( + rows=[ + [ + repo.repo_id, + repo.repo_type, + "{:>12}".format(repo.size_on_disk_str), + repo.nb_files, + repo.last_accessed_str, + repo.last_modified_str, + ", ".join(sorted(repo.refs)), + str(repo.repo_path), + ] + for repo in sorted(self.repos, key=lambda repo: repo.repo_path) + ], + headers=[ + "REPO ID", + "REPO TYPE", + "SIZE ON DISK", + "NB FILES", + "LAST_ACCESSED", + "LAST_MODIFIED", + "REFS", + "LOCAL PATH", + ], + ) + else: + return tabulate( + rows=[ + [ + repo.repo_id, + repo.repo_type, + revision.commit_hash, + "{:>12}".format(revision.size_on_disk_str), + revision.nb_files, + revision.last_modified_str, + ", ".join(sorted(revision.refs)), + str(revision.snapshot_path), + ] + for repo in sorted(self.repos, key=lambda repo: repo.repo_path) + for revision in sorted(repo.revisions, key=lambda revision: revision.commit_hash) + ], + headers=[ + "REPO ID", + "REPO TYPE", + "REVISION", + "SIZE ON DISK", + "NB FILES", + "LAST_MODIFIED", + "REFS", + "LOCAL PATH", + ], + ) + + +def scan_cache_dir(cache_dir: Optional[Union[str, Path]] = None) -> HFCacheInfo: + """Scan the entire HF cache-system and return a [`~HFCacheInfo`] structure. + + Use `scan_cache_dir` in order to programmatically scan your cache-system. The cache + will be scanned repo by repo. If a repo is corrupted, a [`~CorruptedCacheException`] + will be thrown internally but captured and returned in the [`~HFCacheInfo`] + structure. Only valid repos get a proper report. + + ```py + >>> from huggingface_hub import scan_cache_dir + + >>> hf_cache_info = scan_cache_dir() + HFCacheInfo( + size_on_disk=3398085269, + repos=frozenset({ + CachedRepoInfo( + repo_id='t5-small', + repo_type='model', + repo_path=PosixPath(...), + size_on_disk=970726914, + nb_files=11, + revisions=frozenset({ + CachedRevisionInfo( + commit_hash='d78aea13fa7ecd06c29e3e46195d6341255065d5', + size_on_disk=970726339, + snapshot_path=PosixPath(...), + files=frozenset({ + CachedFileInfo( + file_name='config.json', + size_on_disk=1197 + file_path=PosixPath(...), + blob_path=PosixPath(...), + ), + CachedFileInfo(...), + ... + }), + ), + CachedRevisionInfo(...), + ... + }), + ), + CachedRepoInfo(...), + ... + }), + warnings=[ + CorruptedCacheException("Snapshots dir doesn't exist in cached repo: ..."), + CorruptedCacheException(...), + ... + ], + ) + ``` + + You can also print a detailed report directly from the `huggingface-cli` using: + ```text + > huggingface-cli scan-cache + REPO ID REPO TYPE SIZE ON DISK NB FILES REFS LOCAL PATH + --------------------------- --------- ------------ -------- ------------------- ------------------------------------------------------------------------- + glue dataset 116.3K 15 1.17.0, main, 2.4.0 /Users/lucain/.cache/huggingface/hub/datasets--glue + google/fleurs dataset 64.9M 6 main, refs/pr/1 /Users/lucain/.cache/huggingface/hub/datasets--google--fleurs + Jean-Baptiste/camembert-ner model 441.0M 7 main /Users/lucain/.cache/huggingface/hub/models--Jean-Baptiste--camembert-ner + bert-base-cased model 1.9G 13 main /Users/lucain/.cache/huggingface/hub/models--bert-base-cased + t5-base model 10.1K 3 main /Users/lucain/.cache/huggingface/hub/models--t5-base + t5-small model 970.7M 11 refs/pr/1, main /Users/lucain/.cache/huggingface/hub/models--t5-small + + Done in 0.0s. Scanned 6 repo(s) for a total of 3.4G. + Got 1 warning(s) while scanning. Use -vvv to print details. + ``` + + Args: + cache_dir (`str` or `Path`, `optional`): + Cache directory to cache. Defaults to the default HF cache directory. + + + + Raises: + + `CacheNotFound` + If the cache directory does not exist. + + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the cache directory is a file, instead of a directory. + + + + Returns: a [`~HFCacheInfo`] object. + """ + if cache_dir is None: + cache_dir = HF_HUB_CACHE + + cache_dir = Path(cache_dir).expanduser().resolve() + if not cache_dir.exists(): + raise CacheNotFound( + f"Cache directory not found: {cache_dir}. Please use `cache_dir` argument or set `HF_HUB_CACHE` environment variable.", + cache_dir=cache_dir, + ) + + if cache_dir.is_file(): + raise ValueError( + f"Scan cache expects a directory but found a file: {cache_dir}. Please use `cache_dir` argument or set `HF_HUB_CACHE` environment variable." + ) + + repos: Set[CachedRepoInfo] = set() + warnings: List[CorruptedCacheException] = [] + for repo_path in cache_dir.iterdir(): + if repo_path.name == ".locks": # skip './.locks/' folder + continue + try: + repos.add(_scan_cached_repo(repo_path)) + except CorruptedCacheException as e: + warnings.append(e) + + return HFCacheInfo( + repos=frozenset(repos), + size_on_disk=sum(repo.size_on_disk for repo in repos), + warnings=warnings, + ) + + +def _scan_cached_repo(repo_path: Path) -> CachedRepoInfo: + """Scan a single cache repo and return information about it. + + Any unexpected behavior will raise a [`~CorruptedCacheException`]. + """ + if not repo_path.is_dir(): + raise CorruptedCacheException(f"Repo path is not a directory: {repo_path}") + + if "--" not in repo_path.name: + raise CorruptedCacheException(f"Repo path is not a valid HuggingFace cache directory: {repo_path}") + + repo_type, repo_id = repo_path.name.split("--", maxsplit=1) + repo_type = repo_type[:-1] # "models" -> "model" + repo_id = repo_id.replace("--", "/") # google/fleurs -> "google/fleurs" + + if repo_type not in {"dataset", "model", "space"}: + raise CorruptedCacheException( + f"Repo type must be `dataset`, `model` or `space`, found `{repo_type}` ({repo_path})." + ) + + blob_stats: Dict[Path, os.stat_result] = {} # Key is blob_path, value is blob stats + + snapshots_path = repo_path / "snapshots" + refs_path = repo_path / "refs" + + if not snapshots_path.exists() or not snapshots_path.is_dir(): + raise CorruptedCacheException(f"Snapshots dir doesn't exist in cached repo: {snapshots_path}") + + # Scan over `refs` directory + + # key is revision hash, value is set of refs + refs_by_hash: Dict[str, Set[str]] = defaultdict(set) + if refs_path.exists(): + # Example of `refs` directory + # ── refs + # ├── main + # └── refs + # └── pr + # └── 1 + if refs_path.is_file(): + raise CorruptedCacheException(f"Refs directory cannot be a file: {refs_path}") + + for ref_path in refs_path.glob("**/*"): + # glob("**/*") iterates over all files and directories -> skip directories + if ref_path.is_dir() or ref_path.name in FILES_TO_IGNORE: + continue + + ref_name = str(ref_path.relative_to(refs_path)) + with ref_path.open() as f: + commit_hash = f.read() + + refs_by_hash[commit_hash].add(ref_name) + + # Scan snapshots directory + cached_revisions: Set[CachedRevisionInfo] = set() + for revision_path in snapshots_path.iterdir(): + # Ignore OS-created helper files + if revision_path.name in FILES_TO_IGNORE: + continue + if revision_path.is_file(): + raise CorruptedCacheException(f"Snapshots folder corrupted. Found a file: {revision_path}") + + cached_files = set() + for file_path in revision_path.glob("**/*"): + # glob("**/*") iterates over all files and directories -> skip directories + if file_path.is_dir(): + continue + + blob_path = Path(file_path).resolve() + if not blob_path.exists(): + raise CorruptedCacheException(f"Blob missing (broken symlink): {blob_path}") + + if blob_path not in blob_stats: + blob_stats[blob_path] = blob_path.stat() + + cached_files.add( + CachedFileInfo( + file_name=file_path.name, + file_path=file_path, + size_on_disk=blob_stats[blob_path].st_size, + blob_path=blob_path, + blob_last_accessed=blob_stats[blob_path].st_atime, + blob_last_modified=blob_stats[blob_path].st_mtime, + ) + ) + + # Last modified is either the last modified blob file or the revision folder + # itself if it is empty + if len(cached_files) > 0: + revision_last_modified = max(blob_stats[file.blob_path].st_mtime for file in cached_files) + else: + revision_last_modified = revision_path.stat().st_mtime + + cached_revisions.add( + CachedRevisionInfo( + commit_hash=revision_path.name, + files=frozenset(cached_files), + refs=frozenset(refs_by_hash.pop(revision_path.name, set())), + size_on_disk=sum( + blob_stats[blob_path].st_size for blob_path in set(file.blob_path for file in cached_files) + ), + snapshot_path=revision_path, + last_modified=revision_last_modified, + ) + ) + + # Check that all refs referred to an existing revision + if len(refs_by_hash) > 0: + raise CorruptedCacheException( + f"Reference(s) refer to missing commit hashes: {dict(refs_by_hash)} ({repo_path})." + ) + + # Last modified is either the last modified blob file or the repo folder itself if + # no blob files has been found. Same for last accessed. + if len(blob_stats) > 0: + repo_last_accessed = max(stat.st_atime for stat in blob_stats.values()) + repo_last_modified = max(stat.st_mtime for stat in blob_stats.values()) + else: + repo_stats = repo_path.stat() + repo_last_accessed = repo_stats.st_atime + repo_last_modified = repo_stats.st_mtime + + # Build and return frozen structure + return CachedRepoInfo( + nb_files=len(blob_stats), + repo_id=repo_id, + repo_path=repo_path, + repo_type=repo_type, # type: ignore + revisions=frozenset(cached_revisions), + size_on_disk=sum(stat.st_size for stat in blob_stats.values()), + last_accessed=repo_last_accessed, + last_modified=repo_last_modified, + ) + + +def _format_size(num: int) -> str: + """Format size in bytes into a human-readable string. + + Taken from https://stackoverflow.com/a/1094933 + """ + num_f = float(num) + for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]: + if abs(num_f) < 1000.0: + return f"{num_f:3.1f}{unit}" + num_f /= 1000.0 + return f"{num_f:.1f}Y" + + +_TIMESINCE_CHUNKS = ( + # Label, divider, max value + ("second", 1, 60), + ("minute", 60, 60), + ("hour", 60 * 60, 24), + ("day", 60 * 60 * 24, 6), + ("week", 60 * 60 * 24 * 7, 6), + ("month", 60 * 60 * 24 * 30, 11), + ("year", 60 * 60 * 24 * 365, None), +) + + +def _format_timesince(ts: float) -> str: + """Format timestamp in seconds into a human-readable string, relative to now. + + Vaguely inspired by Django's `timesince` formatter. + """ + delta = time.time() - ts + if delta < 20: + return "a few seconds ago" + for label, divider, max_value in _TIMESINCE_CHUNKS: # noqa: B007 + value = round(delta / divider) + if max_value is not None and value <= max_value: + break + return f"{value} {label}{'s' if value > 1 else ''} ago" + + +def _try_delete_path(path: Path, path_type: str) -> None: + """Try to delete a local file or folder. + + If the path does not exists, error is logged as a warning and then ignored. + + Args: + path (`Path`) + Path to delete. Can be a file or a folder. + path_type (`str`) + What path are we deleting ? Only for logging purposes. Example: "snapshot". + """ + logger.info(f"Delete {path_type}: {path}") + try: + if path.is_file(): + os.remove(path) + else: + shutil.rmtree(path) + except FileNotFoundError: + logger.warning(f"Couldn't delete {path_type}: file not found ({path})", exc_info=True) + except PermissionError: + logger.warning(f"Couldn't delete {path_type}: permission denied ({path})", exc_info=True) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_chunk_utils.py b/lib/python3.12/site-packages/huggingface_hub/utils/_chunk_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b0af032ae6a68f03676ad7fdb8e483248d9853f8 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_chunk_utils.py @@ -0,0 +1,65 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains a utility to iterate by chunks over an iterator.""" + +import itertools +from typing import Iterable, TypeVar + + +T = TypeVar("T") + + +def chunk_iterable(iterable: Iterable[T], chunk_size: int) -> Iterable[Iterable[T]]: + """Iterates over an iterator chunk by chunk. + + Taken from https://stackoverflow.com/a/8998040. + See also https://github.com/huggingface/huggingface_hub/pull/920#discussion_r938793088. + + Args: + iterable (`Iterable`): + The iterable on which we want to iterate. + chunk_size (`int`): + Size of the chunks. Must be a strictly positive integer (e.g. >0). + + Example: + + ```python + >>> from huggingface_hub.utils import chunk_iterable + + >>> for items in chunk_iterable(range(17), chunk_size=8): + ... print(items) + # [0, 1, 2, 3, 4, 5, 6, 7] + # [8, 9, 10, 11, 12, 13, 14, 15] + # [16] # smaller last chunk + ``` + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If `chunk_size` <= 0. + + + The last chunk can be smaller than `chunk_size`. + + """ + if not isinstance(chunk_size, int) or chunk_size <= 0: + raise ValueError("`chunk_size` must be a strictly positive integer (>0).") + + iterator = iter(iterable) + while True: + try: + next_item = next(iterator) + except StopIteration: + return + yield itertools.chain((next_item,), itertools.islice(iterator, chunk_size - 1)) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_datetime.py b/lib/python3.12/site-packages/huggingface_hub/utils/_datetime.py new file mode 100644 index 0000000000000000000000000000000000000000..1a7f44285d1c826006c97176ca66c3e9c33f61c0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_datetime.py @@ -0,0 +1,67 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle datetimes in Huggingface Hub.""" + +from datetime import datetime, timezone + + +def parse_datetime(date_string: str) -> datetime: + """ + Parses a date_string returned from the server to a datetime object. + + This parser is a weak-parser is the sense that it handles only a single format of + date_string. It is expected that the server format will never change. The + implementation depends only on the standard lib to avoid an external dependency + (python-dateutil). See full discussion about this decision on PR: + https://github.com/huggingface/huggingface_hub/pull/999. + + Example: + ```py + > parse_datetime('2022-08-19T07:19:38.123Z') + datetime.datetime(2022, 8, 19, 7, 19, 38, 123000, tzinfo=timezone.utc) + ``` + + Args: + date_string (`str`): + A string representing a datetime returned by the Hub server. + String is expected to follow '%Y-%m-%dT%H:%M:%S.%fZ' pattern. + + Returns: + A python datetime object. + + Raises: + :class:`ValueError`: + If `date_string` cannot be parsed. + """ + try: + # Normalize the string to always have 6 digits of fractional seconds + if date_string.endswith("Z"): + # Case 1: No decimal point (e.g., "2024-11-16T00:27:02Z") + if "." not in date_string: + # No fractional seconds - insert .000000 + date_string = date_string[:-1] + ".000000Z" + # Case 2: Has decimal point (e.g., "2022-08-19T07:19:38.123456789Z") + else: + # Get the fractional and base parts + base, fraction = date_string[:-1].split(".") + # fraction[:6] takes first 6 digits and :0<6 pads with zeros if less than 6 digits + date_string = f"{base}.{fraction[:6]:0<6}Z" + + return datetime.strptime(date_string, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=timezone.utc) + except ValueError as e: + raise ValueError( + f"Cannot parse '{date_string}' as a datetime. Date string is expected to" + " follow '%Y-%m-%dT%H:%M:%S.%fZ' pattern." + ) from e diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_deprecation.py b/lib/python3.12/site-packages/huggingface_hub/utils/_deprecation.py new file mode 100644 index 0000000000000000000000000000000000000000..4cb8d6e418c76accd1ecd61158b4bdd265e12f71 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_deprecation.py @@ -0,0 +1,136 @@ +import warnings +from functools import wraps +from inspect import Parameter, signature +from typing import Iterable, Optional + + +def _deprecate_positional_args(*, version: str): + """Decorator for methods that issues warnings for positional arguments. + Using the keyword-only argument syntax in pep 3102, arguments after the + * will issue a warning when passed as a positional argument. + + Args: + version (`str`): + The version when positional arguments will result in error. + """ + + def _inner_deprecate_positional_args(f): + sig = signature(f) + kwonly_args = [] + all_args = [] + for name, param in sig.parameters.items(): + if param.kind == Parameter.POSITIONAL_OR_KEYWORD: + all_args.append(name) + elif param.kind == Parameter.KEYWORD_ONLY: + kwonly_args.append(name) + + @wraps(f) + def inner_f(*args, **kwargs): + extra_args = len(args) - len(all_args) + if extra_args <= 0: + return f(*args, **kwargs) + # extra_args > 0 + args_msg = [ + f"{name}='{arg}'" if isinstance(arg, str) else f"{name}={arg}" + for name, arg in zip(kwonly_args[:extra_args], args[-extra_args:]) + ] + args_msg = ", ".join(args_msg) + warnings.warn( + f"Deprecated positional argument(s) used in '{f.__name__}': pass" + f" {args_msg} as keyword args. From version {version} passing these" + " as positional arguments will result in an error,", + FutureWarning, + ) + kwargs.update(zip(sig.parameters, args)) + return f(**kwargs) + + return inner_f + + return _inner_deprecate_positional_args + + +def _deprecate_arguments( + *, + version: str, + deprecated_args: Iterable[str], + custom_message: Optional[str] = None, +): + """Decorator to issue warnings when using deprecated arguments. + + TODO: could be useful to be able to set a custom error message. + + Args: + version (`str`): + The version when deprecated arguments will result in error. + deprecated_args (`List[str]`): + List of the arguments to be deprecated. + custom_message (`str`, *optional*): + Warning message that is raised. If not passed, a default warning message + will be created. + """ + + def _inner_deprecate_positional_args(f): + sig = signature(f) + + @wraps(f) + def inner_f(*args, **kwargs): + # Check for used deprecated arguments + used_deprecated_args = [] + for _, parameter in zip(args, sig.parameters.values()): + if parameter.name in deprecated_args: + used_deprecated_args.append(parameter.name) + for kwarg_name, kwarg_value in kwargs.items(): + if ( + # If argument is deprecated but still used + kwarg_name in deprecated_args + # And then the value is not the default value + and kwarg_value != sig.parameters[kwarg_name].default + ): + used_deprecated_args.append(kwarg_name) + + # Warn and proceed + if len(used_deprecated_args) > 0: + message = ( + f"Deprecated argument(s) used in '{f.__name__}':" + f" {', '.join(used_deprecated_args)}. Will not be supported from" + f" version '{version}'." + ) + if custom_message is not None: + message += "\n\n" + custom_message + warnings.warn(message, FutureWarning) + return f(*args, **kwargs) + + return inner_f + + return _inner_deprecate_positional_args + + +def _deprecate_method(*, version: str, message: Optional[str] = None): + """Decorator to issue warnings when using a deprecated method. + + Args: + version (`str`): + The version when deprecated arguments will result in error. + message (`str`, *optional*): + Warning message that is raised. If not passed, a default warning message + will be created. + """ + + def _inner_deprecate_method(f): + name = f.__name__ + if name == "__init__": + name = f.__qualname__.split(".")[0] # class name instead of method name + + @wraps(f) + def inner_f(*args, **kwargs): + warning_message = ( + f"'{name}' (from '{f.__module__}') is deprecated and will be removed from version '{version}'." + ) + if message is not None: + warning_message += " " + message + warnings.warn(warning_message, FutureWarning) + return f(*args, **kwargs) + + return inner_f + + return _inner_deprecate_method diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_experimental.py b/lib/python3.12/site-packages/huggingface_hub/utils/_experimental.py new file mode 100644 index 0000000000000000000000000000000000000000..34141eba09123c06fbca55c929a19a0264e5788e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_experimental.py @@ -0,0 +1,66 @@ +# coding=utf-8 +# Copyright 2023-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to flag a feature as "experimental" in Huggingface Hub.""" + +import warnings +from functools import wraps +from typing import Callable + +from .. import constants + + +def experimental(fn: Callable) -> Callable: + """Decorator to flag a feature as experimental. + + An experimental feature trigger a warning when used as it might be subject to breaking changes in the future. + Warnings can be disabled by setting the environment variable `HF_EXPERIMENTAL_WARNING` to `0`. + + Args: + fn (`Callable`): + The function to flag as experimental. + + Returns: + `Callable`: The decorated function. + + Example: + + ```python + >>> from huggingface_hub.utils import experimental + + >>> @experimental + ... def my_function(): + ... print("Hello world!") + + >>> my_function() + UserWarning: 'my_function' is experimental and might be subject to breaking changes in the future. You can disable + this warning by setting `HF_HUB_DISABLE_EXPERIMENTAL_WARNING=1` as environment variable. + Hello world! + ``` + """ + # For classes, put the "experimental" around the "__new__" method => __new__ will be removed in warning message + name = fn.__qualname__[: -len(".__new__")] if fn.__qualname__.endswith(".__new__") else fn.__qualname__ + + @wraps(fn) + def _inner_fn(*args, **kwargs): + if not constants.HF_HUB_DISABLE_EXPERIMENTAL_WARNING: + warnings.warn( + f"'{name}' is experimental and might be subject to breaking changes in the future." + " You can disable this warning by setting `HF_HUB_DISABLE_EXPERIMENTAL_WARNING=1` as environment" + " variable.", + UserWarning, + ) + return fn(*args, **kwargs) + + return _inner_fn diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_fixes.py b/lib/python3.12/site-packages/huggingface_hub/utils/_fixes.py new file mode 100644 index 0000000000000000000000000000000000000000..560003b6222058b03791491b1ce70ea9d7a94404 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_fixes.py @@ -0,0 +1,133 @@ +# JSONDecodeError was introduced in requests=2.27 released in 2022. +# This allows us to support older requests for users +# More information: https://github.com/psf/requests/pull/5856 +try: + from requests import JSONDecodeError # type: ignore # noqa: F401 +except ImportError: + try: + from simplejson import JSONDecodeError # type: ignore # noqa: F401 + except ImportError: + from json import JSONDecodeError # type: ignore # noqa: F401 +import contextlib +import os +import shutil +import stat +import tempfile +import time +from functools import partial +from pathlib import Path +from typing import Callable, Generator, Optional, Union + +import yaml +from filelock import BaseFileLock, FileLock, SoftFileLock, Timeout + +from .. import constants +from . import logging + + +logger = logging.get_logger(__name__) + +# Wrap `yaml.dump` to set `allow_unicode=True` by default. +# +# Example: +# ```py +# >>> yaml.dump({"emoji": "👀", "some unicode": "日本か"}) +# 'emoji: "\\U0001F440"\nsome unicode: "\\u65E5\\u672C\\u304B"\n' +# +# >>> yaml_dump({"emoji": "👀", "some unicode": "日本か"}) +# 'emoji: "👀"\nsome unicode: "日本か"\n' +# ``` +yaml_dump: Callable[..., str] = partial(yaml.dump, stream=None, allow_unicode=True) # type: ignore + + +@contextlib.contextmanager +def SoftTemporaryDirectory( + suffix: Optional[str] = None, + prefix: Optional[str] = None, + dir: Optional[Union[Path, str]] = None, + **kwargs, +) -> Generator[Path, None, None]: + """ + Context manager to create a temporary directory and safely delete it. + + If tmp directory cannot be deleted normally, we set the WRITE permission and retry. + If cleanup still fails, we give up but don't raise an exception. This is equivalent + to `tempfile.TemporaryDirectory(..., ignore_cleanup_errors=True)` introduced in + Python 3.10. + + See https://www.scivision.dev/python-tempfile-permission-error-windows/. + """ + tmpdir = tempfile.TemporaryDirectory(prefix=prefix, suffix=suffix, dir=dir, **kwargs) + yield Path(tmpdir.name).resolve() + + try: + # First once with normal cleanup + shutil.rmtree(tmpdir.name) + except Exception: + # If failed, try to set write permission and retry + try: + shutil.rmtree(tmpdir.name, onerror=_set_write_permission_and_retry) + except Exception: + pass + + # And finally, cleanup the tmpdir. + # If it fails again, give up but do not throw error + try: + tmpdir.cleanup() + except Exception: + pass + + +def _set_write_permission_and_retry(func, path, excinfo): + os.chmod(path, stat.S_IWRITE) + func(path) + + +@contextlib.contextmanager +def WeakFileLock( + lock_file: Union[str, Path], *, timeout: Optional[float] = None +) -> Generator[BaseFileLock, None, None]: + """A filelock with some custom logic. + + This filelock is weaker than the default filelock in that: + 1. It won't raise an exception if release fails. + 2. It will default to a SoftFileLock if the filesystem does not support flock. + + An INFO log message is emitted every 10 seconds if the lock is not acquired immediately. + If a timeout is provided, a `filelock.Timeout` exception is raised if the lock is not acquired within the timeout. + """ + log_interval = constants.FILELOCK_LOG_EVERY_SECONDS + lock = FileLock(lock_file, timeout=log_interval) + start_time = time.time() + + while True: + elapsed_time = time.time() - start_time + if timeout is not None and elapsed_time >= timeout: + raise Timeout(str(lock_file)) + + try: + lock.acquire(timeout=min(log_interval, timeout - elapsed_time) if timeout else log_interval) + except Timeout: + logger.info( + f"Still waiting to acquire lock on {lock_file} (elapsed: {time.time() - start_time:.1f} seconds)" + ) + except NotImplementedError as e: + if "use SoftFileLock instead" in str(e): + logger.warning( + "FileSystem does not appear to support flock. Falling back to SoftFileLock for %s", lock_file + ) + lock = SoftFileLock(lock_file, timeout=log_interval) + continue + else: + break + + try: + yield lock + finally: + try: + lock.release() + except OSError: + try: + Path(lock_file).unlink() + except OSError: + pass diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_git_credential.py b/lib/python3.12/site-packages/huggingface_hub/utils/_git_credential.py new file mode 100644 index 0000000000000000000000000000000000000000..a8ed77f4e49ca88ff4fa9aba48cbf00195036013 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_git_credential.py @@ -0,0 +1,121 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to manage Git credentials.""" + +import re +import subprocess +from typing import List, Optional + +from ..constants import ENDPOINT +from ._subprocess import run_interactive_subprocess, run_subprocess + + +GIT_CREDENTIAL_REGEX = re.compile( + r""" + ^\s* # start of line + credential\.helper # credential.helper value + \s*=\s* # separator + (\w+) # the helper name (group 1) + (\s|$) # whitespace or end of line + """, + flags=re.MULTILINE | re.IGNORECASE | re.VERBOSE, +) + + +def list_credential_helpers(folder: Optional[str] = None) -> List[str]: + """Return the list of git credential helpers configured. + + See https://git-scm.com/docs/gitcredentials. + + Credentials are saved in all configured helpers (store, cache, macOS keychain,...). + Calls "`git credential approve`" internally. See https://git-scm.com/docs/git-credential. + + Args: + folder (`str`, *optional*): + The folder in which to check the configured helpers. + """ + try: + output = run_subprocess("git config --list", folder=folder).stdout + parsed = _parse_credential_output(output) + return parsed + except subprocess.CalledProcessError as exc: + raise EnvironmentError(exc.stderr) + + +def set_git_credential(token: str, username: str = "hf_user", folder: Optional[str] = None) -> None: + """Save a username/token pair in git credential for HF Hub registry. + + Credentials are saved in all configured helpers (store, cache, macOS keychain,...). + Calls "`git credential approve`" internally. See https://git-scm.com/docs/git-credential. + + Args: + username (`str`, defaults to `"hf_user"`): + A git username. Defaults to `"hf_user"`, the default user used in the Hub. + token (`str`, defaults to `"hf_user"`): + A git password. In practice, the User Access Token for the Hub. + See https://huggingface.co/settings/tokens. + folder (`str`, *optional*): + The folder in which to check the configured helpers. + """ + with run_interactive_subprocess("git credential approve", folder=folder) as ( + stdin, + _, + ): + stdin.write(f"url={ENDPOINT}\nusername={username.lower()}\npassword={token}\n\n") + stdin.flush() + + +def unset_git_credential(username: str = "hf_user", folder: Optional[str] = None) -> None: + """Erase credentials from git credential for HF Hub registry. + + Credentials are erased from the configured helpers (store, cache, macOS + keychain,...), if any. If `username` is not provided, any credential configured for + HF Hub endpoint is erased. + Calls "`git credential erase`" internally. See https://git-scm.com/docs/git-credential. + + Args: + username (`str`, defaults to `"hf_user"`): + A git username. Defaults to `"hf_user"`, the default user used in the Hub. + folder (`str`, *optional*): + The folder in which to check the configured helpers. + """ + with run_interactive_subprocess("git credential reject", folder=folder) as ( + stdin, + _, + ): + standard_input = f"url={ENDPOINT}\n" + if username is not None: + standard_input += f"username={username.lower()}\n" + standard_input += "\n" + + stdin.write(standard_input) + stdin.flush() + + +def _parse_credential_output(output: str) -> List[str]: + """Parse the output of `git credential fill` to extract the password. + + Args: + output (`str`): + The output of `git credential fill`. + """ + # NOTE: If user has set an helper for a custom URL, it will not we caught here. + # Example: `credential.https://huggingface.co.helper=store` + # See: https://github.com/huggingface/huggingface_hub/pull/1138#discussion_r1013324508 + return sorted( # Sort for nice printing + set( # Might have some duplicates + match[0] for match in GIT_CREDENTIAL_REGEX.findall(output) + ) + ) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_headers.py b/lib/python3.12/site-packages/huggingface_hub/utils/_headers.py new file mode 100644 index 0000000000000000000000000000000000000000..f0e1ddd87a51e170e7d0518c55cfaa3ab4f0ebb7 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_headers.py @@ -0,0 +1,228 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle headers to send in calls to Huggingface Hub.""" + +from typing import Dict, Optional, Union + +from huggingface_hub.errors import LocalTokenNotFoundError + +from .. import constants +from ._auth import get_token +from ._deprecation import _deprecate_arguments +from ._runtime import ( + get_fastai_version, + get_fastcore_version, + get_hf_hub_version, + get_python_version, + get_tf_version, + get_torch_version, + is_fastai_available, + is_fastcore_available, + is_tf_available, + is_torch_available, +) +from ._validators import validate_hf_hub_args + + +@_deprecate_arguments( + version="1.0", + deprecated_args="is_write_action", + custom_message="This argument is ignored and we let the server handle the permission error instead (if any).", +) +@validate_hf_hub_args +def build_hf_headers( + *, + token: Optional[Union[bool, str]] = None, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, + headers: Optional[Dict[str, str]] = None, + is_write_action: bool = False, +) -> Dict[str, str]: + """ + Build headers dictionary to send in a HF Hub call. + + By default, authorization token is always provided either from argument (explicit + use) or retrieved from the cache (implicit use). To explicitly avoid sending the + token to the Hub, set `token=False` or set the `HF_HUB_DISABLE_IMPLICIT_TOKEN` + environment variable. + + In case of an API call that requires write access, an error is thrown if token is + `None` or token is an organization token (starting with `"api_org***"`). + + In addition to the auth header, a user-agent is added to provide information about + the installed packages (versions of python, huggingface_hub, torch, tensorflow, + fastai and fastcore). + + Args: + token (`str`, `bool`, *optional*): + The token to be sent in authorization header for the Hub call: + - if a string, it is used as the Hugging Face token + - if `True`, the token is read from the machine (cache or env variable) + - if `False`, authorization header is not set + - if `None`, the token is read from the machine only except if + `HF_HUB_DISABLE_IMPLICIT_TOKEN` env variable is set. + library_name (`str`, *optional*): + The name of the library that is making the HTTP request. Will be added to + the user-agent header. + library_version (`str`, *optional*): + The version of the library that is making the HTTP request. Will be added + to the user-agent header. + user_agent (`str`, `dict`, *optional*): + The user agent info in the form of a dictionary or a single string. It will + be completed with information about the installed packages. + headers (`dict`, *optional*): + Additional headers to include in the request. Those headers take precedence + over the ones generated by this function. + is_write_action (`bool`): + Ignored and deprecated argument. + + Returns: + A `Dict` of headers to pass in your API call. + + Example: + ```py + >>> build_hf_headers(token="hf_***") # explicit token + {"authorization": "Bearer hf_***", "user-agent": ""} + + >>> build_hf_headers(token=True) # explicitly use cached token + {"authorization": "Bearer hf_***",...} + + >>> build_hf_headers(token=False) # explicitly don't use cached token + {"user-agent": ...} + + >>> build_hf_headers() # implicit use of the cached token + {"authorization": "Bearer hf_***",...} + + # HF_HUB_DISABLE_IMPLICIT_TOKEN=True # to set as env variable + >>> build_hf_headers() # token is not sent + {"user-agent": ...} + + >>> build_hf_headers(library_name="transformers", library_version="1.2.3") + {"authorization": ..., "user-agent": "transformers/1.2.3; hf_hub/0.10.2; python/3.10.4; tensorflow/1.55"} + ``` + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If organization token is passed and "write" access is required. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If "write" access is required but token is not passed and not saved locally. + [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) + If `token=True` but token is not saved locally. + """ + # Get auth token to send + token_to_send = get_token_to_send(token) + + # Combine headers + hf_headers = { + "user-agent": _http_user_agent( + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + ) + } + if token_to_send is not None: + hf_headers["authorization"] = f"Bearer {token_to_send}" + if headers is not None: + hf_headers.update(headers) + return hf_headers + + +def get_token_to_send(token: Optional[Union[bool, str]]) -> Optional[str]: + """Select the token to send from either `token` or the cache.""" + # Case token is explicitly provided + if isinstance(token, str): + return token + + # Case token is explicitly forbidden + if token is False: + return None + + # Token is not provided: we get it from local cache + cached_token = get_token() + + # Case token is explicitly required + if token is True: + if cached_token is None: + raise LocalTokenNotFoundError( + "Token is required (`token=True`), but no token found. You" + " need to provide a token or be logged in to Hugging Face with" + " `huggingface-cli login` or `huggingface_hub.login`. See" + " https://huggingface.co/settings/tokens." + ) + return cached_token + + # Case implicit use of the token is forbidden by env variable + if constants.HF_HUB_DISABLE_IMPLICIT_TOKEN: + return None + + # Otherwise: we use the cached token as the user has not explicitly forbidden it + return cached_token + + +def _http_user_agent( + *, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, +) -> str: + """Format a user-agent string containing information about the installed packages. + + Args: + library_name (`str`, *optional*): + The name of the library that is making the HTTP request. + library_version (`str`, *optional*): + The version of the library that is making the HTTP request. + user_agent (`str`, `dict`, *optional*): + The user agent info in the form of a dictionary or a single string. + + Returns: + The formatted user-agent string. + """ + if library_name is not None: + ua = f"{library_name}/{library_version}" + else: + ua = "unknown/None" + ua += f"; hf_hub/{get_hf_hub_version()}" + ua += f"; python/{get_python_version()}" + + if not constants.HF_HUB_DISABLE_TELEMETRY: + if is_torch_available(): + ua += f"; torch/{get_torch_version()}" + if is_tf_available(): + ua += f"; tensorflow/{get_tf_version()}" + if is_fastai_available(): + ua += f"; fastai/{get_fastai_version()}" + if is_fastcore_available(): + ua += f"; fastcore/{get_fastcore_version()}" + + if isinstance(user_agent, dict): + ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items()) + elif isinstance(user_agent, str): + ua += "; " + user_agent + + # Retrieve user-agent origin headers from environment variable + origin = constants.HF_HUB_USER_AGENT_ORIGIN + if origin is not None: + ua += "; origin/" + origin + + return _deduplicate_user_agent(ua) + + +def _deduplicate_user_agent(user_agent: str) -> str: + """Deduplicate redundant information in the generated user-agent.""" + # Split around ";" > Strip whitespaces > Store as dict keys (ensure unicity) > format back as string + # Order is implicitly preserved by dictionary structure (see https://stackoverflow.com/a/53657523). + return "; ".join({key.strip(): None for key in user_agent.split(";")}.keys()) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_hf_folder.py b/lib/python3.12/site-packages/huggingface_hub/utils/_hf_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..6418bf2fd2c59b4bcf301c1dd82bc468f2f42ddf --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_hf_folder.py @@ -0,0 +1,68 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# 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. +"""Contain helper class to retrieve/store token from/to local cache.""" + +from pathlib import Path +from typing import Optional + +from .. import constants +from ._auth import get_token + + +class HfFolder: + # TODO: deprecate when adapted in transformers/datasets/gradio + # @_deprecate_method(version="1.0", message="Use `huggingface_hub.login` instead.") + @classmethod + def save_token(cls, token: str) -> None: + """ + Save token, creating folder as needed. + + Token is saved in the huggingface home folder. You can configure it by setting + the `HF_HOME` environment variable. + + Args: + token (`str`): + The token to save to the [`HfFolder`] + """ + path_token = Path(constants.HF_TOKEN_PATH) + path_token.parent.mkdir(parents=True, exist_ok=True) + path_token.write_text(token) + + # TODO: deprecate when adapted in transformers/datasets/gradio + # @_deprecate_method(version="1.0", message="Use `huggingface_hub.get_token` instead.") + @classmethod + def get_token(cls) -> Optional[str]: + """ + Get token or None if not existent. + + This method is deprecated in favor of [`huggingface_hub.get_token`] but is kept for backward compatibility. + Its behavior is the same as [`huggingface_hub.get_token`]. + + Returns: + `str` or `None`: The token, `None` if it doesn't exist. + """ + return get_token() + + # TODO: deprecate when adapted in transformers/datasets/gradio + # @_deprecate_method(version="1.0", message="Use `huggingface_hub.logout` instead.") + @classmethod + def delete_token(cls) -> None: + """ + Deletes the token from storage. Does not fail if token does not exist. + """ + try: + Path(constants.HF_TOKEN_PATH).unlink() + except FileNotFoundError: + pass diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_http.py b/lib/python3.12/site-packages/huggingface_hub/utils/_http.py new file mode 100644 index 0000000000000000000000000000000000000000..5baceb8f8fd511403aa30c93dfe1fd33068c8dfe --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_http.py @@ -0,0 +1,637 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle HTTP requests in Huggingface Hub.""" + +import io +import os +import re +import threading +import time +import uuid +from functools import lru_cache +from http import HTTPStatus +from shlex import quote +from typing import Any, Callable, List, Optional, Tuple, Type, Union + +import requests +from requests import HTTPError, Response +from requests.adapters import HTTPAdapter +from requests.models import PreparedRequest + +from huggingface_hub.errors import OfflineModeIsEnabled + +from .. import constants +from ..errors import ( + BadRequestError, + DisabledRepoError, + EntryNotFoundError, + GatedRepoError, + HfHubHTTPError, + RepositoryNotFoundError, + RevisionNotFoundError, +) +from . import logging +from ._fixes import JSONDecodeError +from ._lfs import SliceFileObj +from ._typing import HTTP_METHOD_T + + +logger = logging.get_logger(__name__) + +# Both headers are used by the Hub to debug failed requests. +# `X_AMZN_TRACE_ID` is better as it also works to debug on Cloudfront and ALB. +# If `X_AMZN_TRACE_ID` is set, the Hub will use it as well. +X_AMZN_TRACE_ID = "X-Amzn-Trace-Id" +X_REQUEST_ID = "x-request-id" + +REPO_API_REGEX = re.compile( + r""" + # staging or production endpoint + ^https://[^/]+ + ( + # on /api/repo_type/repo_id + /api/(models|datasets|spaces)/(.+) + | + # or /repo_id/resolve/revision/... + /(.+)/resolve/(.+) + ) + """, + flags=re.VERBOSE, +) + + +class UniqueRequestIdAdapter(HTTPAdapter): + X_AMZN_TRACE_ID = "X-Amzn-Trace-Id" + + def add_headers(self, request, **kwargs): + super().add_headers(request, **kwargs) + + # Add random request ID => easier for server-side debug + if X_AMZN_TRACE_ID not in request.headers: + request.headers[X_AMZN_TRACE_ID] = request.headers.get(X_REQUEST_ID) or str(uuid.uuid4()) + + # Add debug log + has_token = len(str(request.headers.get("authorization", ""))) > 0 + logger.debug( + f"Request {request.headers[X_AMZN_TRACE_ID]}: {request.method} {request.url} (authenticated: {has_token})" + ) + + def send(self, request: PreparedRequest, *args, **kwargs) -> Response: + """Catch any RequestException to append request id to the error message for debugging.""" + if constants.HF_DEBUG: + logger.debug(f"Send: {_curlify(request)}") + try: + return super().send(request, *args, **kwargs) + except requests.RequestException as e: + request_id = request.headers.get(X_AMZN_TRACE_ID) + if request_id is not None: + # Taken from https://stackoverflow.com/a/58270258 + e.args = (*e.args, f"(Request ID: {request_id})") + raise + + +class OfflineAdapter(HTTPAdapter): + def send(self, request: PreparedRequest, *args, **kwargs) -> Response: + raise OfflineModeIsEnabled( + f"Cannot reach {request.url}: offline mode is enabled. To disable it, please unset the `HF_HUB_OFFLINE` environment variable." + ) + + +def _default_backend_factory() -> requests.Session: + session = requests.Session() + if constants.HF_HUB_OFFLINE: + session.mount("http://", OfflineAdapter()) + session.mount("https://", OfflineAdapter()) + else: + session.mount("http://", UniqueRequestIdAdapter()) + session.mount("https://", UniqueRequestIdAdapter()) + return session + + +BACKEND_FACTORY_T = Callable[[], requests.Session] +_GLOBAL_BACKEND_FACTORY: BACKEND_FACTORY_T = _default_backend_factory + + +def configure_http_backend(backend_factory: BACKEND_FACTORY_T = _default_backend_factory) -> None: + """ + Configure the HTTP backend by providing a `backend_factory`. Any HTTP calls made by `huggingface_hub` will use a + Session object instantiated by this factory. This can be useful if you are running your scripts in a specific + environment requiring custom configuration (e.g. custom proxy or certifications). + + Use [`get_session`] to get a configured Session. Since `requests.Session` is not guaranteed to be thread-safe, + `huggingface_hub` creates 1 Session instance per thread. They are all instantiated using the same `backend_factory` + set in [`configure_http_backend`]. A LRU cache is used to cache the created sessions (and connections) between + calls. Max size is 128 to avoid memory leaks if thousands of threads are spawned. + + See [this issue](https://github.com/psf/requests/issues/2766) to know more about thread-safety in `requests`. + + Example: + ```py + import requests + from huggingface_hub import configure_http_backend, get_session + + # Create a factory function that returns a Session with configured proxies + def backend_factory() -> requests.Session: + session = requests.Session() + session.proxies = {"http": "http://10.10.1.10:3128", "https": "https://10.10.1.11:1080"} + return session + + # Set it as the default session factory + configure_http_backend(backend_factory=backend_factory) + + # In practice, this is mostly done internally in `huggingface_hub` + session = get_session() + ``` + """ + global _GLOBAL_BACKEND_FACTORY + _GLOBAL_BACKEND_FACTORY = backend_factory + reset_sessions() + + +def get_session() -> requests.Session: + """ + Get a `requests.Session` object, using the session factory from the user. + + Use [`get_session`] to get a configured Session. Since `requests.Session` is not guaranteed to be thread-safe, + `huggingface_hub` creates 1 Session instance per thread. They are all instantiated using the same `backend_factory` + set in [`configure_http_backend`]. A LRU cache is used to cache the created sessions (and connections) between + calls. Max size is 128 to avoid memory leaks if thousands of threads are spawned. + + See [this issue](https://github.com/psf/requests/issues/2766) to know more about thread-safety in `requests`. + + Example: + ```py + import requests + from huggingface_hub import configure_http_backend, get_session + + # Create a factory function that returns a Session with configured proxies + def backend_factory() -> requests.Session: + session = requests.Session() + session.proxies = {"http": "http://10.10.1.10:3128", "https": "https://10.10.1.11:1080"} + return session + + # Set it as the default session factory + configure_http_backend(backend_factory=backend_factory) + + # In practice, this is mostly done internally in `huggingface_hub` + session = get_session() + ``` + """ + return _get_session_from_cache(process_id=os.getpid(), thread_id=threading.get_ident()) + + +def reset_sessions() -> None: + """Reset the cache of sessions. + + Mostly used internally when sessions are reconfigured or an SSLError is raised. + See [`configure_http_backend`] for more details. + """ + _get_session_from_cache.cache_clear() + + +@lru_cache +def _get_session_from_cache(process_id: int, thread_id: int) -> requests.Session: + """ + Create a new session per thread using global factory. Using LRU cache (maxsize 128) to avoid memory leaks when + using thousands of threads. Cache is cleared when `configure_http_backend` is called. + """ + return _GLOBAL_BACKEND_FACTORY() + + +def http_backoff( + method: HTTP_METHOD_T, + url: str, + *, + max_retries: int = 5, + base_wait_time: float = 1, + max_wait_time: float = 8, + retry_on_exceptions: Union[Type[Exception], Tuple[Type[Exception], ...]] = ( + requests.Timeout, + requests.ConnectionError, + ), + retry_on_status_codes: Union[int, Tuple[int, ...]] = HTTPStatus.SERVICE_UNAVAILABLE, + **kwargs, +) -> Response: + """Wrapper around requests to retry calls on an endpoint, with exponential backoff. + + Endpoint call is retried on exceptions (ex: connection timeout, proxy error,...) + and/or on specific status codes (ex: service unavailable). If the call failed more + than `max_retries`, the exception is thrown or `raise_for_status` is called on the + response object. + + Re-implement mechanisms from the `backoff` library to avoid adding an external + dependencies to `hugging_face_hub`. See https://github.com/litl/backoff. + + Args: + method (`Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]`): + HTTP method to perform. + url (`str`): + The URL of the resource to fetch. + max_retries (`int`, *optional*, defaults to `5`): + Maximum number of retries, defaults to 5 (no retries). + base_wait_time (`float`, *optional*, defaults to `1`): + Duration (in seconds) to wait before retrying the first time. + Wait time between retries then grows exponentially, capped by + `max_wait_time`. + max_wait_time (`float`, *optional*, defaults to `8`): + Maximum duration (in seconds) to wait before retrying. + retry_on_exceptions (`Type[Exception]` or `Tuple[Type[Exception]]`, *optional*): + Define which exceptions must be caught to retry the request. Can be a single type or a tuple of types. + By default, retry on `requests.Timeout` and `requests.ConnectionError`. + retry_on_status_codes (`int` or `Tuple[int]`, *optional*, defaults to `503`): + Define on which status codes the request must be retried. By default, only + HTTP 503 Service Unavailable is retried. + **kwargs (`dict`, *optional*): + kwargs to pass to `requests.request`. + + Example: + ``` + >>> from huggingface_hub.utils import http_backoff + + # Same usage as "requests.request". + >>> response = http_backoff("GET", "https://www.google.com") + >>> response.raise_for_status() + + # If you expect a Gateway Timeout from time to time + >>> http_backoff("PUT", upload_url, data=data, retry_on_status_codes=504) + >>> response.raise_for_status() + ``` + + + + When using `requests` it is possible to stream data by passing an iterator to the + `data` argument. On http backoff this is a problem as the iterator is not reset + after a failed call. This issue is mitigated for file objects or any IO streams + by saving the initial position of the cursor (with `data.tell()`) and resetting the + cursor between each call (with `data.seek()`). For arbitrary iterators, http backoff + will fail. If this is a hard constraint for you, please let us know by opening an + issue on [Github](https://github.com/huggingface/huggingface_hub). + + + """ + if isinstance(retry_on_exceptions, type): # Tuple from single exception type + retry_on_exceptions = (retry_on_exceptions,) + + if isinstance(retry_on_status_codes, int): # Tuple from single status code + retry_on_status_codes = (retry_on_status_codes,) + + nb_tries = 0 + sleep_time = base_wait_time + + # If `data` is used and is a file object (or any IO), it will be consumed on the + # first HTTP request. We need to save the initial position so that the full content + # of the file is re-sent on http backoff. See warning tip in docstring. + io_obj_initial_pos = None + if "data" in kwargs and isinstance(kwargs["data"], (io.IOBase, SliceFileObj)): + io_obj_initial_pos = kwargs["data"].tell() + + session = get_session() + while True: + nb_tries += 1 + try: + # If `data` is used and is a file object (or any IO), set back cursor to + # initial position. + if io_obj_initial_pos is not None: + kwargs["data"].seek(io_obj_initial_pos) + + # Perform request and return if status_code is not in the retry list. + response = session.request(method=method, url=url, **kwargs) + if response.status_code not in retry_on_status_codes: + return response + + # Wrong status code returned (HTTP 503 for instance) + logger.warning(f"HTTP Error {response.status_code} thrown while requesting {method} {url}") + if nb_tries > max_retries: + response.raise_for_status() # Will raise uncaught exception + # We return response to avoid infinite loop in the corner case where the + # user ask for retry on a status code that doesn't raise_for_status. + return response + + except retry_on_exceptions as err: + logger.warning(f"'{err}' thrown while requesting {method} {url}") + + if isinstance(err, requests.ConnectionError): + reset_sessions() # In case of SSLError it's best to reset the shared requests.Session objects + + if nb_tries > max_retries: + raise err + + # Sleep for X seconds + logger.warning(f"Retrying in {sleep_time}s [Retry {nb_tries}/{max_retries}].") + time.sleep(sleep_time) + + # Update sleep time for next retry + sleep_time = min(max_wait_time, sleep_time * 2) # Exponential backoff + + +def fix_hf_endpoint_in_url(url: str, endpoint: Optional[str]) -> str: + """Replace the default endpoint in a URL by a custom one. + + This is useful when using a proxy and the Hugging Face Hub returns a URL with the default endpoint. + """ + endpoint = endpoint.rstrip("/") if endpoint else constants.ENDPOINT + # check if a proxy has been set => if yes, update the returned URL to use the proxy + if endpoint not in (constants._HF_DEFAULT_ENDPOINT, constants._HF_DEFAULT_STAGING_ENDPOINT): + url = url.replace(constants._HF_DEFAULT_ENDPOINT, endpoint) + url = url.replace(constants._HF_DEFAULT_STAGING_ENDPOINT, endpoint) + return url + + +def hf_raise_for_status(response: Response, endpoint_name: Optional[str] = None) -> None: + """ + Internal version of `response.raise_for_status()` that will refine a + potential HTTPError. Raised exception will be an instance of `HfHubHTTPError`. + + This helper is meant to be the unique method to raise_for_status when making a call + to the Hugging Face Hub. + + + Example: + ```py + import requests + from huggingface_hub.utils import get_session, hf_raise_for_status, HfHubHTTPError + + response = get_session().post(...) + try: + hf_raise_for_status(response) + except HfHubHTTPError as e: + print(str(e)) # formatted message + e.request_id, e.server_message # details returned by server + + # Complete the error message with additional information once it's raised + e.append_to_message("\n`create_commit` expects the repository to exist.") + raise + ``` + + Args: + response (`Response`): + Response from the server. + endpoint_name (`str`, *optional*): + Name of the endpoint that has been called. If provided, the error message + will be more complete. + + + + Raises when the request has failed: + + - [`~utils.RepositoryNotFoundError`] + If the repository to download from cannot be found. This may be because it + doesn't exist, because `repo_type` is not set correctly, or because the repo + is `private` and you do not have access. + - [`~utils.GatedRepoError`] + If the repository exists but is gated and the user is not on the authorized + list. + - [`~utils.RevisionNotFoundError`] + If the repository exists but the revision couldn't be find. + - [`~utils.EntryNotFoundError`] + If the repository exists but the entry (e.g. the requested file) couldn't be + find. + - [`~utils.BadRequestError`] + If request failed with a HTTP 400 BadRequest error. + - [`~utils.HfHubHTTPError`] + If request failed for a reason not listed above. + + + """ + try: + response.raise_for_status() + except HTTPError as e: + error_code = response.headers.get("X-Error-Code") + error_message = response.headers.get("X-Error-Message") + + if error_code == "RevisionNotFound": + message = f"{response.status_code} Client Error." + "\n\n" + f"Revision Not Found for url: {response.url}." + raise _format(RevisionNotFoundError, message, response) from e + + elif error_code == "EntryNotFound": + message = f"{response.status_code} Client Error." + "\n\n" + f"Entry Not Found for url: {response.url}." + raise _format(EntryNotFoundError, message, response) from e + + elif error_code == "GatedRepo": + message = ( + f"{response.status_code} Client Error." + "\n\n" + f"Cannot access gated repo for url {response.url}." + ) + raise _format(GatedRepoError, message, response) from e + + elif error_message == "Access to this resource is disabled.": + message = ( + f"{response.status_code} Client Error." + + "\n\n" + + f"Cannot access repository for url {response.url}." + + "\n" + + "Access to this resource is disabled." + ) + raise _format(DisabledRepoError, message, response) from e + + elif error_code == "RepoNotFound" or ( + response.status_code == 401 + and error_message != "Invalid credentials in Authorization header" + and response.request is not None + and response.request.url is not None + and REPO_API_REGEX.search(response.request.url) is not None + ): + # 401 is misleading as it is returned for: + # - private and gated repos if user is not authenticated + # - missing repos + # => for now, we process them as `RepoNotFound` anyway. + # See https://gist.github.com/Wauplin/46c27ad266b15998ce56a6603796f0b9 + message = ( + f"{response.status_code} Client Error." + + "\n\n" + + f"Repository Not Found for url: {response.url}." + + "\nPlease make sure you specified the correct `repo_id` and" + " `repo_type`.\nIf you are trying to access a private or gated repo," + " make sure you are authenticated. For more details, see" + " https://huggingface.co/docs/huggingface_hub/authentication" + ) + raise _format(RepositoryNotFoundError, message, response) from e + + elif response.status_code == 400: + message = ( + f"\n\nBad request for {endpoint_name} endpoint:" if endpoint_name is not None else "\n\nBad request:" + ) + raise _format(BadRequestError, message, response) from e + + elif response.status_code == 403: + message = ( + f"\n\n{response.status_code} Forbidden: {error_message}." + + f"\nCannot access content at: {response.url}." + + "\nMake sure your token has the correct permissions." + ) + raise _format(HfHubHTTPError, message, response) from e + + elif response.status_code == 416: + range_header = response.request.headers.get("Range") + message = f"{e}. Requested range: {range_header}. Content-Range: {response.headers.get('Content-Range')}." + raise _format(HfHubHTTPError, message, response) from e + + # Convert `HTTPError` into a `HfHubHTTPError` to display request information + # as well (request id and/or server error message) + raise _format(HfHubHTTPError, str(e), response) from e + + +def _format(error_type: Type[HfHubHTTPError], custom_message: str, response: Response) -> HfHubHTTPError: + server_errors = [] + + # Retrieve server error from header + from_headers = response.headers.get("X-Error-Message") + if from_headers is not None: + server_errors.append(from_headers) + + # Retrieve server error from body + try: + # Case errors are returned in a JSON format + data = response.json() + + error = data.get("error") + if error is not None: + if isinstance(error, list): + # Case {'error': ['my error 1', 'my error 2']} + server_errors.extend(error) + else: + # Case {'error': 'my error'} + server_errors.append(error) + + errors = data.get("errors") + if errors is not None: + # Case {'errors': [{'message': 'my error 1'}, {'message': 'my error 2'}]} + for error in errors: + if "message" in error: + server_errors.append(error["message"]) + + except JSONDecodeError: + # If content is not JSON and not HTML, append the text + content_type = response.headers.get("Content-Type", "") + if response.text and "html" not in content_type.lower(): + server_errors.append(response.text) + + # Strip all server messages + server_errors = [str(line).strip() for line in server_errors if str(line).strip()] + + # Deduplicate server messages (keep order) + # taken from https://stackoverflow.com/a/17016257 + server_errors = list(dict.fromkeys(server_errors)) + + # Format server error + server_message = "\n".join(server_errors) + + # Add server error to custom message + final_error_message = custom_message + if server_message and server_message.lower() not in custom_message.lower(): + if "\n\n" in custom_message: + final_error_message += "\n" + server_message + else: + final_error_message += "\n\n" + server_message + # Add Request ID + request_id = str(response.headers.get(X_REQUEST_ID, "")) + if request_id: + request_id_message = f" (Request ID: {request_id})" + else: + # Fallback to X-Amzn-Trace-Id + request_id = str(response.headers.get(X_AMZN_TRACE_ID, "")) + if request_id: + request_id_message = f" (Amzn Trace ID: {request_id})" + if request_id and request_id.lower() not in final_error_message.lower(): + if "\n" in final_error_message: + newline_index = final_error_message.index("\n") + final_error_message = ( + final_error_message[:newline_index] + request_id_message + final_error_message[newline_index:] + ) + else: + final_error_message += request_id_message + + # Return + return error_type(final_error_message.strip(), response=response, server_message=server_message or None) + + +def _curlify(request: requests.PreparedRequest) -> str: + """Convert a `requests.PreparedRequest` into a curl command (str). + + Used for debug purposes only. + + Implementation vendored from https://github.com/ofw/curlify/blob/master/curlify.py. + MIT License Copyright (c) 2016 Egor. + """ + parts: List[Tuple[Any, Any]] = [ + ("curl", None), + ("-X", request.method), + ] + + for k, v in sorted(request.headers.items()): + if k.lower() == "authorization": + v = "" # Hide authorization header, no matter its value (can be Bearer, Key, etc.) + parts += [("-H", "{0}: {1}".format(k, v))] + + if request.body: + body = request.body + if isinstance(body, bytes): + body = body.decode("utf-8", errors="ignore") + elif hasattr(body, "read"): + body = "" # Don't try to read it to avoid consuming the stream + if len(body) > 1000: + body = body[:1000] + " ... [truncated]" + parts += [("-d", body.replace("\n", ""))] + + parts += [(None, request.url)] + + flat_parts = [] + for k, v in parts: + if k: + flat_parts.append(quote(k)) + if v: + flat_parts.append(quote(v)) + + return " ".join(flat_parts) + + +# Regex to parse HTTP Range header +RANGE_REGEX = re.compile(r"^\s*bytes\s*=\s*(\d*)\s*-\s*(\d*)\s*$", re.IGNORECASE) + + +def _adjust_range_header(original_range: Optional[str], resume_size: int) -> Optional[str]: + """ + Adjust HTTP Range header to account for resume position. + """ + if not original_range: + return f"bytes={resume_size}-" + + if "," in original_range: + raise ValueError(f"Multiple ranges detected - {original_range!r}, not supported yet.") + + match = RANGE_REGEX.match(original_range) + if not match: + raise RuntimeError(f"Invalid range format - {original_range!r}.") + start, end = match.groups() + + if not start: + if not end: + raise RuntimeError(f"Invalid range format - {original_range!r}.") + + new_suffix = int(end) - resume_size + new_range = f"bytes=-{new_suffix}" + if new_suffix <= 0: + raise RuntimeError(f"Empty new range - {new_range!r}.") + return new_range + + start = int(start) + new_start = start + resume_size + if end: + end = int(end) + new_range = f"bytes={new_start}-{end}" + if new_start > end: + raise RuntimeError(f"Empty new range - {new_range!r}.") + return new_range + + return f"bytes={new_start}-" diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_lfs.py b/lib/python3.12/site-packages/huggingface_hub/utils/_lfs.py new file mode 100644 index 0000000000000000000000000000000000000000..307f371ffa79a8ae726ee03458c52e230a792898 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_lfs.py @@ -0,0 +1,110 @@ +# coding=utf-8 +# Copyright 2019-present, the HuggingFace Inc. team. +# +# 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. +"""Git LFS related utilities""" + +import io +import os +from contextlib import AbstractContextManager +from typing import BinaryIO + + +class SliceFileObj(AbstractContextManager): + """ + Utility context manager to read a *slice* of a seekable file-like object as a seekable, file-like object. + + This is NOT thread safe + + Inspired by stackoverflow.com/a/29838711/593036 + + Credits to @julien-c + + Args: + fileobj (`BinaryIO`): + A file-like object to slice. MUST implement `tell()` and `seek()` (and `read()` of course). + `fileobj` will be reset to its original position when exiting the context manager. + seek_from (`int`): + The start of the slice (offset from position 0 in bytes). + read_limit (`int`): + The maximum number of bytes to read from the slice. + + Attributes: + previous_position (`int`): + The previous position + + Examples: + + Reading 200 bytes with an offset of 128 bytes from a file (ie bytes 128 to 327): + ```python + >>> with open("path/to/file", "rb") as file: + ... with SliceFileObj(file, seek_from=128, read_limit=200) as fslice: + ... fslice.read(...) + ``` + + Reading a file in chunks of 512 bytes + ```python + >>> import os + >>> chunk_size = 512 + >>> file_size = os.getsize("path/to/file") + >>> with open("path/to/file", "rb") as file: + ... for chunk_idx in range(ceil(file_size / chunk_size)): + ... with SliceFileObj(file, seek_from=chunk_idx * chunk_size, read_limit=chunk_size) as fslice: + ... chunk = fslice.read(...) + + ``` + """ + + def __init__(self, fileobj: BinaryIO, seek_from: int, read_limit: int): + self.fileobj = fileobj + self.seek_from = seek_from + self.read_limit = read_limit + + def __enter__(self): + self._previous_position = self.fileobj.tell() + end_of_stream = self.fileobj.seek(0, os.SEEK_END) + self._len = min(self.read_limit, end_of_stream - self.seek_from) + # ^^ The actual number of bytes that can be read from the slice + self.fileobj.seek(self.seek_from, io.SEEK_SET) + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.fileobj.seek(self._previous_position, io.SEEK_SET) + + def read(self, n: int = -1): + pos = self.tell() + if pos >= self._len: + return b"" + remaining_amount = self._len - pos + data = self.fileobj.read(remaining_amount if n < 0 else min(n, remaining_amount)) + return data + + def tell(self) -> int: + return self.fileobj.tell() - self.seek_from + + def seek(self, offset: int, whence: int = os.SEEK_SET) -> int: + start = self.seek_from + end = start + self._len + if whence in (os.SEEK_SET, os.SEEK_END): + offset = start + offset if whence == os.SEEK_SET else end + offset + offset = max(start, min(offset, end)) + whence = os.SEEK_SET + elif whence == os.SEEK_CUR: + cur_pos = self.fileobj.tell() + offset = max(start - cur_pos, min(offset, end - cur_pos)) + else: + raise ValueError(f"whence value {whence} is not supported") + return self.fileobj.seek(offset, whence) - self.seek_from + + def __iter__(self): + yield self.read(n=4 * 1024 * 1024) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_pagination.py b/lib/python3.12/site-packages/huggingface_hub/utils/_pagination.py new file mode 100644 index 0000000000000000000000000000000000000000..3ef2b6668ba09d4c6a715509131d157139a1fac0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_pagination.py @@ -0,0 +1,52 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle pagination on Huggingface Hub.""" + +from typing import Dict, Iterable, Optional + +import requests + +from . import get_session, hf_raise_for_status, http_backoff, logging + + +logger = logging.get_logger(__name__) + + +def paginate(path: str, params: Dict, headers: Dict) -> Iterable: + """Fetch a list of models/datasets/spaces and paginate through results. + + This is using the same "Link" header format as GitHub. + See: + - https://requests.readthedocs.io/en/latest/api/#requests.Response.links + - https://docs.github.com/en/rest/guides/traversing-with-pagination#link-header + """ + session = get_session() + r = session.get(path, params=params, headers=headers) + hf_raise_for_status(r) + yield from r.json() + + # Follow pages + # Next link already contains query params + next_page = _get_next_page(r) + while next_page is not None: + logger.debug(f"Pagination detected. Requesting next page: {next_page}") + r = http_backoff("GET", next_page, max_retries=20, retry_on_status_codes=429, headers=headers) + hf_raise_for_status(r) + yield from r.json() + next_page = _get_next_page(r) + + +def _get_next_page(response: requests.Response) -> Optional[str]: + return response.links.get("next", {}).get("url") diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_paths.py b/lib/python3.12/site-packages/huggingface_hub/utils/_paths.py new file mode 100644 index 0000000000000000000000000000000000000000..4f2c0ebce070bbde4900e919a3aca7cfc331e747 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_paths.py @@ -0,0 +1,141 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to handle paths in Huggingface Hub.""" + +from fnmatch import fnmatch +from pathlib import Path +from typing import Callable, Generator, Iterable, List, Optional, TypeVar, Union + + +T = TypeVar("T") + +# Always ignore `.git` and `.cache/huggingface` folders in commits +DEFAULT_IGNORE_PATTERNS = [ + ".git", + ".git/*", + "*/.git", + "**/.git/**", + ".cache/huggingface", + ".cache/huggingface/*", + "*/.cache/huggingface", + "**/.cache/huggingface/**", +] +# Forbidden to commit these folders +FORBIDDEN_FOLDERS = [".git", ".cache"] + + +def filter_repo_objects( + items: Iterable[T], + *, + allow_patterns: Optional[Union[List[str], str]] = None, + ignore_patterns: Optional[Union[List[str], str]] = None, + key: Optional[Callable[[T], str]] = None, +) -> Generator[T, None, None]: + """Filter repo objects based on an allowlist and a denylist. + + Input must be a list of paths (`str` or `Path`) or a list of arbitrary objects. + In the later case, `key` must be provided and specifies a function of one argument + that is used to extract a path from each element in iterable. + + Patterns are Unix shell-style wildcards which are NOT regular expressions. See + https://docs.python.org/3/library/fnmatch.html for more details. + + Args: + items (`Iterable`): + List of items to filter. + allow_patterns (`str` or `List[str]`, *optional*): + Patterns constituting the allowlist. If provided, item paths must match at + least one pattern from the allowlist. + ignore_patterns (`str` or `List[str]`, *optional*): + Patterns constituting the denylist. If provided, item paths must not match + any patterns from the denylist. + key (`Callable[[T], str]`, *optional*): + Single-argument function to extract a path from each item. If not provided, + the `items` must already be `str` or `Path`. + + Returns: + Filtered list of objects, as a generator. + + Raises: + :class:`ValueError`: + If `key` is not provided and items are not `str` or `Path`. + + Example usage with paths: + ```python + >>> # Filter only PDFs that are not hidden. + >>> list(filter_repo_objects( + ... ["aaa.PDF", "bbb.jpg", ".ccc.pdf", ".ddd.png"], + ... allow_patterns=["*.pdf"], + ... ignore_patterns=[".*"], + ... )) + ["aaa.pdf"] + ``` + + Example usage with objects: + ```python + >>> list(filter_repo_objects( + ... [ + ... CommitOperationAdd(path_or_fileobj="/tmp/aaa.pdf", path_in_repo="aaa.pdf") + ... CommitOperationAdd(path_or_fileobj="/tmp/bbb.jpg", path_in_repo="bbb.jpg") + ... CommitOperationAdd(path_or_fileobj="/tmp/.ccc.pdf", path_in_repo=".ccc.pdf") + ... CommitOperationAdd(path_or_fileobj="/tmp/.ddd.png", path_in_repo=".ddd.png") + ... ], + ... allow_patterns=["*.pdf"], + ... ignore_patterns=[".*"], + ... key=lambda x: x.repo_in_path + ... )) + [CommitOperationAdd(path_or_fileobj="/tmp/aaa.pdf", path_in_repo="aaa.pdf")] + ``` + """ + if isinstance(allow_patterns, str): + allow_patterns = [allow_patterns] + + if isinstance(ignore_patterns, str): + ignore_patterns = [ignore_patterns] + + if allow_patterns is not None: + allow_patterns = [_add_wildcard_to_directories(p) for p in allow_patterns] + if ignore_patterns is not None: + ignore_patterns = [_add_wildcard_to_directories(p) for p in ignore_patterns] + + if key is None: + + def _identity(item: T) -> str: + if isinstance(item, str): + return item + if isinstance(item, Path): + return str(item) + raise ValueError(f"Please provide `key` argument in `filter_repo_objects`: `{item}` is not a string.") + + key = _identity # Items must be `str` or `Path`, otherwise raise ValueError + + for item in items: + path = key(item) + + # Skip if there's an allowlist and path doesn't match any + if allow_patterns is not None and not any(fnmatch(path, r) for r in allow_patterns): + continue + + # Skip if there's a denylist and path matches any + if ignore_patterns is not None and any(fnmatch(path, r) for r in ignore_patterns): + continue + + yield item + + +def _add_wildcard_to_directories(pattern: str) -> str: + if pattern[-1] == "/": + return pattern + "*" + return pattern diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_runtime.py b/lib/python3.12/site-packages/huggingface_hub/utils/_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..1ffd25b3d86e16975c7ba3f3d7ca2892e0a936e4 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_runtime.py @@ -0,0 +1,394 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# 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. +"""Check presence of installed packages at runtime.""" + +import importlib.metadata +import os +import platform +import sys +import warnings +from typing import Any, Dict + +from .. import __version__, constants + + +_PY_VERSION: str = sys.version.split()[0].rstrip("+") + +_package_versions = {} + +_CANDIDATES = { + "aiohttp": {"aiohttp"}, + "fastai": {"fastai"}, + "fastapi": {"fastapi"}, + "fastcore": {"fastcore"}, + "gradio": {"gradio"}, + "graphviz": {"graphviz"}, + "hf_transfer": {"hf_transfer"}, + "hf_xet": {"hf_xet"}, + "jinja": {"Jinja2"}, + "keras": {"keras"}, + "numpy": {"numpy"}, + "pillow": {"Pillow"}, + "pydantic": {"pydantic"}, + "pydot": {"pydot"}, + "safetensors": {"safetensors"}, + "tensorboard": {"tensorboardX"}, + "tensorflow": ( + "tensorflow", + "tensorflow-cpu", + "tensorflow-gpu", + "tf-nightly", + "tf-nightly-cpu", + "tf-nightly-gpu", + "intel-tensorflow", + "intel-tensorflow-avx512", + "tensorflow-rocm", + "tensorflow-macos", + ), + "torch": {"torch"}, +} + +# Check once at runtime +for candidate_name, package_names in _CANDIDATES.items(): + _package_versions[candidate_name] = "N/A" + for name in package_names: + try: + _package_versions[candidate_name] = importlib.metadata.version(name) + break + except importlib.metadata.PackageNotFoundError: + pass + + +def _get_version(package_name: str) -> str: + return _package_versions.get(package_name, "N/A") + + +def is_package_available(package_name: str) -> bool: + return _get_version(package_name) != "N/A" + + +# Python +def get_python_version() -> str: + return _PY_VERSION + + +# Huggingface Hub +def get_hf_hub_version() -> str: + return __version__ + + +# aiohttp +def is_aiohttp_available() -> bool: + return is_package_available("aiohttp") + + +def get_aiohttp_version() -> str: + return _get_version("aiohttp") + + +# FastAI +def is_fastai_available() -> bool: + return is_package_available("fastai") + + +def get_fastai_version() -> str: + return _get_version("fastai") + + +# FastAPI +def is_fastapi_available() -> bool: + return is_package_available("fastapi") + + +def get_fastapi_version() -> str: + return _get_version("fastapi") + + +# Fastcore +def is_fastcore_available() -> bool: + return is_package_available("fastcore") + + +def get_fastcore_version() -> str: + return _get_version("fastcore") + + +# FastAI +def is_gradio_available() -> bool: + return is_package_available("gradio") + + +def get_gradio_version() -> str: + return _get_version("gradio") + + +# Graphviz +def is_graphviz_available() -> bool: + return is_package_available("graphviz") + + +def get_graphviz_version() -> str: + return _get_version("graphviz") + + +# hf_transfer +def is_hf_transfer_available() -> bool: + return is_package_available("hf_transfer") + + +def get_hf_transfer_version() -> str: + return _get_version("hf_transfer") + + +# xet +def is_xet_available() -> bool: + # since hf_xet is automatically used if available, allow explicit disabling via environment variable + if constants._is_true(os.environ.get("HF_HUB_DISABLE_XET")): # type: ignore + return False + + return is_package_available("hf_xet") + + +def get_xet_version() -> str: + return _get_version("hf_xet") + + +# keras +def is_keras_available() -> bool: + return is_package_available("keras") + + +def get_keras_version() -> str: + return _get_version("keras") + + +# Numpy +def is_numpy_available() -> bool: + return is_package_available("numpy") + + +def get_numpy_version() -> str: + return _get_version("numpy") + + +# Jinja +def is_jinja_available() -> bool: + return is_package_available("jinja") + + +def get_jinja_version() -> str: + return _get_version("jinja") + + +# Pillow +def is_pillow_available() -> bool: + return is_package_available("pillow") + + +def get_pillow_version() -> str: + return _get_version("pillow") + + +# Pydantic +def is_pydantic_available() -> bool: + if not is_package_available("pydantic"): + return False + # For Pydantic, we add an extra check to test whether it is correctly installed or not. If both pydantic 2.x and + # typing_extensions<=4.5.0 are installed, then pydantic will fail at import time. This should not happen when + # it is installed with `pip install huggingface_hub[inference]` but it can happen when it is installed manually + # by the user in an environment that we don't control. + # + # Usually we won't need to do this kind of check on optional dependencies. However, pydantic is a special case + # as it is automatically imported when doing `from huggingface_hub import ...` even if the user doesn't use it. + # + # See https://github.com/huggingface/huggingface_hub/pull/1829 for more details. + try: + from pydantic import validator # noqa: F401 + except ImportError: + # Example: "ImportError: cannot import name 'TypeAliasType' from 'typing_extensions'" + warnings.warn( + "Pydantic is installed but cannot be imported. Please check your installation. `huggingface_hub` will " + "default to not using Pydantic. Error message: '{e}'" + ) + return False + return True + + +def get_pydantic_version() -> str: + return _get_version("pydantic") + + +# Pydot +def is_pydot_available() -> bool: + return is_package_available("pydot") + + +def get_pydot_version() -> str: + return _get_version("pydot") + + +# Tensorboard +def is_tensorboard_available() -> bool: + return is_package_available("tensorboard") + + +def get_tensorboard_version() -> str: + return _get_version("tensorboard") + + +# Tensorflow +def is_tf_available() -> bool: + return is_package_available("tensorflow") + + +def get_tf_version() -> str: + return _get_version("tensorflow") + + +# Torch +def is_torch_available() -> bool: + return is_package_available("torch") + + +def get_torch_version() -> str: + return _get_version("torch") + + +# Safetensors +def is_safetensors_available() -> bool: + return is_package_available("safetensors") + + +# Shell-related helpers +try: + # Set to `True` if script is running in a Google Colab notebook. + # If running in Google Colab, git credential store is set globally which makes the + # warning disappear. See https://github.com/huggingface/huggingface_hub/issues/1043 + # + # Taken from https://stackoverflow.com/a/63519730. + _is_google_colab = "google.colab" in str(get_ipython()) # type: ignore # noqa: F821 +except NameError: + _is_google_colab = False + + +def is_notebook() -> bool: + """Return `True` if code is executed in a notebook (Jupyter, Colab, QTconsole). + + Taken from https://stackoverflow.com/a/39662359. + Adapted to make it work with Google colab as well. + """ + try: + shell_class = get_ipython().__class__ # type: ignore # noqa: F821 + for parent_class in shell_class.__mro__: # e.g. "is subclass of" + if parent_class.__name__ == "ZMQInteractiveShell": + return True # Jupyter notebook, Google colab or qtconsole + return False + except NameError: + return False # Probably standard Python interpreter + + +def is_google_colab() -> bool: + """Return `True` if code is executed in a Google colab. + + Taken from https://stackoverflow.com/a/63519730. + """ + return _is_google_colab + + +def is_colab_enterprise() -> bool: + """Return `True` if code is executed in a Google Colab Enterprise environment.""" + return os.environ.get("VERTEX_PRODUCT") == "COLAB_ENTERPRISE" + + +def dump_environment_info() -> Dict[str, Any]: + """Dump information about the machine to help debugging issues. + + Similar helper exist in: + - `datasets` (https://github.com/huggingface/datasets/blob/main/src/datasets/commands/env.py) + - `diffusers` (https://github.com/huggingface/diffusers/blob/main/src/diffusers/commands/env.py) + - `transformers` (https://github.com/huggingface/transformers/blob/main/src/transformers/commands/env.py) + """ + from huggingface_hub import get_token, whoami + from huggingface_hub.utils import list_credential_helpers + + token = get_token() + + # Generic machine info + info: Dict[str, Any] = { + "huggingface_hub version": get_hf_hub_version(), + "Platform": platform.platform(), + "Python version": get_python_version(), + } + + # Interpreter info + try: + shell_class = get_ipython().__class__ # type: ignore # noqa: F821 + info["Running in iPython ?"] = "Yes" + info["iPython shell"] = shell_class.__name__ + except NameError: + info["Running in iPython ?"] = "No" + info["Running in notebook ?"] = "Yes" if is_notebook() else "No" + info["Running in Google Colab ?"] = "Yes" if is_google_colab() else "No" + info["Running in Google Colab Enterprise ?"] = "Yes" if is_colab_enterprise() else "No" + # Login info + info["Token path ?"] = constants.HF_TOKEN_PATH + info["Has saved token ?"] = token is not None + if token is not None: + try: + info["Who am I ?"] = whoami()["name"] + except Exception: + pass + + try: + info["Configured git credential helpers"] = ", ".join(list_credential_helpers()) + except Exception: + pass + + # Installed dependencies + info["FastAI"] = get_fastai_version() + info["Tensorflow"] = get_tf_version() + info["Torch"] = get_torch_version() + info["Jinja2"] = get_jinja_version() + info["Graphviz"] = get_graphviz_version() + info["keras"] = get_keras_version() + info["Pydot"] = get_pydot_version() + info["Pillow"] = get_pillow_version() + info["hf_transfer"] = get_hf_transfer_version() + info["gradio"] = get_gradio_version() + info["tensorboard"] = get_tensorboard_version() + info["numpy"] = get_numpy_version() + info["pydantic"] = get_pydantic_version() + info["aiohttp"] = get_aiohttp_version() + info["hf_xet"] = get_xet_version() + + # Environment variables + info["ENDPOINT"] = constants.ENDPOINT + info["HF_HUB_CACHE"] = constants.HF_HUB_CACHE + info["HF_ASSETS_CACHE"] = constants.HF_ASSETS_CACHE + info["HF_TOKEN_PATH"] = constants.HF_TOKEN_PATH + info["HF_STORED_TOKENS_PATH"] = constants.HF_STORED_TOKENS_PATH + info["HF_HUB_OFFLINE"] = constants.HF_HUB_OFFLINE + info["HF_HUB_DISABLE_TELEMETRY"] = constants.HF_HUB_DISABLE_TELEMETRY + info["HF_HUB_DISABLE_PROGRESS_BARS"] = constants.HF_HUB_DISABLE_PROGRESS_BARS + info["HF_HUB_DISABLE_SYMLINKS_WARNING"] = constants.HF_HUB_DISABLE_SYMLINKS_WARNING + info["HF_HUB_DISABLE_EXPERIMENTAL_WARNING"] = constants.HF_HUB_DISABLE_EXPERIMENTAL_WARNING + info["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = constants.HF_HUB_DISABLE_IMPLICIT_TOKEN + info["HF_HUB_ENABLE_HF_TRANSFER"] = constants.HF_HUB_ENABLE_HF_TRANSFER + info["HF_HUB_ETAG_TIMEOUT"] = constants.HF_HUB_ETAG_TIMEOUT + info["HF_HUB_DOWNLOAD_TIMEOUT"] = constants.HF_HUB_DOWNLOAD_TIMEOUT + + print("\nCopy-and-paste the text below in your GitHub issue.\n") + print("\n".join([f"- {prop}: {val}" for prop, val in info.items()]) + "\n") + return info diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_safetensors.py b/lib/python3.12/site-packages/huggingface_hub/utils/_safetensors.py new file mode 100644 index 0000000000000000000000000000000000000000..38546c6d34db786c62861e1706f747a21b7012bf --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_safetensors.py @@ -0,0 +1,111 @@ +import functools +import operator +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Dict, List, Literal, Optional, Tuple + + +FILENAME_T = str +TENSOR_NAME_T = str +DTYPE_T = Literal["F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8", "BOOL"] + + +@dataclass +class TensorInfo: + """Information about a tensor. + + For more details regarding the safetensors format, check out https://huggingface.co/docs/safetensors/index#format. + + Attributes: + dtype (`str`): + The data type of the tensor ("F64", "F32", "F16", "BF16", "I64", "I32", "I16", "I8", "U8", "BOOL"). + shape (`List[int]`): + The shape of the tensor. + data_offsets (`Tuple[int, int]`): + The offsets of the data in the file as a tuple `[BEGIN, END]`. + parameter_count (`int`): + The number of parameters in the tensor. + """ + + dtype: DTYPE_T + shape: List[int] + data_offsets: Tuple[int, int] + parameter_count: int = field(init=False) + + def __post_init__(self) -> None: + # Taken from https://stackoverflow.com/a/13840436 + try: + self.parameter_count = functools.reduce(operator.mul, self.shape) + except TypeError: + self.parameter_count = 1 # scalar value has no shape + + +@dataclass +class SafetensorsFileMetadata: + """Metadata for a Safetensors file hosted on the Hub. + + This class is returned by [`parse_safetensors_file_metadata`]. + + For more details regarding the safetensors format, check out https://huggingface.co/docs/safetensors/index#format. + + Attributes: + metadata (`Dict`): + The metadata contained in the file. + tensors (`Dict[str, TensorInfo]`): + A map of all tensors. Keys are tensor names and values are information about the corresponding tensor, as a + [`TensorInfo`] object. + parameter_count (`Dict[str, int]`): + A map of the number of parameters per data type. Keys are data types and values are the number of parameters + of that data type. + """ + + metadata: Dict[str, str] + tensors: Dict[TENSOR_NAME_T, TensorInfo] + parameter_count: Dict[DTYPE_T, int] = field(init=False) + + def __post_init__(self) -> None: + parameter_count: Dict[DTYPE_T, int] = defaultdict(int) + for tensor in self.tensors.values(): + parameter_count[tensor.dtype] += tensor.parameter_count + self.parameter_count = dict(parameter_count) + + +@dataclass +class SafetensorsRepoMetadata: + """Metadata for a Safetensors repo. + + A repo is considered to be a Safetensors repo if it contains either a 'model.safetensors' weight file (non-shared + model) or a 'model.safetensors.index.json' index file (sharded model) at its root. + + This class is returned by [`get_safetensors_metadata`]. + + For more details regarding the safetensors format, check out https://huggingface.co/docs/safetensors/index#format. + + Attributes: + metadata (`Dict`, *optional*): + The metadata contained in the 'model.safetensors.index.json' file, if it exists. Only populated for sharded + models. + sharded (`bool`): + Whether the repo contains a sharded model or not. + weight_map (`Dict[str, str]`): + A map of all weights. Keys are tensor names and values are filenames of the files containing the tensors. + files_metadata (`Dict[str, SafetensorsFileMetadata]`): + A map of all files metadata. Keys are filenames and values are the metadata of the corresponding file, as + a [`SafetensorsFileMetadata`] object. + parameter_count (`Dict[str, int]`): + A map of the number of parameters per data type. Keys are data types and values are the number of parameters + of that data type. + """ + + metadata: Optional[Dict] + sharded: bool + weight_map: Dict[TENSOR_NAME_T, FILENAME_T] # tensor name -> filename + files_metadata: Dict[FILENAME_T, SafetensorsFileMetadata] # filename -> metadata + parameter_count: Dict[DTYPE_T, int] = field(init=False) + + def __post_init__(self) -> None: + parameter_count: Dict[DTYPE_T, int] = defaultdict(int) + for file_metadata in self.files_metadata.values(): + for dtype, nb_parameters_ in file_metadata.parameter_count.items(): + parameter_count[dtype] += nb_parameters_ + self.parameter_count = dict(parameter_count) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_subprocess.py b/lib/python3.12/site-packages/huggingface_hub/utils/_subprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..fdabf1c4df3b61dc610ae08eb7842df6af3552f3 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_subprocess.py @@ -0,0 +1,144 @@ +# coding=utf-8 +# Copyright 2021 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 +"""Contains utilities to easily handle subprocesses in `huggingface_hub`.""" + +import os +import subprocess +import sys +from contextlib import contextmanager +from io import StringIO +from pathlib import Path +from typing import IO, Generator, List, Optional, Tuple, Union + +from .logging import get_logger + + +logger = get_logger(__name__) + + +@contextmanager +def capture_output() -> Generator[StringIO, None, None]: + """Capture output that is printed to terminal. + + Taken from https://stackoverflow.com/a/34738440 + + Example: + ```py + >>> with capture_output() as output: + ... print("hello world") + >>> assert output.getvalue() == "hello world\n" + ``` + """ + output = StringIO() + previous_output = sys.stdout + sys.stdout = output + try: + yield output + finally: + sys.stdout = previous_output + + +def run_subprocess( + command: Union[str, List[str]], + folder: Optional[Union[str, Path]] = None, + check=True, + **kwargs, +) -> subprocess.CompletedProcess: + """ + Method to run subprocesses. Calling this will capture the `stderr` and `stdout`, + please call `subprocess.run` manually in case you would like for them not to + be captured. + + Args: + command (`str` or `List[str]`): + The command to execute as a string or list of strings. + folder (`str`, *optional*): + The folder in which to run the command. Defaults to current working + directory (from `os.getcwd()`). + check (`bool`, *optional*, defaults to `True`): + Setting `check` to `True` will raise a `subprocess.CalledProcessError` + when the subprocess has a non-zero exit code. + kwargs (`Dict[str]`): + Keyword arguments to be passed to the `subprocess.run` underlying command. + + Returns: + `subprocess.CompletedProcess`: The completed process. + """ + if isinstance(command, str): + command = command.split() + + if isinstance(folder, Path): + folder = str(folder) + + return subprocess.run( + command, + stderr=subprocess.PIPE, + stdout=subprocess.PIPE, + check=check, + encoding="utf-8", + errors="replace", # if not utf-8, replace char by � + cwd=folder or os.getcwd(), + **kwargs, + ) + + +@contextmanager +def run_interactive_subprocess( + command: Union[str, List[str]], + folder: Optional[Union[str, Path]] = None, + **kwargs, +) -> Generator[Tuple[IO[str], IO[str]], None, None]: + """Run a subprocess in an interactive mode in a context manager. + + Args: + command (`str` or `List[str]`): + The command to execute as a string or list of strings. + folder (`str`, *optional*): + The folder in which to run the command. Defaults to current working + directory (from `os.getcwd()`). + kwargs (`Dict[str]`): + Keyword arguments to be passed to the `subprocess.run` underlying command. + + Returns: + `Tuple[IO[str], IO[str]]`: A tuple with `stdin` and `stdout` to interact + with the process (input and output are utf-8 encoded). + + Example: + ```python + with _interactive_subprocess("git credential-store get") as (stdin, stdout): + # Write to stdin + stdin.write("url=hf.co\nusername=obama\n".encode("utf-8")) + stdin.flush() + + # Read from stdout + output = stdout.read().decode("utf-8") + ``` + """ + if isinstance(command, str): + command = command.split() + + with subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + encoding="utf-8", + errors="replace", # if not utf-8, replace char by � + cwd=folder or os.getcwd(), + **kwargs, + ) as process: + assert process.stdin is not None, "subprocess is opened as subprocess.PIPE" + assert process.stdout is not None, "subprocess is opened as subprocess.PIPE" + yield process.stdin, process.stdout diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_telemetry.py b/lib/python3.12/site-packages/huggingface_hub/utils/_telemetry.py new file mode 100644 index 0000000000000000000000000000000000000000..2ba4a6349a8de1c565263ec73d235d36f88b68cf --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_telemetry.py @@ -0,0 +1,126 @@ +from queue import Queue +from threading import Lock, Thread +from typing import Dict, Optional, Union +from urllib.parse import quote + +from .. import constants, logging +from . import build_hf_headers, get_session, hf_raise_for_status + + +logger = logging.get_logger(__name__) + +# Telemetry is sent by a separate thread to avoid blocking the main thread. +# A daemon thread is started once and consume tasks from the _TELEMETRY_QUEUE. +# If the thread stops for some reason -shouldn't happen-, we restart a new one. +_TELEMETRY_THREAD: Optional[Thread] = None +_TELEMETRY_THREAD_LOCK = Lock() # Lock to avoid starting multiple threads in parallel +_TELEMETRY_QUEUE: Queue = Queue() + + +def send_telemetry( + topic: str, + *, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, +) -> None: + """ + Sends telemetry that helps tracking usage of different HF libraries. + + This usage data helps us debug issues and prioritize new features. However, we understand that not everyone wants + to share additional information, and we respect your privacy. You can disable telemetry collection by setting the + `HF_HUB_DISABLE_TELEMETRY=1` as environment variable. Telemetry is also disabled in offline mode (i.e. when setting + `HF_HUB_OFFLINE=1`). + + Telemetry collection is run in a separate thread to minimize impact for the user. + + Args: + topic (`str`): + Name of the topic that is monitored. The topic is directly used to build the URL. If you want to monitor + subtopics, just use "/" separation. Examples: "gradio", "transformers/examples",... + library_name (`str`, *optional*): + The name of the library that is making the HTTP request. Will be added to the user-agent header. + library_version (`str`, *optional*): + The version of the library that is making the HTTP request. Will be added to the user-agent header. + user_agent (`str`, `dict`, *optional*): + The user agent info in the form of a dictionary or a single string. It will be completed with information about the installed packages. + + Example: + ```py + >>> from huggingface_hub.utils import send_telemetry + + # Send telemetry without library information + >>> send_telemetry("ping") + + # Send telemetry to subtopic with library information + >>> send_telemetry("gradio/local_link", library_name="gradio", library_version="3.22.1") + + # Send telemetry with additional data + >>> send_telemetry( + ... topic="examples", + ... library_name="transformers", + ... library_version="4.26.0", + ... user_agent={"pipeline": "text_classification", "framework": "flax"}, + ... ) + ``` + """ + if constants.HF_HUB_OFFLINE or constants.HF_HUB_DISABLE_TELEMETRY: + return + + _start_telemetry_thread() # starts thread only if doesn't exist yet + _TELEMETRY_QUEUE.put( + {"topic": topic, "library_name": library_name, "library_version": library_version, "user_agent": user_agent} + ) + + +def _start_telemetry_thread(): + """Start a daemon thread to consume tasks from the telemetry queue. + + If the thread is interrupted, start a new one. + """ + with _TELEMETRY_THREAD_LOCK: # avoid to start multiple threads if called concurrently + global _TELEMETRY_THREAD + if _TELEMETRY_THREAD is None or not _TELEMETRY_THREAD.is_alive(): + _TELEMETRY_THREAD = Thread(target=_telemetry_worker, daemon=True) + _TELEMETRY_THREAD.start() + + +def _telemetry_worker(): + """Wait for a task and consume it.""" + while True: + kwargs = _TELEMETRY_QUEUE.get() + _send_telemetry_in_thread(**kwargs) + _TELEMETRY_QUEUE.task_done() + + +def _send_telemetry_in_thread( + topic: str, + *, + library_name: Optional[str] = None, + library_version: Optional[str] = None, + user_agent: Union[Dict, str, None] = None, +) -> None: + """Contains the actual data sending data to the Hub. + + This function is called directly in gradio's analytics because + it is not possible to send telemetry from a daemon thread. + + See here: https://github.com/gradio-app/gradio/pull/8180 + + Please do not rename or remove this function. + """ + path = "/".join(quote(part) for part in topic.split("/") if len(part) > 0) + try: + r = get_session().head( + f"{constants.ENDPOINT}/api/telemetry/{path}", + headers=build_hf_headers( + token=False, # no need to send a token for telemetry + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + ), + ) + hf_raise_for_status(r) + except Exception as e: + # We don't want to error in case of connection errors of any kind. + logger.debug(f"Error while sending telemetry: {e}") diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_typing.py b/lib/python3.12/site-packages/huggingface_hub/utils/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..b8388ca0c003e4741a44b298f02ef69932737f93 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_typing.py @@ -0,0 +1,75 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# 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. +"""Handle typing imports based on system compatibility.""" + +import sys +from typing import Any, Callable, List, Literal, Type, TypeVar, Union, get_args, get_origin + + +UNION_TYPES: List[Any] = [Union] +if sys.version_info >= (3, 10): + from types import UnionType + + UNION_TYPES += [UnionType] + + +HTTP_METHOD_T = Literal["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"] + +# type hint meaning "function signature not changed by decorator" +CallableT = TypeVar("CallableT", bound=Callable) + +_JSON_SERIALIZABLE_TYPES = (int, float, str, bool, type(None)) + + +def is_jsonable(obj: Any) -> bool: + """Check if an object is JSON serializable. + + This is a weak check, as it does not check for the actual JSON serialization, but only for the types of the object. + It works correctly for basic use cases but do not guarantee an exhaustive check. + + Object is considered to be recursively json serializable if: + - it is an instance of int, float, str, bool, or NoneType + - it is a list or tuple and all its items are json serializable + - it is a dict and all its keys are strings and all its values are json serializable + """ + try: + if isinstance(obj, _JSON_SERIALIZABLE_TYPES): + return True + if isinstance(obj, (list, tuple)): + return all(is_jsonable(item) for item in obj) + if isinstance(obj, dict): + return all(isinstance(key, _JSON_SERIALIZABLE_TYPES) and is_jsonable(value) for key, value in obj.items()) + if hasattr(obj, "__json__"): + return True + return False + except RecursionError: + return False + + +def is_simple_optional_type(type_: Type) -> bool: + """Check if a type is optional, i.e. Optional[Type] or Union[Type, None] or Type | None, where Type is a non-composite type.""" + if get_origin(type_) in UNION_TYPES: + union_args = get_args(type_) + if len(union_args) == 2 and type(None) in union_args: + return True + return False + + +def unwrap_simple_optional_type(optional_type: Type) -> Type: + """Unwraps a simple optional type, i.e. returns Type from Optional[Type].""" + for arg in get_args(optional_type): + if arg is not type(None): + return arg + raise ValueError(f"'{optional_type}' is not an optional type") diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py b/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..27833f28e3e2030680fb72b95a547521bc08831b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py @@ -0,0 +1,226 @@ +# coding=utf-8 +# Copyright 2022-present, the HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contains utilities to validate argument values in `huggingface_hub`.""" + +import inspect +import re +import warnings +from functools import wraps +from itertools import chain +from typing import Any, Dict + +from huggingface_hub.errors import HFValidationError + +from ._typing import CallableT + + +REPO_ID_REGEX = re.compile( + r""" + ^ + (\b[\w\-.]+\b/)? # optional namespace (username or organization) + \b # starts with a word boundary + [\w\-.]{1,96} # repo_name: alphanumeric + . _ - + \b # ends with a word boundary + $ + """, + flags=re.VERBOSE, +) + + +def validate_hf_hub_args(fn: CallableT) -> CallableT: + """Validate values received as argument for any public method of `huggingface_hub`. + + The goal of this decorator is to harmonize validation of arguments reused + everywhere. By default, all defined validators are tested. + + Validators: + - [`~utils.validate_repo_id`]: `repo_id` must be `"repo_name"` + or `"namespace/repo_name"`. Namespace is a username or an organization. + - [`~utils.smoothly_deprecate_use_auth_token`]: Use `token` instead of + `use_auth_token` (only if `use_auth_token` is not expected by the decorated + function - in practice, always the case in `huggingface_hub`). + + Example: + ```py + >>> from huggingface_hub.utils import validate_hf_hub_args + + >>> @validate_hf_hub_args + ... def my_cool_method(repo_id: str): + ... print(repo_id) + + >>> my_cool_method(repo_id="valid_repo_id") + valid_repo_id + + >>> my_cool_method("other..repo..id") + huggingface_hub.utils._validators.HFValidationError: Cannot have -- or .. in repo_id: 'other..repo..id'. + + >>> my_cool_method(repo_id="other..repo..id") + huggingface_hub.utils._validators.HFValidationError: Cannot have -- or .. in repo_id: 'other..repo..id'. + + >>> @validate_hf_hub_args + ... def my_cool_auth_method(token: str): + ... print(token) + + >>> my_cool_auth_method(token="a token") + "a token" + + >>> my_cool_auth_method(use_auth_token="a use_auth_token") + "a use_auth_token" + + >>> my_cool_auth_method(token="a token", use_auth_token="a use_auth_token") + UserWarning: Both `token` and `use_auth_token` are passed (...) + "a token" + ``` + + Raises: + [`~utils.HFValidationError`]: + If an input is not valid. + """ + # TODO: add an argument to opt-out validation for specific argument? + signature = inspect.signature(fn) + + # Should the validator switch `use_auth_token` values to `token`? In practice, always + # True in `huggingface_hub`. Might not be the case in a downstream library. + check_use_auth_token = "use_auth_token" not in signature.parameters and "token" in signature.parameters + + @wraps(fn) + def _inner_fn(*args, **kwargs): + has_token = False + for arg_name, arg_value in chain( + zip(signature.parameters, args), # Args values + kwargs.items(), # Kwargs values + ): + if arg_name in ["repo_id", "from_id", "to_id"]: + validate_repo_id(arg_value) + + elif arg_name == "token" and arg_value is not None: + has_token = True + + if check_use_auth_token: + kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.__name__, has_token=has_token, kwargs=kwargs) + + return fn(*args, **kwargs) + + return _inner_fn # type: ignore + + +def validate_repo_id(repo_id: str) -> None: + """Validate `repo_id` is valid. + + This is not meant to replace the proper validation made on the Hub but rather to + avoid local inconsistencies whenever possible (example: passing `repo_type` in the + `repo_id` is forbidden). + + Rules: + - Between 1 and 96 characters. + - Either "repo_name" or "namespace/repo_name" + - [a-zA-Z0-9] or "-", "_", "." + - "--" and ".." are forbidden + + Valid: `"foo"`, `"foo/bar"`, `"123"`, `"Foo-BAR_foo.bar123"` + + Not valid: `"datasets/foo/bar"`, `".repo_id"`, `"foo--bar"`, `"foo.git"` + + Example: + ```py + >>> from huggingface_hub.utils import validate_repo_id + >>> validate_repo_id(repo_id="valid_repo_id") + >>> validate_repo_id(repo_id="other..repo..id") + huggingface_hub.utils._validators.HFValidationError: Cannot have -- or .. in repo_id: 'other..repo..id'. + ``` + + Discussed in https://github.com/huggingface/huggingface_hub/issues/1008. + In moon-landing (internal repository): + - https://github.com/huggingface/moon-landing/blob/main/server/lib/Names.ts#L27 + - https://github.com/huggingface/moon-landing/blob/main/server/views/components/NewRepoForm/NewRepoForm.svelte#L138 + """ + if not isinstance(repo_id, str): + # Typically, a Path is not a repo_id + raise HFValidationError(f"Repo id must be a string, not {type(repo_id)}: '{repo_id}'.") + + if repo_id.count("/") > 1: + raise HFValidationError( + "Repo id must be in the form 'repo_name' or 'namespace/repo_name':" + f" '{repo_id}'. Use `repo_type` argument if needed." + ) + + if not REPO_ID_REGEX.match(repo_id): + raise HFValidationError( + "Repo id must use alphanumeric chars or '-', '_', '.', '--' and '..' are" + " forbidden, '-' and '.' cannot start or end the name, max length is 96:" + f" '{repo_id}'." + ) + + if "--" in repo_id or ".." in repo_id: + raise HFValidationError(f"Cannot have -- or .. in repo_id: '{repo_id}'.") + + if repo_id.endswith(".git"): + raise HFValidationError(f"Repo_id cannot end by '.git': '{repo_id}'.") + + +def smoothly_deprecate_use_auth_token(fn_name: str, has_token: bool, kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Smoothly deprecate `use_auth_token` in the `huggingface_hub` codebase. + + The long-term goal is to remove any mention of `use_auth_token` in the codebase in + favor of a unique and less verbose `token` argument. This will be done a few steps: + + 0. Step 0: methods that require a read-access to the Hub use the `use_auth_token` + argument (`str`, `bool` or `None`). Methods requiring write-access have a `token` + argument (`str`, `None`). This implicit rule exists to be able to not send the + token when not necessary (`use_auth_token=False`) even if logged in. + + 1. Step 1: we want to harmonize everything and use `token` everywhere (supporting + `token=False` for read-only methods). In order not to break existing code, if + `use_auth_token` is passed to a function, the `use_auth_token` value is passed + as `token` instead, without any warning. + a. Corner case: if both `use_auth_token` and `token` values are passed, a warning + is thrown and the `use_auth_token` value is ignored. + + 2. Step 2: Once it is release, we should push downstream libraries to switch from + `use_auth_token` to `token` as much as possible, but without throwing a warning + (e.g. manually create issues on the corresponding repos). + + 3. Step 3: After a transitional period (6 months e.g. until April 2023?), we update + `huggingface_hub` to throw a warning on `use_auth_token`. Hopefully, very few + users will be impacted as it would have already been fixed. + In addition, unit tests in `huggingface_hub` must be adapted to expect warnings + to be thrown (but still use `use_auth_token` as before). + + 4. Step 4: After a normal deprecation cycle (3 releases ?), remove this validator. + `use_auth_token` will definitely not be supported. + In addition, we update unit tests in `huggingface_hub` to use `token` everywhere. + + This has been discussed in: + - https://github.com/huggingface/huggingface_hub/issues/1094. + - https://github.com/huggingface/huggingface_hub/pull/928 + - (related) https://github.com/huggingface/huggingface_hub/pull/1064 + """ + new_kwargs = kwargs.copy() # do not mutate input ! + + use_auth_token = new_kwargs.pop("use_auth_token", None) # remove from kwargs + if use_auth_token is not None: + if has_token: + warnings.warn( + "Both `token` and `use_auth_token` are passed to" + f" `{fn_name}` with non-None values. `token` is now the" + " preferred argument to pass a User Access Token." + " `use_auth_token` value will be ignored." + ) + else: + # `token` argument is not passed and a non-None value is passed in + # `use_auth_token` => use `use_auth_token` value as `token` kwarg. + new_kwargs["token"] = use_auth_token + + return new_kwargs diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/_xet.py b/lib/python3.12/site-packages/huggingface_hub/utils/_xet.py new file mode 100644 index 0000000000000000000000000000000000000000..e80b0e804b61849eba327bdff28ed47b6b133ad0 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/_xet.py @@ -0,0 +1,188 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Dict, Optional + +import requests + +from .. import constants +from . import get_session, hf_raise_for_status, validate_hf_hub_args + + +class XetTokenType(str, Enum): + READ = "read" + WRITE = "write" + + +@dataclass(frozen=True) +class XetFileData: + file_hash: str + refresh_route: str + + +@dataclass(frozen=True) +class XetConnectionInfo: + access_token: str + expiration_unix_epoch: int + endpoint: str + + +def parse_xet_file_data_from_response(response: requests.Response) -> Optional[XetFileData]: + """ + Parse XET file metadata from an HTTP response. + + This function extracts XET file metadata from the HTTP headers or HTTP links + of a given response object. If the required metadata is not found, it returns `None`. + + Args: + response (`requests.Response`): + The HTTP response object containing headers dict and links dict to extract the XET metadata from. + Returns: + `Optional[XetFileData]`: + An instance of `XetFileData` containing the file hash and refresh route if the metadata + is found. Returns `None` if the required metadata is missing. + """ + if response is None: + return None + try: + file_hash = response.headers[constants.HUGGINGFACE_HEADER_X_XET_HASH] + + if constants.HUGGINGFACE_HEADER_LINK_XET_AUTH_KEY in response.links: + refresh_route = response.links[constants.HUGGINGFACE_HEADER_LINK_XET_AUTH_KEY]["url"] + else: + refresh_route = response.headers[constants.HUGGINGFACE_HEADER_X_XET_REFRESH_ROUTE] + except KeyError: + return None + + return XetFileData( + file_hash=file_hash, + refresh_route=refresh_route, + ) + + +def parse_xet_connection_info_from_headers(headers: Dict[str, str]) -> Optional[XetConnectionInfo]: + """ + Parse XET connection info from the HTTP headers or return None if not found. + Args: + headers (`Dict`): + HTTP headers to extract the XET metadata from. + Returns: + `XetConnectionInfo` or `None`: + The information needed to connect to the XET storage service. + Returns `None` if the headers do not contain the XET connection info. + """ + try: + endpoint = headers[constants.HUGGINGFACE_HEADER_X_XET_ENDPOINT] + access_token = headers[constants.HUGGINGFACE_HEADER_X_XET_ACCESS_TOKEN] + expiration_unix_epoch = int(headers[constants.HUGGINGFACE_HEADER_X_XET_EXPIRATION]) + except (KeyError, ValueError, TypeError): + return None + + return XetConnectionInfo( + endpoint=endpoint, + access_token=access_token, + expiration_unix_epoch=expiration_unix_epoch, + ) + + +@validate_hf_hub_args +def refresh_xet_connection_info( + *, + file_data: XetFileData, + headers: Dict[str, str], +) -> XetConnectionInfo: + """ + Utilizes the information in the parsed metadata to request the Hub xet connection information. + This includes the access token, expiration, and XET service URL. + Args: + file_data: (`XetFileData`): + The file data needed to refresh the xet connection information. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + Returns: + `XetConnectionInfo`: + The connection information needed to make the request to the xet storage service. + Raises: + [`~utils.HfHubHTTPError`] + If the Hub API returned an error. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the Hub API response is improperly formatted. + """ + if file_data.refresh_route is None: + raise ValueError("The provided xet metadata does not contain a refresh endpoint.") + return _fetch_xet_connection_info_with_url(file_data.refresh_route, headers) + + +@validate_hf_hub_args +def fetch_xet_connection_info_from_repo_info( + *, + token_type: XetTokenType, + repo_id: str, + repo_type: str, + revision: Optional[str] = None, + headers: Dict[str, str], + endpoint: Optional[str] = None, + params: Optional[Dict[str, str]] = None, +) -> XetConnectionInfo: + """ + Uses the repo info to request a xet access token from Hub. + Args: + token_type (`XetTokenType`): + Type of the token to request: `"read"` or `"write"`. + repo_id (`str`): + A namespace (user or an organization) and a repo name separated by a `/`. + repo_type (`str`): + Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`. + revision (`str`, `optional`): + The revision of the repo to get the token for. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + endpoint (`str`, `optional`): + The endpoint to use for the request. Defaults to the Hub endpoint. + params (`Dict[str, str]`, `optional`): + Additional parameters to pass with the request. + Returns: + `XetConnectionInfo`: + The connection information needed to make the request to the xet storage service. + Raises: + [`~utils.HfHubHTTPError`] + If the Hub API returned an error. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the Hub API response is improperly formatted. + """ + endpoint = endpoint if endpoint is not None else constants.ENDPOINT + url = f"{endpoint}/api/{repo_type}s/{repo_id}/xet-{token_type.value}-token/{revision}" + return _fetch_xet_connection_info_with_url(url, headers, params) + + +@validate_hf_hub_args +def _fetch_xet_connection_info_with_url( + url: str, + headers: Dict[str, str], + params: Optional[Dict[str, str]] = None, +) -> XetConnectionInfo: + """ + Requests the xet connection info from the supplied URL. This includes the + access token, expiration time, and endpoint to use for the xet storage service. + Args: + url: (`str`): + The access token endpoint URL. + headers (`Dict[str, str]`): + Headers to use for the request, including authorization headers and user agent. + params (`Dict[str, str]`, `optional`): + Additional parameters to pass with the request. + Returns: + `XetConnectionInfo`: + The connection information needed to make the request to the xet storage service. + Raises: + [`~utils.HfHubHTTPError`] + If the Hub API returned an error. + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) + If the Hub API response is improperly formatted. + """ + resp = get_session().get(headers=headers, url=url, params=params) + hf_raise_for_status(resp) + + metadata = parse_xet_connection_info_from_headers(resp.headers) # type: ignore + if metadata is None: + raise ValueError("Xet headers have not been correctly set by the server.") + return metadata diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/endpoint_helpers.py b/lib/python3.12/site-packages/huggingface_hub/utils/endpoint_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..85cd86011b78bcdc57034aeebc3c01e9e721ab50 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/endpoint_helpers.py @@ -0,0 +1,66 @@ +# 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. +""" +Helpful utility functions and classes in relation to exploring API endpoints +with the aim for a user-friendly interface. +""" + +import math +import re +from typing import TYPE_CHECKING + +from ..repocard_data import ModelCardData + + +if TYPE_CHECKING: + from ..hf_api import ModelInfo + + +def _is_emission_within_threshold(model_info: "ModelInfo", minimum_threshold: float, maximum_threshold: float) -> bool: + """Checks if a model's emission is within a given threshold. + + Args: + model_info (`ModelInfo`): + A model info object containing the model's emission information. + minimum_threshold (`float`): + A minimum carbon threshold to filter by, such as 1. + maximum_threshold (`float`): + A maximum carbon threshold to filter by, such as 10. + + Returns: + `bool`: Whether the model's emission is within the given threshold. + """ + if minimum_threshold is None and maximum_threshold is None: + raise ValueError("Both `minimum_threshold` and `maximum_threshold` cannot both be `None`") + if minimum_threshold is None: + minimum_threshold = -1 + if maximum_threshold is None: + maximum_threshold = math.inf + + card_data = getattr(model_info, "card_data", None) + if card_data is None or not isinstance(card_data, (dict, ModelCardData)): + return False + + # Get CO2 emission metadata + emission = card_data.get("co2_eq_emissions", None) + if isinstance(emission, dict): + emission = emission["emissions"] + if not emission: + return False + + # Filter out if value is missing or out of range + matched = re.search(r"\d+\.\d+|\d+", str(emission)) + if matched is None: + return False + + emission_value = float(matched.group(0)) + return minimum_threshold <= emission_value <= maximum_threshold diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/insecure_hashlib.py b/lib/python3.12/site-packages/huggingface_hub/utils/insecure_hashlib.py new file mode 100644 index 0000000000000000000000000000000000000000..f232ee0adcfc52dcc18b5ea4d9c913b206521f71 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/insecure_hashlib.py @@ -0,0 +1,34 @@ +# Taken from https://github.com/mlflow/mlflow/pull/10119 +# +# DO NOT use this function for security purposes (e.g., password hashing). +# +# In Python >= 3.9, insecure hashing algorithms such as MD5 fail in FIPS-compliant +# environments unless `usedforsecurity=False` is explicitly passed. +# +# References: +# - https://github.com/mlflow/mlflow/issues/9905 +# - https://github.com/mlflow/mlflow/pull/10119 +# - https://docs.python.org/3/library/hashlib.html +# - https://github.com/huggingface/transformers/pull/27038 +# +# Usage: +# ```python +# # Use +# from huggingface_hub.utils.insecure_hashlib import sha256 +# # instead of +# from hashlib import sha256 +# +# # Use +# from huggingface_hub.utils import insecure_hashlib +# # instead of +# import hashlib +# ``` +import functools +import hashlib +import sys + + +_kwargs = {"usedforsecurity": False} if sys.version_info >= (3, 9) else {} +md5 = functools.partial(hashlib.md5, **_kwargs) +sha1 = functools.partial(hashlib.sha1, **_kwargs) +sha256 = functools.partial(hashlib.sha256, **_kwargs) diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/logging.py b/lib/python3.12/site-packages/huggingface_hub/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..813719683a54cc65768bab5488e7ea153ad08d7e --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/logging.py @@ -0,0 +1,188 @@ +# coding=utf-8 +# Copyright 2020 Optuna, Hugging Face +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Logging utilities.""" + +import logging +import os +from logging import ( + CRITICAL, # NOQA + DEBUG, # NOQA + ERROR, # NOQA + FATAL, # NOQA + INFO, # NOQA + NOTSET, # NOQA + WARN, # NOQA + WARNING, # NOQA +) +from typing import Optional + +from .. import constants + + +log_levels = { + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + +_default_log_level = logging.WARNING + + +def _get_library_name() -> str: + return __name__.split(".")[0] + + +def _get_library_root_logger() -> logging.Logger: + return logging.getLogger(_get_library_name()) + + +def _get_default_logging_level(): + """ + If `HF_HUB_VERBOSITY` env var is set to one of the valid choices return that as the new default level. If it is not + - fall back to `_default_log_level` + """ + env_level_str = os.getenv("HF_HUB_VERBOSITY", None) + if env_level_str: + if env_level_str in log_levels: + return log_levels[env_level_str] + else: + logging.getLogger().warning( + f"Unknown option HF_HUB_VERBOSITY={env_level_str}, has to be one of: {', '.join(log_levels.keys())}" + ) + return _default_log_level + + +def _configure_library_root_logger() -> None: + library_root_logger = _get_library_root_logger() + library_root_logger.addHandler(logging.StreamHandler()) + library_root_logger.setLevel(_get_default_logging_level()) + + +def _reset_library_root_logger() -> None: + library_root_logger = _get_library_root_logger() + library_root_logger.setLevel(logging.NOTSET) + + +def get_logger(name: Optional[str] = None) -> logging.Logger: + """ + Returns a logger with the specified name. This function is not supposed + to be directly accessed by library users. + + Args: + name (`str`, *optional*): + The name of the logger to get, usually the filename + + Example: + + ```python + >>> from huggingface_hub import get_logger + + >>> logger = get_logger(__file__) + >>> logger.set_verbosity_info() + ``` + """ + + if name is None: + name = _get_library_name() + + return logging.getLogger(name) + + +def get_verbosity() -> int: + """Return the current level for the HuggingFace Hub's root logger. + + Returns: + Logging level, e.g., `huggingface_hub.logging.DEBUG` and + `huggingface_hub.logging.INFO`. + + + + HuggingFace Hub has following logging levels: + + - `huggingface_hub.logging.CRITICAL`, `huggingface_hub.logging.FATAL` + - `huggingface_hub.logging.ERROR` + - `huggingface_hub.logging.WARNING`, `huggingface_hub.logging.WARN` + - `huggingface_hub.logging.INFO` + - `huggingface_hub.logging.DEBUG` + + + """ + return _get_library_root_logger().getEffectiveLevel() + + +def set_verbosity(verbosity: int) -> None: + """ + Sets the level for the HuggingFace Hub's root logger. + + Args: + verbosity (`int`): + Logging level, e.g., `huggingface_hub.logging.DEBUG` and + `huggingface_hub.logging.INFO`. + """ + _get_library_root_logger().setLevel(verbosity) + + +def set_verbosity_info(): + """ + Sets the verbosity to `logging.INFO`. + """ + return set_verbosity(INFO) + + +def set_verbosity_warning(): + """ + Sets the verbosity to `logging.WARNING`. + """ + return set_verbosity(WARNING) + + +def set_verbosity_debug(): + """ + Sets the verbosity to `logging.DEBUG`. + """ + return set_verbosity(DEBUG) + + +def set_verbosity_error(): + """ + Sets the verbosity to `logging.ERROR`. + """ + return set_verbosity(ERROR) + + +def disable_propagation() -> None: + """ + Disable propagation of the library log outputs. Note that log propagation is + disabled by default. + """ + _get_library_root_logger().propagate = False + + +def enable_propagation() -> None: + """ + Enable propagation of the library log outputs. Please disable the + HuggingFace Hub's default handler to prevent double logging if the root + logger has been configured. + """ + _get_library_root_logger().propagate = True + + +_configure_library_root_logger() + +if constants.HF_DEBUG: + # If `HF_DEBUG` environment variable is set, set the verbosity of `huggingface_hub` logger to `DEBUG`. + set_verbosity_debug() diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/sha.py b/lib/python3.12/site-packages/huggingface_hub/utils/sha.py new file mode 100644 index 0000000000000000000000000000000000000000..001c3fe8b2f37a64e890888ca3d521c10ec8f03b --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/sha.py @@ -0,0 +1,64 @@ +"""Utilities to efficiently compute the SHA 256 hash of a bunch of bytes.""" + +from typing import BinaryIO, Optional + +from .insecure_hashlib import sha1, sha256 + + +def sha_fileobj(fileobj: BinaryIO, chunk_size: Optional[int] = None) -> bytes: + """ + Computes the sha256 hash of the given file object, by chunks of size `chunk_size`. + + Args: + fileobj (file-like object): + The File object to compute sha256 for, typically obtained with `open(path, "rb")` + chunk_size (`int`, *optional*): + The number of bytes to read from `fileobj` at once, defaults to 1MB. + + Returns: + `bytes`: `fileobj`'s sha256 hash as bytes + """ + chunk_size = chunk_size if chunk_size is not None else 1024 * 1024 + + sha = sha256() + while True: + chunk = fileobj.read(chunk_size) + sha.update(chunk) + if not chunk: + break + return sha.digest() + + +def git_hash(data: bytes) -> str: + """ + Computes the git-sha1 hash of the given bytes, using the same algorithm as git. + + This is equivalent to running `git hash-object`. See https://git-scm.com/docs/git-hash-object + for more details. + + Note: this method is valid for regular files. For LFS files, the proper git hash is supposed to be computed on the + pointer file content, not the actual file content. However, for simplicity, we directly compare the sha256 of + the LFS file content when we want to compare LFS files. + + Args: + data (`bytes`): + The data to compute the git-hash for. + + Returns: + `str`: the git-hash of `data` as an hexadecimal string. + + Example: + ```python + >>> from huggingface_hub.utils.sha import git_hash + >>> git_hash(b"Hello, World!") + 'b45ef6fec89518d314f546fd6c3025367b721684' + ``` + """ + # Taken from https://gist.github.com/msabramo/763200 + # Note: no need to optimize by reading the file in chunks as we're not supposed to hash huge files (5MB maximum). + sha = sha1() + sha.update(b"blob ") + sha.update(str(len(data)).encode()) + sha.update(b"\0") + sha.update(data) + return sha.hexdigest() diff --git a/lib/python3.12/site-packages/huggingface_hub/utils/tqdm.py b/lib/python3.12/site-packages/huggingface_hub/utils/tqdm.py new file mode 100644 index 0000000000000000000000000000000000000000..4c1fcef4beb73bae13c57b3f66c5828e775b7cd9 --- /dev/null +++ b/lib/python3.12/site-packages/huggingface_hub/utils/tqdm.py @@ -0,0 +1,307 @@ +# coding=utf-8 +# Copyright 2021 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 +"""Utility helpers to handle progress bars in `huggingface_hub`. + +Example: + 1. Use `huggingface_hub.utils.tqdm` as you would use `tqdm.tqdm` or `tqdm.auto.tqdm`. + 2. To disable progress bars, either use `disable_progress_bars()` helper or set the + environment variable `HF_HUB_DISABLE_PROGRESS_BARS` to 1. + 3. To re-enable progress bars, use `enable_progress_bars()`. + 4. To check whether progress bars are disabled, use `are_progress_bars_disabled()`. + +NOTE: Environment variable `HF_HUB_DISABLE_PROGRESS_BARS` has the priority. + +Example: + ```py + >>> from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm + + # Disable progress bars globally + >>> disable_progress_bars() + + # Use as normal `tqdm` + >>> for _ in tqdm(range(5)): + ... pass + + # Still not showing progress bars, as `disable=False` is overwritten to `True`. + >>> for _ in tqdm(range(5), disable=False): + ... pass + + >>> are_progress_bars_disabled() + True + + # Re-enable progress bars globally + >>> enable_progress_bars() + + # Progress bar will be shown ! + >>> for _ in tqdm(range(5)): + ... pass + 100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s] + ``` + +Group-based control: + ```python + # Disable progress bars for a specific group + >>> disable_progress_bars("peft.foo") + + # Check state of different groups + >>> assert not are_progress_bars_disabled("peft")) + >>> assert not are_progress_bars_disabled("peft.something") + >>> assert are_progress_bars_disabled("peft.foo")) + >>> assert are_progress_bars_disabled("peft.foo.bar")) + + # Enable progress bars for a subgroup + >>> enable_progress_bars("peft.foo.bar") + + # Check if enabling a subgroup affects the parent group + >>> assert are_progress_bars_disabled("peft.foo")) + >>> assert not are_progress_bars_disabled("peft.foo.bar")) + + # No progress bar for `name="peft.foo"` + >>> for _ in tqdm(range(5), name="peft.foo"): + ... pass + + # Progress bar will be shown for `name="peft.foo.bar"` + >>> for _ in tqdm(range(5), name="peft.foo.bar"): + ... pass + 100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s] + + ``` +""" + +import io +import logging +import os +import warnings +from contextlib import contextmanager, nullcontext +from pathlib import Path +from typing import ContextManager, Dict, Iterator, Optional, Union + +from tqdm.auto import tqdm as old_tqdm + +from ..constants import HF_HUB_DISABLE_PROGRESS_BARS + + +# The `HF_HUB_DISABLE_PROGRESS_BARS` environment variable can be True, False, or not set (None), +# allowing for control over progress bar visibility. When set, this variable takes precedence +# over programmatic settings, dictating whether progress bars should be shown or hidden globally. +# Essentially, the environment variable's setting overrides any code-based configurations. +# +# If `HF_HUB_DISABLE_PROGRESS_BARS` is not defined (None), it implies that users can manage +# progress bar visibility through code. By default, progress bars are turned on. + + +progress_bar_states: Dict[str, bool] = {} + + +def disable_progress_bars(name: Optional[str] = None) -> None: + """ + Disable progress bars either globally or for a specified group. + + This function updates the state of progress bars based on a group name. + If no group name is provided, all progress bars are disabled. The operation + respects the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable's setting. + + Args: + name (`str`, *optional*): + The name of the group for which to disable the progress bars. If None, + progress bars are disabled globally. + + Raises: + Warning: If the environment variable precludes changes. + """ + if HF_HUB_DISABLE_PROGRESS_BARS is False: + warnings.warn( + "Cannot disable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=0` is set and has priority." + ) + return + + if name is None: + progress_bar_states.clear() + progress_bar_states["_global"] = False + else: + keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")] + for key in keys_to_remove: + del progress_bar_states[key] + progress_bar_states[name] = False + + +def enable_progress_bars(name: Optional[str] = None) -> None: + """ + Enable progress bars either globally or for a specified group. + + This function sets the progress bars to enabled for the specified group or globally + if no group is specified. The operation is subject to the `HF_HUB_DISABLE_PROGRESS_BARS` + environment setting. + + Args: + name (`str`, *optional*): + The name of the group for which to enable the progress bars. If None, + progress bars are enabled globally. + + Raises: + Warning: If the environment variable precludes changes. + """ + if HF_HUB_DISABLE_PROGRESS_BARS is True: + warnings.warn( + "Cannot enable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=1` is set and has priority." + ) + return + + if name is None: + progress_bar_states.clear() + progress_bar_states["_global"] = True + else: + keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")] + for key in keys_to_remove: + del progress_bar_states[key] + progress_bar_states[name] = True + + +def are_progress_bars_disabled(name: Optional[str] = None) -> bool: + """ + Check if progress bars are disabled globally or for a specific group. + + This function returns whether progress bars are disabled for a given group or globally. + It checks the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable first, then the programmatic + settings. + + Args: + name (`str`, *optional*): + The group name to check; if None, checks the global setting. + + Returns: + `bool`: True if progress bars are disabled, False otherwise. + """ + if HF_HUB_DISABLE_PROGRESS_BARS is True: + return True + + if name is None: + return not progress_bar_states.get("_global", True) + + while name: + if name in progress_bar_states: + return not progress_bar_states[name] + name = ".".join(name.split(".")[:-1]) + + return not progress_bar_states.get("_global", True) + + +def is_tqdm_disabled(log_level: int) -> Optional[bool]: + """ + Determine if tqdm progress bars should be disabled based on logging level and environment settings. + + see https://github.com/huggingface/huggingface_hub/pull/2000 and https://github.com/huggingface/huggingface_hub/pull/2698. + """ + if log_level == logging.NOTSET: + return True + if os.getenv("TQDM_POSITION") == "-1": + return False + return None + + +class tqdm(old_tqdm): + """ + Class to override `disable` argument in case progress bars are globally disabled. + + Taken from https://github.com/tqdm/tqdm/issues/619#issuecomment-619639324. + """ + + def __init__(self, *args, **kwargs): + name = kwargs.pop("name", None) # do not pass `name` to `tqdm` + if are_progress_bars_disabled(name): + kwargs["disable"] = True + super().__init__(*args, **kwargs) + + def __delattr__(self, attr: str) -> None: + """Fix for https://github.com/huggingface/huggingface_hub/issues/1603""" + try: + super().__delattr__(attr) + except AttributeError: + if attr != "_lock": + raise + + +@contextmanager +def tqdm_stream_file(path: Union[Path, str]) -> Iterator[io.BufferedReader]: + """ + Open a file as binary and wrap the `read` method to display a progress bar when it's streamed. + + First implemented in `transformers` in 2019 but removed when switched to git-lfs. Used in `huggingface_hub` to show + progress bar when uploading an LFS file to the Hub. See github.com/huggingface/transformers/pull/2078#discussion_r354739608 + for implementation details. + + Note: currently implementation handles only files stored on disk as it is the most common use case. Could be + extended to stream any `BinaryIO` object but we might have to debug some corner cases. + + Example: + ```py + >>> with tqdm_stream_file("config.json") as f: + >>> requests.put(url, data=f) + config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s] + ``` + """ + if isinstance(path, str): + path = Path(path) + + with path.open("rb") as f: + total_size = path.stat().st_size + pbar = tqdm( + unit="B", + unit_scale=True, + total=total_size, + initial=0, + desc=path.name, + ) + + f_read = f.read + + def _inner_read(size: Optional[int] = -1) -> bytes: + data = f_read(size) + pbar.update(len(data)) + return data + + f.read = _inner_read # type: ignore + + yield f + + pbar.close() + + +def _get_progress_bar_context( + *, + desc: str, + log_level: int, + total: Optional[int] = None, + initial: int = 0, + unit: str = "B", + unit_scale: bool = True, + name: Optional[str] = None, + _tqdm_bar: Optional[tqdm] = None, +) -> ContextManager[tqdm]: + if _tqdm_bar is not None: + return nullcontext(_tqdm_bar) + # ^ `contextlib.nullcontext` mimics a context manager that does nothing + # Makes it easier to use the same code path for both cases but in the later + # case, the progress bar is not closed when exiting the context manager. + + return tqdm( + unit=unit, + unit_scale=unit_scale, + total=total, + initial=initial, + desc=desc, + disable=is_tqdm_disabled(log_level=log_level), + name=name, + ) diff --git a/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/INSTALLER b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/LICENSE.txt b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..7b190ca6712aa09eede3e6de79f68d7fa29072da --- /dev/null +++ b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/LICENSE.txt @@ -0,0 +1,28 @@ +Copyright 2011 Pallets + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/METADATA b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ddf54648499557c652181f6126362ffd5751c273 --- /dev/null +++ b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/METADATA @@ -0,0 +1,60 @@ +Metadata-Version: 2.1 +Name: itsdangerous +Version: 2.2.0 +Summary: Safely pass data to untrusted environments and back. +Maintainer-email: Pallets +Requires-Python: >=3.8 +Description-Content-Type: text/markdown +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: BSD License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Typing :: Typed +Project-URL: Changes, https://itsdangerous.palletsprojects.com/changes/ +Project-URL: Chat, https://discord.gg/pallets +Project-URL: Documentation, https://itsdangerous.palletsprojects.com/ +Project-URL: Donate, https://palletsprojects.com/donate +Project-URL: Source, https://github.com/pallets/itsdangerous/ + +# ItsDangerous + +... so better sign this + +Various helpers to pass data to untrusted environments and to get it +back safe and sound. Data is cryptographically signed to ensure that a +token has not been tampered with. + +It's possible to customize how data is serialized. Data is compressed as +needed. A timestamp can be added and verified automatically while +loading a token. + + +## A Simple Example + +Here's how you could generate a token for transmitting a user's id and +name between web requests. + +```python +from itsdangerous import URLSafeSerializer +auth_s = URLSafeSerializer("secret key", "auth") +token = auth_s.dumps({"id": 5, "name": "itsdangerous"}) + +print(token) +# eyJpZCI6NSwibmFtZSI6Iml0c2Rhbmdlcm91cyJ9.6YP6T0BaO67XP--9UzTrmurXSmg + +data = auth_s.loads(token) +print(data["name"]) +# itsdangerous +``` + + +## Donate + +The Pallets organization develops and supports ItsDangerous and other +popular packages. In order to grow the community of contributors and +users, and allow the maintainers to devote more time to the projects, +[please donate today][]. + +[please donate today]: https://palletsprojects.com/donate + diff --git a/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/RECORD b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..245f43e83264f9b2d801757b931bce38fccdb0be --- /dev/null +++ b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/RECORD @@ -0,0 +1,22 @@ +itsdangerous-2.2.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +itsdangerous-2.2.0.dist-info/LICENSE.txt,sha256=Y68JiRtr6K0aQlLtQ68PTvun_JSOIoNnvtfzxa4LCdc,1475 +itsdangerous-2.2.0.dist-info/METADATA,sha256=0rk0-1ZwihuU5DnwJVwPWoEI4yWOyCexih3JyZHblhE,1924 +itsdangerous-2.2.0.dist-info/RECORD,, +itsdangerous-2.2.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81 +itsdangerous/__init__.py,sha256=4SK75sCe29xbRgQE1ZQtMHnKUuZYAf3bSpZOrff1IAY,1427 +itsdangerous/__pycache__/__init__.cpython-312.pyc,, +itsdangerous/__pycache__/_json.cpython-312.pyc,, +itsdangerous/__pycache__/encoding.cpython-312.pyc,, +itsdangerous/__pycache__/exc.cpython-312.pyc,, +itsdangerous/__pycache__/serializer.cpython-312.pyc,, +itsdangerous/__pycache__/signer.cpython-312.pyc,, +itsdangerous/__pycache__/timed.cpython-312.pyc,, +itsdangerous/__pycache__/url_safe.cpython-312.pyc,, +itsdangerous/_json.py,sha256=wPQGmge2yZ9328EHKF6gadGeyGYCJQKxtU-iLKE6UnA,473 +itsdangerous/encoding.py,sha256=wwTz5q_3zLcaAdunk6_vSoStwGqYWe307Zl_U87aRFM,1409 +itsdangerous/exc.py,sha256=Rr3exo0MRFEcPZltwecyK16VV1bE2K9_F1-d-ljcUn4,3201 +itsdangerous/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +itsdangerous/serializer.py,sha256=PmdwADLqkSyQLZ0jOKAgDsAW4k_H0TlA71Ei3z0C5aI,15601 +itsdangerous/signer.py,sha256=YO0CV7NBvHA6j549REHJFUjUojw2pHqwcUpQnU7yNYQ,9647 +itsdangerous/timed.py,sha256=6RvDMqNumGMxf0-HlpaZdN9PUQQmRvrQGplKhxuivUs,8083 +itsdangerous/url_safe.py,sha256=az4e5fXi_vs-YbWj8YZwn4wiVKfeD--GEKRT5Ueu4P4,2505 diff --git a/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/WHEEL b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..3b5e64b5e6c4a210201d1676a891fd57b15cda99 --- /dev/null +++ b/lib/python3.12/site-packages/itsdangerous-2.2.0.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: flit 3.9.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/INSTALLER b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/METADATA b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..7cd74a25cfa5548a3b42f0c97ac5f572dc2df62a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/METADATA @@ -0,0 +1,797 @@ +Metadata-Version: 2.4 +Name: pydantic +Version: 2.11.10 +Summary: Data validation using Python type hints +Project-URL: Homepage, https://github.com/pydantic/pydantic +Project-URL: Documentation, https://docs.pydantic.dev +Project-URL: Funding, https://github.com/sponsors/samuelcolvin +Project-URL: Source, https://github.com/pydantic/pydantic +Project-URL: Changelog, https://docs.pydantic.dev/latest/changelog/ +Author-email: Samuel Colvin , Eric Jolibois , Hasan Ramezani , Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>, Terrence Dorsey , David Montague , Serge Matveenko , Marcelo Trylesinski , Sydney Runkle , David Hewitt , Alex Hall , Victorien Plot +License-Expression: MIT +License-File: LICENSE +Classifier: Development Status :: 5 - Production/Stable +Classifier: Framework :: Hypothesis +Classifier: Framework :: Pydantic +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Information Technology +Classifier: License :: OSI Approved :: MIT License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Internet +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Requires-Python: >=3.9 +Requires-Dist: annotated-types>=0.6.0 +Requires-Dist: pydantic-core==2.33.2 +Requires-Dist: typing-extensions>=4.12.2 +Requires-Dist: typing-inspection>=0.4.0 +Provides-Extra: email +Requires-Dist: email-validator>=2.0.0; extra == 'email' +Provides-Extra: timezone +Requires-Dist: tzdata; (python_version >= '3.9' and platform_system == 'Windows') and extra == 'timezone' +Description-Content-Type: text/markdown + +# Pydantic +[![CI](https://img.shields.io/github/actions/workflow/status/pydantic/pydantic/ci.yml?branch=main&logo=github&label=CI)](https://github.com/pydantic/pydantic/actions?query=event%3Apush+branch%3Amain+workflow%3ACI) +[![Coverage](https://coverage-badge.samuelcolvin.workers.dev/pydantic/pydantic.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/pydantic/pydantic) +[![pypi](https://img.shields.io/pypi/v/pydantic.svg)](https://pypi.python.org/pypi/pydantic) +[![CondaForge](https://img.shields.io/conda/v/conda-forge/pydantic.svg)](https://anaconda.org/conda-forge/pydantic) +[![downloads](https://static.pepy.tech/badge/pydantic/month)](https://pepy.tech/project/pydantic) +[![versions](https://img.shields.io/pypi/pyversions/pydantic.svg)](https://github.com/pydantic/pydantic) +[![license](https://img.shields.io/github/license/pydantic/pydantic.svg)](https://github.com/pydantic/pydantic/blob/main/LICENSE) +[![Pydantic v2](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/pydantic/pydantic/main/docs/badge/v2.json)](https://docs.pydantic.dev/latest/contributing/#badges) +[![llms.txt](https://img.shields.io/badge/llms.txt-green)](https://docs.pydantic.dev/latest/llms.txt) + + +Data validation using Python type hints. + +Fast and extensible, Pydantic plays nicely with your linters/IDE/brain. +Define how data should be in pure, canonical Python 3.9+; validate it with Pydantic. + +## Pydantic Logfire :fire: + +We've recently launched Pydantic Logfire to help you monitor your applications. +[Learn more](https://pydantic.dev/articles/logfire-announcement) + +## Pydantic V1.10 vs. V2 + +Pydantic V2 is a ground-up rewrite that offers many new features, performance improvements, and some breaking changes compared to Pydantic V1. + +If you're using Pydantic V1 you may want to look at the +[pydantic V1.10 Documentation](https://docs.pydantic.dev/) or, +[`1.10.X-fixes` git branch](https://github.com/pydantic/pydantic/tree/1.10.X-fixes). Pydantic V2 also ships with the latest version of Pydantic V1 built in so that you can incrementally upgrade your code base and projects: `from pydantic import v1 as pydantic_v1`. + +## Help + +See [documentation](https://docs.pydantic.dev/) for more details. + +## Installation + +Install using `pip install -U pydantic` or `conda install pydantic -c conda-forge`. +For more installation options to make Pydantic even faster, +see the [Install](https://docs.pydantic.dev/install/) section in the documentation. + +## A Simple Example + +```python +from datetime import datetime +from typing import Optional +from pydantic import BaseModel + +class User(BaseModel): + id: int + name: str = 'John Doe' + signup_ts: Optional[datetime] = None + friends: list[int] = [] + +external_data = {'id': '123', 'signup_ts': '2017-06-01 12:22', 'friends': [1, '2', b'3']} +user = User(**external_data) +print(user) +#> User id=123 name='John Doe' signup_ts=datetime.datetime(2017, 6, 1, 12, 22) friends=[1, 2, 3] +print(user.id) +#> 123 +``` + +## Contributing + +For guidance on setting up a development environment and how to make a +contribution to Pydantic, see +[Contributing to Pydantic](https://docs.pydantic.dev/contributing/). + +## Reporting a Security Vulnerability + +See our [security policy](https://github.com/pydantic/pydantic/security/policy). + +## Changelog + +## v2.11.10 (2025-10-04) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.10) + +### What's Changed + +#### Fixes + +* Backport v1.10.24 changes by [@Viicos](https://github.com/Viicos) + +## v2.11.9 (2025-09-13) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.9) + +### What's Changed + +#### Fixes + +* Backport v1.10.23 changes by [@Viicos](https://github.com/Viicos) + +## v2.11.8 (2025-09-13) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.8) + +### What's Changed + +#### Fixes + +* Fix mypy plugin for mypy 1.18 by [@cdce8p](https://github.com/cdce8p) in [#12209](https://github.com/pydantic/pydantic/pull/12209) + +## v2.11.7 (2025-06-14) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.7) + +### What's Changed + +#### Fixes + +* Copy `FieldInfo` instance if necessary during `FieldInfo` build by [@Viicos](https://github.com/Viicos) in [#11898](https://github.com/pydantic/pydantic/pull/11898) + +## v2.11.6 (2025-06-13) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.6) + +### What's Changed + +#### Fixes + +* Rebuild dataclass fields before schema generation by [@Viicos](https://github.com/Viicos) in [#11949](https://github.com/pydantic/pydantic/pull/11949) +* Always store the original field assignment on `FieldInfo` by [@Viicos](https://github.com/Viicos) in [#11946](https://github.com/pydantic/pydantic/pull/11946) + +## v2.11.5 (2025-05-22) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.5) + +### What's Changed + +#### Fixes + +* Check if `FieldInfo` is complete after applying type variable map by [@Viicos](https://github.com/Viicos) in [#11855](https://github.com/pydantic/pydantic/pull/11855) +* Do not delete mock validator/serializer in `model_rebuild()` by [@Viicos](https://github.com/Viicos) in [#11890](https://github.com/pydantic/pydantic/pull/11890) +* Do not duplicate metadata on model rebuild by [@Viicos](https://github.com/Viicos) in [#11902](https://github.com/pydantic/pydantic/pull/11902) + +## v2.11.4 (2025-04-29) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.4) + +### What's Changed + +#### Packaging + +* Bump `mkdocs-llmstxt` to v0.2.0 by [@Viicos](https://github.com/Viicos) in [#11725](https://github.com/pydantic/pydantic/pull/11725) + +#### Changes + +* Allow config and bases to be specified together in `create_model()` by [@Viicos](https://github.com/Viicos) in [#11714](https://github.com/pydantic/pydantic/pull/11714). + This change was backported as it was previously possible (although not meant to be supported) + to provide `model_config` as a field, which would make it possible to provide both configuration + and bases. + +#### Fixes + +* Remove generics cache workaround by [@Viicos](https://github.com/Viicos) in [#11755](https://github.com/pydantic/pydantic/pull/11755) +* Remove coercion of decimal constraints by [@Viicos](https://github.com/Viicos) in [#11772](https://github.com/pydantic/pydantic/pull/11772) +* Fix crash when expanding root type in the mypy plugin by [@Viicos](https://github.com/Viicos) in [#11735](https://github.com/pydantic/pydantic/pull/11735) +* Fix issue with recursive generic models by [@Viicos](https://github.com/Viicos) in [#11775](https://github.com/pydantic/pydantic/pull/11775) +* Traverse `function-before` schemas during schema gathering by [@Viicos](https://github.com/Viicos) in [#11801](https://github.com/pydantic/pydantic/pull/11801) + +## v2.11.3 (2025-04-08) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.3) + +### What's Changed + +#### Packaging + +* Update V1 copy to v1.10.21 by [@Viicos](https://github.com/Viicos) in [#11706](https://github.com/pydantic/pydantic/pull/11706) + +#### Fixes + +* Preserve field description when rebuilding model fields by [@Viicos](https://github.com/Viicos) in [#11698](https://github.com/pydantic/pydantic/pull/11698) + +## v2.11.2 (2025-04-03) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.2) + +### What's Changed + +#### Fixes + +* Bump `pydantic-core` to v2.33.1 by [@Viicos](https://github.com/Viicos) in [#11678](https://github.com/pydantic/pydantic/pull/11678) +* Make sure `__pydantic_private__` exists before setting private attributes by [@Viicos](https://github.com/Viicos) in [#11666](https://github.com/pydantic/pydantic/pull/11666) +* Do not override `FieldInfo._complete` when using field from parent class by [@Viicos](https://github.com/Viicos) in [#11668](https://github.com/pydantic/pydantic/pull/11668) +* Provide the available definitions when applying discriminated unions by [@Viicos](https://github.com/Viicos) in [#11670](https://github.com/pydantic/pydantic/pull/11670) +* Do not expand root type in the mypy plugin for variables by [@Viicos](https://github.com/Viicos) in [#11676](https://github.com/pydantic/pydantic/pull/11676) +* Mention the attribute name in model fields deprecation message by [@Viicos](https://github.com/Viicos) in [#11674](https://github.com/pydantic/pydantic/pull/11674) +* Properly validate parameterized mappings by [@Viicos](https://github.com/Viicos) in [#11658](https://github.com/pydantic/pydantic/pull/11658) + +## v2.11.1 (2025-03-28) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.1) + +### What's Changed + +#### Fixes + +* Do not override `'definitions-ref'` schemas containing serialization schemas or metadata by [@Viicos](https://github.com/Viicos) in [#11644](https://github.com/pydantic/pydantic/pull/11644) + +## v2.11.0 (2025-03-27) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.0) + +### What's Changed + +Pydantic v2.11 is a version strongly focused on build time performance of Pydantic models (and core schema generation in general). +See the [blog post](https://pydantic.dev/articles/pydantic-v2-11-release) for more details. + +#### Packaging + +* Bump `pydantic-core` to v2.33.0 by [@Viicos](https://github.com/Viicos) in [#11631](https://github.com/pydantic/pydantic/pull/11631) + +#### New Features + +* Add `encoded_string()` method to the URL types by [@YassinNouh21](https://github.com/YassinNouh21) in [#11580](https://github.com/pydantic/pydantic/pull/11580) +* Add support for `defer_build` with `@validate_call` decorator by [@Viicos](https://github.com/Viicos) in [#11584](https://github.com/pydantic/pydantic/pull/11584) +* Allow `@with_config` decorator to be used with keyword arguments by [@Viicos](https://github.com/Viicos) in [#11608](https://github.com/pydantic/pydantic/pull/11608) +* Simplify customization of default value inclusion in JSON Schema generation by [@Viicos](https://github.com/Viicos) in [#11634](https://github.com/pydantic/pydantic/pull/11634) +* Add `generate_arguments_schema()` function by [@Viicos](https://github.com/Viicos) in [#11572](https://github.com/pydantic/pydantic/pull/11572) + +#### Fixes + +* Allow generic typed dictionaries to be used for unpacked variadic keyword parameters by [@Viicos](https://github.com/Viicos) in [#11571](https://github.com/pydantic/pydantic/pull/11571) +* Fix runtime error when computing model string representation involving cached properties and self-referenced models by [@Viicos](https://github.com/Viicos) in [#11579](https://github.com/pydantic/pydantic/pull/11579) +* Preserve other steps when using the ellipsis in the pipeline API by [@Viicos](https://github.com/Viicos) in [#11626](https://github.com/pydantic/pydantic/pull/11626) +* Fix deferred discriminator application logic by [@Viicos](https://github.com/Viicos) in [#11591](https://github.com/pydantic/pydantic/pull/11591) + +### New Contributors + +* [@cmenon12](https://github.com/cmenon12) made their first contribution in [#11562](https://github.com/pydantic/pydantic/pull/11562) +* [@Jeukoh](https://github.com/Jeukoh) made their first contribution in [#11611](https://github.com/pydantic/pydantic/pull/11611) + +## v2.11.0b2 (2025-03-17) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.0b2) + +### What's Changed + +#### Packaging + +* Bump `pydantic-core` to v2.32.0 by [@Viicos](https://github.com/Viicos) in [#11567](https://github.com/pydantic/pydantic/pull/11567) + +#### New Features + +* Add experimental support for free threading by [@Viicos](https://github.com/Viicos) in [#11516](https://github.com/pydantic/pydantic/pull/11516) + +#### Fixes + +* Fix `NotRequired` qualifier not taken into account in stringified annotation by [@Viicos](https://github.com/Viicos) in [#11559](https://github.com/pydantic/pydantic/pull/11559) + +### New Contributors + +* [@joren485](https://github.com/joren485) made their first contribution in [#11547](https://github.com/pydantic/pydantic/pull/11547) + +## v2.11.0b1 (2025-03-06) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.0b1) + +### What's Changed + +#### Packaging + +* Add a `check_pydantic_core_version()` function by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11324 +* Remove `greenlet` development dependency by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11351 +* Use the `typing-inspection` library by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11479 +* Bump `pydantic-core` to `v2.31.1` by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11526 + +#### New Features + +* Support unsubstituted type variables with both a default and a bound or constraints by [@FyZzyss](https://github.com/FyZzyss) in https://github.com/pydantic/pydantic/pull/10789 +* Add a `default_factory_takes_validated_data` property to `FieldInfo` by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11034 +* Raise a better error when a generic alias is used inside `type[]` by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11088 +* Properly support PEP 695 generics syntax by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11189 +* Properly support type variable defaults by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11332 +* Add support for validating v6, v7, v8 UUIDs by [@astei](https://github.com/astei) in https://github.com/pydantic/pydantic/pull/11436 +* Improve alias configuration APIs by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11468 + +#### Changes + +* Rework `create_model` field definitions format by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11032 +* Raise a deprecation warning when a field is annotated as final with a default value by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11168 +* Deprecate accessing `model_fields` and `model_computed_fields` on instances by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11169 +* **Breaking Change:** Move core schema generation logic for path types inside the `GenerateSchema` class by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/10846 +* Remove Python 3.8 Support by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11258 +* Optimize calls to `get_type_ref` by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/10863 +* Disable `pydantic-core` core schema validation by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11271 + +#### Performance + +* Only evaluate `FieldInfo` annotations if required during schema building by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/10769 +* Improve `__setattr__` performance of Pydantic models by caching setter functions by [@MarkusSintonen](https://github.com/MarkusSintonen) in https://github.com/pydantic/pydantic/pull/10868 +* Improve annotation application performance by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11186 +* Improve performance of `_typing_extra` module by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11255 +* Refactor and optimize schema cleaning logic by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11244 +* Create a single dictionary when creating a `CoreConfig` instance by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11384 +* Bump `pydantic-core` and thus use `SchemaValidator` and `SchemaSerializer` caching by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11402 +* Reuse cached core schemas for parametrized generic Pydantic models by [@MarkusSintonen](https://github.com/MarkusSintonen) in https://github.com/pydantic/pydantic/pull/11434 + +#### Fixes + +* Improve `TypeAdapter` instance repr by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/10872 +* Use the correct frame when instantiating a parametrized `TypeAdapter` by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/10893 +* Infer final fields with a default value as class variables in the mypy plugin by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11121 +* Recursively unpack `Literal` values if using PEP 695 type aliases by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11114 +* Override `__subclasscheck__` on `ModelMetaclass` to avoid memory leak and performance issues by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11116 +* Remove unused `_extract_get_pydantic_json_schema()` parameter by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11155 +* Improve discriminated union error message for invalid union variants by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11161 +* Unpack PEP 695 type aliases if using the `Annotated` form by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11109 +* Add missing stacklevel in `deprecated_instance_property` warning by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11200 +* Copy `WithJsonSchema` schema to avoid sharing mutated data by [@thejcannon](https://github.com/thejcannon) in https://github.com/pydantic/pydantic/pull/11014 +* Do not cache parametrized models when in the process of parametrizing another model by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/10704 +* Add discriminated union related metadata entries to the `CoreMetadata` definition by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11216 +* Consolidate schema definitions logic in the `_Definitions` class by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11208 +* Support initializing root model fields with values of the `root` type in the mypy plugin by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11212 +* Fix various issues with dataclasses and `use_attribute_docstrings` by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11246 +* Only compute normalized decimal places if necessary in `decimal_places_validator` by [@misrasaurabh1](https://github.com/misrasaurabh1) in https://github.com/pydantic/pydantic/pull/11281 +* Add support for `validation_alias` in the mypy plugin by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11295 +* Fix JSON Schema reference collection with `"examples"` keys by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11305 +* Do not transform model serializer functions as class methods in the mypy plugin by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11298 +* Simplify `GenerateJsonSchema.literal_schema()` implementation by [@misrasaurabh1](https://github.com/misrasaurabh1) in https://github.com/pydantic/pydantic/pull/11321 +* Add additional allowed schemes for `ClickHouseDsn` by [@Maze21127](https://github.com/Maze21127) in https://github.com/pydantic/pydantic/pull/11319 +* Coerce decimal constraints to `Decimal` instances by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11350 +* Use the correct JSON Schema mode when handling function schemas by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11367 +* Improve exception message when encountering recursion errors during type evaluation by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11356 +* Always include `additionalProperties: True` for arbitrary dictionary schemas by [@austinyu](https://github.com/austinyu) in https://github.com/pydantic/pydantic/pull/11392 +* Expose `fallback` parameter in serialization methods by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11398 +* Fix path serialization behavior by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11416 +* Do not reuse validators and serializers during model rebuild by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11429 +* Collect model fields when rebuilding a model by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11388 +* Allow cached properties to be altered on frozen models by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11432 +* Fix tuple serialization for `Sequence` types by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11435 +* Fix: do not check for `__get_validators__` on classes where `__get_pydantic_core_schema__` is also defined by [@tlambert03](https://github.com/tlambert03) in https://github.com/pydantic/pydantic/pull/11444 +* Allow callable instances to be used as serializers by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11451 +* Improve error thrown when overriding field with a property by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11459 +* Fix JSON Schema generation with referenceable core schemas holding JSON metadata by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11475 +* Support strict specification on union member types by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11481 +* Implicitly set `validate_by_name` to `True` when `validate_by_alias` is `False` by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic/pull/11503 +* Change type of `Any` when synthesizing `BaseSettings.__init__` signature in the mypy plugin by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11497 +* Support type variable defaults referencing other type variables by [@Viicos](https://github.com/Viicos) in https://github.com/pydantic/pydantic/pull/11520 +* Fix `ValueError` on year zero by [@davidhewitt](https://github.com/davidhewitt) in https://github.com/pydantic/pydantic-core/pull/1583 +* `dataclass` `InitVar` shouldn't be required on serialization by [@sydney-runkle](https://github.com/sydney-runkle) in https://github.com/pydantic/pydantic-core/pull/1602 + +## New Contributors +* [@FyZzyss](https://github.com/FyZzyss) made their first contribution in https://github.com/pydantic/pydantic/pull/10789 +* [@tamird](https://github.com/tamird) made their first contribution in https://github.com/pydantic/pydantic/pull/10948 +* [@felixxm](https://github.com/felixxm) made their first contribution in https://github.com/pydantic/pydantic/pull/11077 +* [@alexprabhat99](https://github.com/alexprabhat99) made their first contribution in https://github.com/pydantic/pydantic/pull/11082 +* [@Kharianne](https://github.com/Kharianne) made their first contribution in https://github.com/pydantic/pydantic/pull/11111 +* [@mdaffad](https://github.com/mdaffad) made their first contribution in https://github.com/pydantic/pydantic/pull/11177 +* [@thejcannon](https://github.com/thejcannon) made their first contribution in https://github.com/pydantic/pydantic/pull/11014 +* [@thomasfrimannkoren](https://github.com/thomasfrimannkoren) made their first contribution in https://github.com/pydantic/pydantic/pull/11251 +* [@usernameMAI](https://github.com/usernameMAI) made their first contribution in https://github.com/pydantic/pydantic/pull/11275 +* [@ananiavito](https://github.com/ananiavito) made their first contribution in https://github.com/pydantic/pydantic/pull/11302 +* [@pawamoy](https://github.com/pawamoy) made their first contribution in https://github.com/pydantic/pydantic/pull/11311 +* [@Maze21127](https://github.com/Maze21127) made their first contribution in https://github.com/pydantic/pydantic/pull/11319 +* [@kauabh](https://github.com/kauabh) made their first contribution in https://github.com/pydantic/pydantic/pull/11369 +* [@jaceklaskowski](https://github.com/jaceklaskowski) made their first contribution in https://github.com/pydantic/pydantic/pull/11353 +* [@tmpbeing](https://github.com/tmpbeing) made their first contribution in https://github.com/pydantic/pydantic/pull/11375 +* [@petyosi](https://github.com/petyosi) made their first contribution in https://github.com/pydantic/pydantic/pull/11405 +* [@austinyu](https://github.com/austinyu) made their first contribution in https://github.com/pydantic/pydantic/pull/11392 +* [@mikeedjones](https://github.com/mikeedjones) made their first contribution in https://github.com/pydantic/pydantic/pull/11402 +* [@astei](https://github.com/astei) made their first contribution in https://github.com/pydantic/pydantic/pull/11436 +* [@dsayling](https://github.com/dsayling) made their first contribution in https://github.com/pydantic/pydantic/pull/11522 +* [@sobolevn](https://github.com/sobolevn) made their first contribution in https://github.com/pydantic/pydantic-core/pull/1645 + +## v2.11.0a2 (2025-02-10) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.0a2) + +### What's Changed + +Pydantic v2.11 is a version strongly focused on build time performance of Pydantic models (and core schema generation in general). +This is another early alpha release, meant to collect early feedback from users having issues with core schema builds. + +#### Packaging + +* Bump `ruff` from 0.9.2 to 0.9.5 by [@Viicos](https://github.com/Viicos) in [#11407](https://github.com/pydantic/pydantic/pull/11407) +* Bump `pydantic-core` to v2.29.0 by [@mikeedjones](https://github.com/mikeedjones) in [#11402](https://github.com/pydantic/pydantic/pull/11402) +* Use locally-built rust with symbols & pgo by [@davidhewitt](https://github.com/davidhewitt) in [#11403](https://github.com/pydantic/pydantic/pull/11403) + + +#### Performance + +* Create a single dictionary when creating a `CoreConfig` instance by [@sydney-runkle](https://github.com/sydney-runkle) in [#11384](https://github.com/pydantic/pydantic/pull/11384) + +#### Fixes + +* Use the correct JSON Schema mode when handling function schemas by [@Viicos](https://github.com/Viicos) in [#11367](https://github.com/pydantic/pydantic/pull/11367) +* Fix JSON Schema reference logic with `examples` keys by [@Viicos](https://github.com/Viicos) in [#11366](https://github.com/pydantic/pydantic/pull/11366) +* Improve exception message when encountering recursion errors during type evaluation by [@Viicos](https://github.com/Viicos) in [#11356](https://github.com/pydantic/pydantic/pull/11356) +* Always include `additionalProperties: True` for arbitrary dictionary schemas by [@austinyu](https://github.com/austinyu) in [#11392](https://github.com/pydantic/pydantic/pull/11392) +* Expose `fallback` parameter in serialization methods by [@Viicos](https://github.com/Viicos) in [#11398](https://github.com/pydantic/pydantic/pull/11398) +* Fix path serialization behavior by [@sydney-runkle](https://github.com/sydney-runkle) in [#11416](https://github.com/pydantic/pydantic/pull/11416) + +### New Contributors + +* [@kauabh](https://github.com/kauabh) made their first contribution in [#11369](https://github.com/pydantic/pydantic/pull/11369) +* [@jaceklaskowski](https://github.com/jaceklaskowski) made their first contribution in [#11353](https://github.com/pydantic/pydantic/pull/11353) +* [@tmpbeing](https://github.com/tmpbeing) made their first contribution in [#11375](https://github.com/pydantic/pydantic/pull/11375) +* [@petyosi](https://github.com/petyosi) made their first contribution in [#11405](https://github.com/pydantic/pydantic/pull/11405) +* [@austinyu](https://github.com/austinyu) made their first contribution in [#11392](https://github.com/pydantic/pydantic/pull/11392) +* [@mikeedjones](https://github.com/mikeedjones) made their first contribution in [#11402](https://github.com/pydantic/pydantic/pull/11402) + +## v2.11.0a1 (2025-01-30) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.11.0a1) + +### What's Changed + +Pydantic v2.11 is a version strongly focused on build time performance of Pydantic models (and core schema generation in general). +This is an early alpha release, meant to collect early feedback from users having issues with core schema builds. + +#### Packaging + +* Bump dawidd6/action-download-artifact from 6 to 7 by [@dependabot](https://github.com/dependabot) in [#11018](https://github.com/pydantic/pydantic/pull/11018) +* Re-enable memray related tests on Python 3.12+ by [@Viicos](https://github.com/Viicos) in [#11191](https://github.com/pydantic/pydantic/pull/11191) +* Bump astral-sh/setup-uv to 5 by [@dependabot](https://github.com/dependabot) in [#11205](https://github.com/pydantic/pydantic/pull/11205) +* Bump `ruff` to v0.9.0 by [@sydney-runkle](https://github.com/sydney-runkle) in [#11254](https://github.com/pydantic/pydantic/pull/11254) +* Regular `uv.lock` deps update by [@sydney-runkle](https://github.com/sydney-runkle) in [#11333](https://github.com/pydantic/pydantic/pull/11333) +* Add a `check_pydantic_core_version()` function by [@Viicos](https://github.com/Viicos) in [#11324](https://github.com/pydantic/pydantic/pull/11324) +* Remove `greenlet` development dependency by [@Viicos](https://github.com/Viicos) in [#11351](https://github.com/pydantic/pydantic/pull/11351) +* Bump `pydantic-core` to v2.28.0 by [@Viicos](https://github.com/Viicos) in [#11364](https://github.com/pydantic/pydantic/pull/11364) + +#### New Features + +* Support unsubstituted type variables with both a default and a bound or constraints by [@FyZzyss](https://github.com/FyZzyss) in [#10789](https://github.com/pydantic/pydantic/pull/10789) +* Add a `default_factory_takes_validated_data` property to `FieldInfo` by [@Viicos](https://github.com/Viicos) in [#11034](https://github.com/pydantic/pydantic/pull/11034) +* Raise a better error when a generic alias is used inside `type[]` by [@Viicos](https://github.com/Viicos) in [#11088](https://github.com/pydantic/pydantic/pull/11088) +* Properly support PEP 695 generics syntax by [@Viicos](https://github.com/Viicos) in [#11189](https://github.com/pydantic/pydantic/pull/11189) +* Properly support type variable defaults by [@Viicos](https://github.com/Viicos) in [#11332](https://github.com/pydantic/pydantic/pull/11332) + +#### Changes + +* Rework `create_model` field definitions format by [@Viicos](https://github.com/Viicos) in [#11032](https://github.com/pydantic/pydantic/pull/11032) +* Raise a deprecation warning when a field is annotated as final with a default value by [@Viicos](https://github.com/Viicos) in [#11168](https://github.com/pydantic/pydantic/pull/11168) +* Deprecate accessing `model_fields` and `model_computed_fields` on instances by [@Viicos](https://github.com/Viicos) in [#11169](https://github.com/pydantic/pydantic/pull/11169) +* Move core schema generation logic for path types inside the `GenerateSchema` class by [@sydney-runkle](https://github.com/sydney-runkle) in [#10846](https://github.com/pydantic/pydantic/pull/10846) +* Move `deque` schema gen to `GenerateSchema` class by [@sydney-runkle](https://github.com/sydney-runkle) in [#11239](https://github.com/pydantic/pydantic/pull/11239) +* Move `Mapping` schema gen to `GenerateSchema` to complete removal of `prepare_annotations_for_known_type` workaround by [@sydney-runkle](https://github.com/sydney-runkle) in [#11247](https://github.com/pydantic/pydantic/pull/11247) +* Remove Python 3.8 Support by [@sydney-runkle](https://github.com/sydney-runkle) in [#11258](https://github.com/pydantic/pydantic/pull/11258) +* Disable `pydantic-core` core schema validation by [@sydney-runkle](https://github.com/sydney-runkle) in [#11271](https://github.com/pydantic/pydantic/pull/11271) + +#### Performance + +* Only evaluate `FieldInfo` annotations if required during schema building by [@Viicos](https://github.com/Viicos) in [#10769](https://github.com/pydantic/pydantic/pull/10769) +* Optimize calls to `get_type_ref` by [@Viicos](https://github.com/Viicos) in [#10863](https://github.com/pydantic/pydantic/pull/10863) +* Improve `__setattr__` performance of Pydantic models by caching setter functions by [@MarkusSintonen](https://github.com/MarkusSintonen) in [#10868](https://github.com/pydantic/pydantic/pull/10868) +* Improve annotation application performance by [@Viicos](https://github.com/Viicos) in [#11186](https://github.com/pydantic/pydantic/pull/11186) +* Improve performance of `_typing_extra` module by [@Viicos](https://github.com/Viicos) in [#11255](https://github.com/pydantic/pydantic/pull/11255) +* Refactor and optimize schema cleaning logic by [@Viicos](https://github.com/Viicos) and [@MarkusSintonen](https://github.com/MarkusSintonen) in [#11244](https://github.com/pydantic/pydantic/pull/11244) + +#### Fixes + +* Add validation tests for `_internal/_validators.py` by [@tkasuz](https://github.com/tkasuz) in [#10763](https://github.com/pydantic/pydantic/pull/10763) +* Improve `TypeAdapter` instance repr by [@sydney-runkle](https://github.com/sydney-runkle) in [#10872](https://github.com/pydantic/pydantic/pull/10872) +* Revert "ci: use locally built pydantic-core with debug symbols by [@sydney-runkle](https://github.com/sydney-runkle) in [#10942](https://github.com/pydantic/pydantic/pull/10942) +* Re-enable all FastAPI tests by [@tamird](https://github.com/tamird) in [#10948](https://github.com/pydantic/pydantic/pull/10948) +* Fix typo in HISTORY.md. by [@felixxm](https://github.com/felixxm) in [#11077](https://github.com/pydantic/pydantic/pull/11077) +* Infer final fields with a default value as class variables in the mypy plugin by [@Viicos](https://github.com/Viicos) in [#11121](https://github.com/pydantic/pydantic/pull/11121) +* Recursively unpack `Literal` values if using PEP 695 type aliases by [@Viicos](https://github.com/Viicos) in [#11114](https://github.com/pydantic/pydantic/pull/11114) +* Override `__subclasscheck__` on `ModelMetaclass` to avoid memory leak and performance issues by [@Viicos](https://github.com/Viicos) in [#11116](https://github.com/pydantic/pydantic/pull/11116) +* Remove unused `_extract_get_pydantic_json_schema()` parameter by [@Viicos](https://github.com/Viicos) in [#11155](https://github.com/pydantic/pydantic/pull/11155) +* Add FastAPI and SQLModel to third-party tests by [@sydney-runkle](https://github.com/sydney-runkle) in [#11044](https://github.com/pydantic/pydantic/pull/11044) +* Fix conditional expressions syntax for third-party tests by [@Viicos](https://github.com/Viicos) in [#11162](https://github.com/pydantic/pydantic/pull/11162) +* Move FastAPI tests to third-party workflow by [@Viicos](https://github.com/Viicos) in [#11164](https://github.com/pydantic/pydantic/pull/11164) +* Improve discriminated union error message for invalid union variants by [@Viicos](https://github.com/Viicos) in [#11161](https://github.com/pydantic/pydantic/pull/11161) +* Unpack PEP 695 type aliases if using the `Annotated` form by [@Viicos](https://github.com/Viicos) in [#11109](https://github.com/pydantic/pydantic/pull/11109) +* Include `openapi-python-client` check in issue creation for third-party failures, use `main` branch by [@sydney-runkle](https://github.com/sydney-runkle) in [#11182](https://github.com/pydantic/pydantic/pull/11182) +* Add pandera third-party tests by [@Viicos](https://github.com/Viicos) in [#11193](https://github.com/pydantic/pydantic/pull/11193) +* Add ODMantic third-party tests by [@sydney-runkle](https://github.com/sydney-runkle) in [#11197](https://github.com/pydantic/pydantic/pull/11197) +* Add missing stacklevel in `deprecated_instance_property` warning by [@Viicos](https://github.com/Viicos) in [#11200](https://github.com/pydantic/pydantic/pull/11200) +* Copy `WithJsonSchema` schema to avoid sharing mutated data by [@thejcannon](https://github.com/thejcannon) in [#11014](https://github.com/pydantic/pydantic/pull/11014) +* Do not cache parametrized models when in the process of parametrizing another model by [@Viicos](https://github.com/Viicos) in [#10704](https://github.com/pydantic/pydantic/pull/10704) +* Re-enable Beanie third-party tests by [@Viicos](https://github.com/Viicos) in [#11214](https://github.com/pydantic/pydantic/pull/11214) +* Add discriminated union related metadata entries to the `CoreMetadata` definition by [@Viicos](https://github.com/Viicos) in [#11216](https://github.com/pydantic/pydantic/pull/11216) +* Consolidate schema definitions logic in the `_Definitions` class by [@Viicos](https://github.com/Viicos) in [#11208](https://github.com/pydantic/pydantic/pull/11208) +* Support initializing root model fields with values of the `root` type in the mypy plugin by [@Viicos](https://github.com/Viicos) in [#11212](https://github.com/pydantic/pydantic/pull/11212) +* Fix various issues with dataclasses and `use_attribute_docstrings` by [@Viicos](https://github.com/Viicos) in [#11246](https://github.com/pydantic/pydantic/pull/11246) +* Only compute normalized decimal places if necessary in `decimal_places_validator` by [@misrasaurabh1](https://github.com/misrasaurabh1) in [#11281](https://github.com/pydantic/pydantic/pull/11281) +* Fix two misplaced sentences in validation errors documentation by [@ananiavito](https://github.com/ananiavito) in [#11302](https://github.com/pydantic/pydantic/pull/11302) +* Fix mkdocstrings inventory example in documentation by [@pawamoy](https://github.com/pawamoy) in [#11311](https://github.com/pydantic/pydantic/pull/11311) +* Add support for `validation_alias` in the mypy plugin by [@Viicos](https://github.com/Viicos) in [#11295](https://github.com/pydantic/pydantic/pull/11295) +* Do not transform model serializer functions as class methods in the mypy plugin by [@Viicos](https://github.com/Viicos) in [#11298](https://github.com/pydantic/pydantic/pull/11298) +* Simplify `GenerateJsonSchema.literal_schema()` implementation by [@misrasaurabh1](https://github.com/misrasaurabh1) in [#11321](https://github.com/pydantic/pydantic/pull/11321) +* Add additional allowed schemes for `ClickHouseDsn` by [@Maze21127](https://github.com/Maze21127) in [#11319](https://github.com/pydantic/pydantic/pull/11319) +* Coerce decimal constraints to `Decimal` instances by [@Viicos](https://github.com/Viicos) in [#11350](https://github.com/pydantic/pydantic/pull/11350) +* Fix `ValueError` on year zero by [@davidhewitt](https://github.com/davidhewitt) in [pydantic-core#1583](https://github.com/pydantic/pydantic-core/pull/1583) + +### New Contributors + +* [@FyZzyss](https://github.com/FyZzyss) made their first contribution in [#10789](https://github.com/pydantic/pydantic/pull/10789) +* [@tamird](https://github.com/tamird) made their first contribution in [#10948](https://github.com/pydantic/pydantic/pull/10948) +* [@felixxm](https://github.com/felixxm) made their first contribution in [#11077](https://github.com/pydantic/pydantic/pull/11077) +* [@alexprabhat99](https://github.com/alexprabhat99) made their first contribution in [#11082](https://github.com/pydantic/pydantic/pull/11082) +* [@Kharianne](https://github.com/Kharianne) made their first contribution in [#11111](https://github.com/pydantic/pydantic/pull/11111) +* [@mdaffad](https://github.com/mdaffad) made their first contribution in [#11177](https://github.com/pydantic/pydantic/pull/11177) +* [@thejcannon](https://github.com/thejcannon) made their first contribution in [#11014](https://github.com/pydantic/pydantic/pull/11014) +* [@thomasfrimannkoren](https://github.com/thomasfrimannkoren) made their first contribution in [#11251](https://github.com/pydantic/pydantic/pull/11251) +* [@usernameMAI](https://github.com/usernameMAI) made their first contribution in [#11275](https://github.com/pydantic/pydantic/pull/11275) +* [@ananiavito](https://github.com/ananiavito) made their first contribution in [#11302](https://github.com/pydantic/pydantic/pull/11302) +* [@pawamoy](https://github.com/pawamoy) made their first contribution in [#11311](https://github.com/pydantic/pydantic/pull/11311) +* [@Maze21127](https://github.com/Maze21127) made their first contribution in [#11319](https://github.com/pydantic/pydantic/pull/11319) + +## v2.10.6 (2025-01-23) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.6) + +### What's Changed + +#### Fixes + +* Fix JSON Schema reference collection with `'examples'` keys by [@Viicos](https://github.com/Viicos) in [#11325](https://github.com/pydantic/pydantic/pull/11325) +* Fix url python serialization by [@sydney-runkle](https://github.com/sydney-runkle) in [#11331](https://github.com/pydantic/pydantic/pull/11331) + +## v2.10.5 (2025-01-08) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.5) + +### What's Changed + +#### Fixes + +* Remove custom MRO implementation of Pydantic models by [@Viicos](https://github.com/Viicos) in [#11184](https://github.com/pydantic/pydantic/pull/11184) +* Fix URL serialization for unions by [@sydney-runkle](https://github.com/sydney-runkle) in [#11233](https://github.com/pydantic/pydantic/pull/11233) + +## v2.10.4 (2024-12-18) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.4) + +### What's Changed + +#### Packaging + +* Bump `pydantic-core` to v2.27.2 by [@davidhewitt](https://github.com/davidhewitt) in [#11138](https://github.com/pydantic/pydantic/pull/11138) + +#### Fixes + +* Fix for comparison of `AnyUrl` objects by [@alexprabhat99](https://github.com/alexprabhat99) in [#11082](https://github.com/pydantic/pydantic/pull/11082) +* Properly fetch PEP 695 type params for functions, do not fetch annotations from signature by [@Viicos](https://github.com/Viicos) in [#11093](https://github.com/pydantic/pydantic/pull/11093) +* Include JSON Schema input core schema in function schemas by [@Viicos](https://github.com/Viicos) in [#11085](https://github.com/pydantic/pydantic/pull/11085) +* Add `len` to `_BaseUrl` to avoid TypeError by [@Kharianne](https://github.com/Kharianne) in [#11111](https://github.com/pydantic/pydantic/pull/11111) +* Make sure the type reference is removed from the seen references by [@Viicos](https://github.com/Viicos) in [#11143](https://github.com/pydantic/pydantic/pull/11143) + +### New Contributors + +* [@FyZzyss](https://github.com/FyZzyss) made their first contribution in [#10789](https://github.com/pydantic/pydantic/pull/10789) +* [@tamird](https://github.com/tamird) made their first contribution in [#10948](https://github.com/pydantic/pydantic/pull/10948) +* [@felixxm](https://github.com/felixxm) made their first contribution in [#11077](https://github.com/pydantic/pydantic/pull/11077) +* [@alexprabhat99](https://github.com/alexprabhat99) made their first contribution in [#11082](https://github.com/pydantic/pydantic/pull/11082) +* [@Kharianne](https://github.com/Kharianne) made their first contribution in [#11111](https://github.com/pydantic/pydantic/pull/11111) + +## v2.10.3 (2024-12-03) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.3) + +### What's Changed + +#### Fixes + +* Set fields when `defer_build` is set on Pydantic dataclasses by [@Viicos](https://github.com/Viicos) in [#10984](https://github.com/pydantic/pydantic/pull/10984) +* Do not resolve the JSON Schema reference for `dict` core schema keys by [@Viicos](https://github.com/Viicos) in [#10989](https://github.com/pydantic/pydantic/pull/10989) +* Use the globals of the function when evaluating the return type for `PlainSerializer` and `WrapSerializer` functions by [@Viicos](https://github.com/Viicos) in [#11008](https://github.com/pydantic/pydantic/pull/11008) +* Fix host required enforcement for urls to be compatible with v2.9 behavior by [@sydney-runkle](https://github.com/sydney-runkle) in [#11027](https://github.com/pydantic/pydantic/pull/11027) +* Add a `default_factory_takes_validated_data` property to `FieldInfo` by [@Viicos](https://github.com/Viicos) in [#11034](https://github.com/pydantic/pydantic/pull/11034) +* Fix url json schema in `serialization` mode by [@sydney-runkle](https://github.com/sydney-runkle) in [#11035](https://github.com/pydantic/pydantic/pull/11035) + +## v2.10.2 (2024-11-25) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.2) + +### What's Changed + +#### Fixes + +* Only evaluate FieldInfo annotations if required during schema building by [@Viicos](https://github.com/Viicos) in [#10769](https://github.com/pydantic/pydantic/pull/10769) +* Do not evaluate annotations for private fields by [@Viicos](https://github.com/Viicos) in [#10962](https://github.com/pydantic/pydantic/pull/10962) +* Support serialization as any for `Secret` types and `Url` types by [@sydney-runkle](https://github.com/sydney-runkle) in [#10947](https://github.com/pydantic/pydantic/pull/10947) +* Fix type hint of `Field.default` to be compatible with Python 3.8 and 3.9 by [@Viicos](https://github.com/Viicos) in [#10972](https://github.com/pydantic/pydantic/pull/10972) +* Add hashing support for URL types by [@sydney-runkle](https://github.com/sydney-runkle) in [#10975](https://github.com/pydantic/pydantic/pull/10975) +* Hide `BaseModel.__replace__` definition from type checkers by [@Viicos](https://github.com/Viicos) in [#10979](https://github.com/pydantic/pydantic/pull/10979) + +## v2.10.1 (2024-11-21) + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.1) + +### What's Changed + +#### Packaging + +* Bump `pydantic-core` version to `v2.27.1` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10938](https://github.com/pydantic/pydantic/pull/10938) + +#### Fixes + +* Use the correct frame when instantiating a parametrized `TypeAdapter` by [@Viicos](https://github.com/Viicos) in [#10893](https://github.com/pydantic/pydantic/pull/10893) +* Relax check for validated data in `default_factory` utils by [@sydney-runkle](https://github.com/sydney-runkle) in [#10909](https://github.com/pydantic/pydantic/pull/10909) +* Fix type checking issue with `model_fields` and `model_computed_fields` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10911](https://github.com/pydantic/pydantic/pull/10911) +* Use the parent configuration during schema generation for stdlib `dataclass`es by [@sydney-runkle](https://github.com/sydney-runkle) in [#10928](https://github.com/pydantic/pydantic/pull/10928) +* Use the `globals` of the function when evaluating the return type of serializers and `computed_field`s by [@Viicos](https://github.com/Viicos) in [#10929](https://github.com/pydantic/pydantic/pull/10929) +* Fix URL constraint application by [@sydney-runkle](https://github.com/sydney-runkle) in [#10922](https://github.com/pydantic/pydantic/pull/10922) +* Fix URL equality with different validation methods by [@sydney-runkle](https://github.com/sydney-runkle) in [#10934](https://github.com/pydantic/pydantic/pull/10934) +* Fix JSON schema title when specified as `''` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10936](https://github.com/pydantic/pydantic/pull/10936) +* Fix `python` mode serialization for `complex` inference by [@sydney-runkle](https://github.com/sydney-runkle) in [pydantic-core#1549](https://github.com/pydantic/pydantic-core/pull/1549) + +### New Contributors + +## v2.10.0 (2024-11-20) + +The code released in v2.10.0 is practically identical to that of v2.10.0b2. + +[GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.0) + +See the [v2.10 release blog post](https://pydantic.dev/articles/pydantic-v2-10-release) for the highlights! + +### What's Changed + +#### Packaging + +* Bump `pydantic-core` to `v2.27.0` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10825](https://github.com/pydantic/pydantic/pull/10825) +* Replaced pdm with uv by [@frfahim](https://github.com/frfahim) in [#10727](https://github.com/pydantic/pydantic/pull/10727) + +#### New Features + +* Support `fractions.Fraction` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10318](https://github.com/pydantic/pydantic/pull/10318) +* Support `Hashable` for json validation by [@sydney-runkle](https://github.com/sydney-runkle) in [#10324](https://github.com/pydantic/pydantic/pull/10324) +* Add a `SocketPath` type for `linux` systems by [@theunkn0wn1](https://github.com/theunkn0wn1) in [#10378](https://github.com/pydantic/pydantic/pull/10378) +* Allow arbitrary refs in JSON schema `examples` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10417](https://github.com/pydantic/pydantic/pull/10417) +* Support `defer_build` for Pydantic dataclasses by [@Viicos](https://github.com/Viicos) in [#10313](https://github.com/pydantic/pydantic/pull/10313) +* Adding v1 / v2 incompatibility warning for nested v1 model by [@sydney-runkle](https://github.com/sydney-runkle) in [#10431](https://github.com/pydantic/pydantic/pull/10431) +* Add support for unpacked `TypedDict` to type hint variadic keyword arguments with `@validate_call` by [@Viicos](https://github.com/Viicos) in [#10416](https://github.com/pydantic/pydantic/pull/10416) +* Support compiled patterns in `protected_namespaces` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10522](https://github.com/pydantic/pydantic/pull/10522) +* Add support for `propertyNames` in JSON schema by [@FlorianSW](https://github.com/FlorianSW) in [#10478](https://github.com/pydantic/pydantic/pull/10478) +* Adding `__replace__` protocol for Python 3.13+ support by [@sydney-runkle](https://github.com/sydney-runkle) in [#10596](https://github.com/pydantic/pydantic/pull/10596) +* Expose public `sort` method for JSON schema generation by [@sydney-runkle](https://github.com/sydney-runkle) in [#10595](https://github.com/pydantic/pydantic/pull/10595) +* Add runtime validation of `@validate_call` callable argument by [@kc0506](https://github.com/kc0506) in [#10627](https://github.com/pydantic/pydantic/pull/10627) +* Add `experimental_allow_partial` support by [@samuelcolvin](https://github.com/samuelcolvin) in [#10748](https://github.com/pydantic/pydantic/pull/10748) +* Support default factories taking validated data as an argument by [@Viicos](https://github.com/Viicos) in [#10678](https://github.com/pydantic/pydantic/pull/10678) +* Allow subclassing `ValidationError` and `PydanticCustomError` by [@Youssefares](https://github.com/Youssefares) in [pydantic/pydantic-core#1413](https://github.com/pydantic/pydantic-core/pull/1413) +* Add `trailing-strings` support to `experimental_allow_partial` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10825](https://github.com/pydantic/pydantic/pull/10825) +* Add `rebuild()` method for `TypeAdapter` and simplify `defer_build` patterns by [@sydney-runkle](https://github.com/sydney-runkle) in [#10537](https://github.com/pydantic/pydantic/pull/10537) +* Improve `TypeAdapter` instance repr by [@sydney-runkle](https://github.com/sydney-runkle) in [#10872](https://github.com/pydantic/pydantic/pull/10872) + +#### Changes + +* Don't allow customization of `SchemaGenerator` until interface is more stable by [@sydney-runkle](https://github.com/sydney-runkle) in [#10303](https://github.com/pydantic/pydantic/pull/10303) +* Cleanly `defer_build` on `TypeAdapters`, removing experimental flag by [@sydney-runkle](https://github.com/sydney-runkle) in [#10329](https://github.com/pydantic/pydantic/pull/10329) +* Fix `mro` of generic subclass by [@kc0506](https://github.com/kc0506) in [#10100](https://github.com/pydantic/pydantic/pull/10100) +* Strip whitespaces on JSON Schema title generation by [@sydney-runkle](https://github.com/sydney-runkle) in [#10404](https://github.com/pydantic/pydantic/pull/10404) +* Use `b64decode` and `b64encode` for `Base64Bytes` type by [@sydney-runkle](https://github.com/sydney-runkle) in [#10486](https://github.com/pydantic/pydantic/pull/10486) +* Relax protected namespace config default by [@sydney-runkle](https://github.com/sydney-runkle) in [#10441](https://github.com/pydantic/pydantic/pull/10441) +* Revalidate parametrized generics if instance's origin is subclass of OG class by [@sydney-runkle](https://github.com/sydney-runkle) in [#10666](https://github.com/pydantic/pydantic/pull/10666) +* Warn if configuration is specified on the `@dataclass` decorator and with the `__pydantic_config__` attribute by [@sydney-runkle](https://github.com/sydney-runkle) in [#10406](https://github.com/pydantic/pydantic/pull/10406) +* Recommend against using `Ellipsis` (...) with `Field` by [@Viicos](https://github.com/Viicos) in [#10661](https://github.com/pydantic/pydantic/pull/10661) +* Migrate to subclassing instead of annotated approach for pydantic url types by [@sydney-runkle](https://github.com/sydney-runkle) in [#10662](https://github.com/pydantic/pydantic/pull/10662) +* Change JSON schema generation of `Literal`s and `Enums` by [@Viicos](https://github.com/Viicos) in [#10692](https://github.com/pydantic/pydantic/pull/10692) +* Simplify unions involving `Any` or `Never` when replacing type variables by [@Viicos](https://github.com/Viicos) in [#10338](https://github.com/pydantic/pydantic/pull/10338) +* Do not require padding when decoding `base64` bytes by [@bschoenmaeckers](https://github.com/bschoenmaeckers) in [pydantic/pydantic-core#1448](https://github.com/pydantic/pydantic-core/pull/1448) +* Support dates all the way to 1BC by [@changhc](https://github.com/changhc) in [pydantic/speedate#77](https://github.com/pydantic/speedate/pull/77) + +#### Performance + +* Schema cleaning: skip unnecessary copies during schema walking by [@Viicos](https://github.com/Viicos) in [#10286](https://github.com/pydantic/pydantic/pull/10286) +* Refactor namespace logic for annotations evaluation by [@Viicos](https://github.com/Viicos) in [#10530](https://github.com/pydantic/pydantic/pull/10530) +* Improve email regexp on edge cases by [@AlekseyLobanov](https://github.com/AlekseyLobanov) in [#10601](https://github.com/pydantic/pydantic/pull/10601) +* `CoreMetadata` refactor with an emphasis on documentation, schema build time performance, and reducing complexity by [@sydney-runkle](https://github.com/sydney-runkle) in [#10675](https://github.com/pydantic/pydantic/pull/10675) + +#### Fixes + +* Remove guarding check on `computed_field` with `field_serializer` by [@nix010](https://github.com/nix010) in [#10390](https://github.com/pydantic/pydantic/pull/10390) +* Fix `Predicate` issue in `v2.9.0` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10321](https://github.com/pydantic/pydantic/pull/10321) +* Fixing `annotated-types` bound by [@sydney-runkle](https://github.com/sydney-runkle) in [#10327](https://github.com/pydantic/pydantic/pull/10327) +* Turn `tzdata` install requirement into optional `timezone` dependency by [@jakob-keller](https://github.com/jakob-keller) in [#10331](https://github.com/pydantic/pydantic/pull/10331) +* Use correct types namespace when building `namedtuple` core schemas by [@Viicos](https://github.com/Viicos) in [#10337](https://github.com/pydantic/pydantic/pull/10337) +* Fix evaluation of stringified annotations during namespace inspection by [@Viicos](https://github.com/Viicos) in [#10347](https://github.com/pydantic/pydantic/pull/10347) +* Fix `IncEx` type alias definition by [@Viicos](https://github.com/Viicos) in [#10339](https://github.com/pydantic/pydantic/pull/10339) +* Do not error when trying to evaluate annotations of private attributes by [@Viicos](https://github.com/Viicos) in [#10358](https://github.com/pydantic/pydantic/pull/10358) +* Fix nested type statement by [@kc0506](https://github.com/kc0506) in [#10369](https://github.com/pydantic/pydantic/pull/10369) +* Improve typing of `ModelMetaclass.mro` by [@Viicos](https://github.com/Viicos) in [#10372](https://github.com/pydantic/pydantic/pull/10372) +* Fix class access of deprecated `computed_field`s by [@Viicos](https://github.com/Viicos) in [#10391](https://github.com/pydantic/pydantic/pull/10391) +* Make sure `inspect.iscoroutinefunction` works on coroutines decorated with `@validate_call` by [@MovisLi](https://github.com/MovisLi) in [#10374](https://github.com/pydantic/pydantic/pull/10374) +* Fix `NameError` when using `validate_call` with PEP 695 on a class by [@kc0506](https://github.com/kc0506) in [#10380](https://github.com/pydantic/pydantic/pull/10380) +* Fix `ZoneInfo` with various invalid types by [@sydney-runkle](https://github.com/sydney-runkle) in [#10408](https://github.com/pydantic/pydantic/pull/10408) +* Fix `PydanticUserError` on empty `model_config` with annotations by [@cdwilson](https://github.com/cdwilson) in [#10412](https://github.com/pydantic/pydantic/pull/10412) +* Fix variance issue in `_IncEx` type alias, only allow `True` by [@Viicos](https://github.com/Viicos) in [#10414](https://github.com/pydantic/pydantic/pull/10414) +* Fix serialization schema generation when using `PlainValidator` by [@Viicos](https://github.com/Viicos) in [#10427](https://github.com/pydantic/pydantic/pull/10427) +* Fix schema generation error when serialization schema holds references by [@Viicos](https://github.com/Viicos) in [#10444](https://github.com/pydantic/pydantic/pull/10444) +* Inline references if possible when generating schema for `json_schema_input_type` by [@Viicos](https://github.com/Viicos) in [#10439](https://github.com/pydantic/pydantic/pull/10439) +* Fix recursive arguments in `Representation` by [@Viicos](https://github.com/Viicos) in [#10480](https://github.com/pydantic/pydantic/pull/10480) +* Fix representation for builtin function types by [@kschwab](https://github.com/kschwab) in [#10479](https://github.com/pydantic/pydantic/pull/10479) +* Add python validators for decimal constraints (`max_digits` and `decimal_places`) by [@sydney-runkle](https://github.com/sydney-runkle) in [#10506](https://github.com/pydantic/pydantic/pull/10506) +* Only fetch `__pydantic_core_schema__` from the current class during schema generation by [@Viicos](https://github.com/Viicos) in [#10518](https://github.com/pydantic/pydantic/pull/10518) +* Fix `stacklevel` on deprecation warnings for `BaseModel` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10520](https://github.com/pydantic/pydantic/pull/10520) +* Fix warning `stacklevel` in `BaseModel.__init__` by [@Viicos](https://github.com/Viicos) in [#10526](https://github.com/pydantic/pydantic/pull/10526) +* Improve error handling for in-evaluable refs for discriminator application by [@sydney-runkle](https://github.com/sydney-runkle) in [#10440](https://github.com/pydantic/pydantic/pull/10440) +* Change the signature of `ConfigWrapper.core_config` to take the title directly by [@Viicos](https://github.com/Viicos) in [#10562](https://github.com/pydantic/pydantic/pull/10562) +* Do not use the previous config from the stack for dataclasses without config by [@Viicos](https://github.com/Viicos) in [#10576](https://github.com/pydantic/pydantic/pull/10576) +* Fix serialization for IP types with `mode='python'` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10594](https://github.com/pydantic/pydantic/pull/10594) +* Support constraint application for `Base64Etc` types by [@sydney-runkle](https://github.com/sydney-runkle) in [#10584](https://github.com/pydantic/pydantic/pull/10584) +* Fix `validate_call` ignoring `Field` in `Annotated` by [@kc0506](https://github.com/kc0506) in [#10610](https://github.com/pydantic/pydantic/pull/10610) +* Raise an error when `Self` is invalid by [@kc0506](https://github.com/kc0506) in [#10609](https://github.com/pydantic/pydantic/pull/10609) +* Using `core_schema.InvalidSchema` instead of metadata injection + checks by [@sydney-runkle](https://github.com/sydney-runkle) in [#10523](https://github.com/pydantic/pydantic/pull/10523) +* Tweak type alias logic by [@kc0506](https://github.com/kc0506) in [#10643](https://github.com/pydantic/pydantic/pull/10643) +* Support usage of `type` with `typing.Self` and type aliases by [@kc0506](https://github.com/kc0506) in [#10621](https://github.com/pydantic/pydantic/pull/10621) +* Use overloads for `Field` and `PrivateAttr` functions by [@Viicos](https://github.com/Viicos) in [#10651](https://github.com/pydantic/pydantic/pull/10651) +* Clean up the `mypy` plugin implementation by [@Viicos](https://github.com/Viicos) in [#10669](https://github.com/pydantic/pydantic/pull/10669) +* Properly check for `typing_extensions` variant of `TypeAliasType` by [@Daraan](https://github.com/Daraan) in [#10713](https://github.com/pydantic/pydantic/pull/10713) +* Allow any mapping in `BaseModel.model_copy()` by [@Viicos](https://github.com/Viicos) in [#10751](https://github.com/pydantic/pydantic/pull/10751) +* Fix `isinstance` behavior for urls by [@sydney-runkle](https://github.com/sydney-runkle) in [#10766](https://github.com/pydantic/pydantic/pull/10766) +* Ensure `cached_property` can be set on Pydantic models by [@Viicos](https://github.com/Viicos) in [#10774](https://github.com/pydantic/pydantic/pull/10774) +* Fix equality checks for primitives in literals by [@sydney-runkle](https://github.com/sydney-runkle) in [pydantic/pydantic-core#1459](https://github.com/pydantic/pydantic-core/pull/1459) +* Properly enforce `host_required` for URLs by [@Viicos](https://github.com/Viicos) in [pydantic/pydantic-core#1488](https://github.com/pydantic/pydantic-core/pull/1488) +* Fix when `coerce_numbers_to_str` enabled and string has invalid Unicode character by [@andrey-berenda](https://github.com/andrey-berenda) in [pydantic/pydantic-core#1515](https://github.com/pydantic/pydantic-core/pull/1515) +* Fix serializing `complex` values in `Enum`s by [@changhc](https://github.com/changhc) in [pydantic/pydantic-core#1524](https://github.com/pydantic/pydantic-core/pull/1524) +* Refactor `_typing_extra` module by [@Viicos](https://github.com/Viicos) in [#10725](https://github.com/pydantic/pydantic/pull/10725) +* Support intuitive equality for urls by [@sydney-runkle](https://github.com/sydney-runkle) in [#10798](https://github.com/pydantic/pydantic/pull/10798) +* Add `bytearray` to `TypeAdapter.validate_json` signature by [@samuelcolvin](https://github.com/samuelcolvin) in [#10802](https://github.com/pydantic/pydantic/pull/10802) +* Ensure class access of method descriptors is performed when used as a default with `Field` by [@Viicos](https://github.com/Viicos) in [#10816](https://github.com/pydantic/pydantic/pull/10816) +* Fix circular import with `validate_call` by [@sydney-runkle](https://github.com/sydney-runkle) in [#10807](https://github.com/pydantic/pydantic/pull/10807) +* Fix error when using type aliases referencing other type aliases by [@Viicos](https://github.com/Viicos) in [#10809](https://github.com/pydantic/pydantic/pull/10809) +* Fix `IncEx` type alias to be compatible with mypy by [@Viicos](https://github.com/Viicos) in [#10813](https://github.com/pydantic/pydantic/pull/10813) +* Make `__signature__` a lazy property, do not deepcopy defaults by [@Viicos](https://github.com/Viicos) in [#10818](https://github.com/pydantic/pydantic/pull/10818) +* Make `__signature__` lazy for dataclasses, too by [@sydney-runkle](https://github.com/sydney-runkle) in [#10832](https://github.com/pydantic/pydantic/pull/10832) +* Subclass all single host url classes from `AnyUrl` to preserve behavior from v2.9 by [@sydney-runkle](https://github.com/sydney-runkle) in [#10856](https://github.com/pydantic/pydantic/pull/10856) + +### New Contributors + +* [@jakob-keller](https://github.com/jakob-keller) made their first contribution in [#10331](https://github.com/pydantic/pydantic/pull/10331) +* [@MovisLi](https://github.com/MovisLi) made their first contribution in [#10374](https://github.com/pydantic/pydantic/pull/10374) +* [@joaopalmeiro](https://github.com/joaopalmeiro) made their first contribution in [#10405](https://github.com/pydantic/pydantic/pull/10405) +* [@theunkn0wn1](https://github.com/theunkn0wn1) made their first contribution in [#10378](https://github.com/pydantic/pydantic/pull/10378) +* [@cdwilson](https://github.com/cdwilson) made their first contribution in [#10412](https://github.com/pydantic/pydantic/pull/10412) +* [@dlax](https://github.com/dlax) made their first contribution in [#10421](https://github.com/pydantic/pydantic/pull/10421) +* [@kschwab](https://github.com/kschwab) made their first contribution in [#10479](https://github.com/pydantic/pydantic/pull/10479) +* [@santibreo](https://github.com/santibreo) made their first contribution in [#10453](https://github.com/pydantic/pydantic/pull/10453) +* [@FlorianSW](https://github.com/FlorianSW) made their first contribution in [#10478](https://github.com/pydantic/pydantic/pull/10478) +* [@tkasuz](https://github.com/tkasuz) made their first contribution in [#10555](https://github.com/pydantic/pydantic/pull/10555) +* [@AlekseyLobanov](https://github.com/AlekseyLobanov) made their first contribution in [#10601](https://github.com/pydantic/pydantic/pull/10601) +* [@NiclasvanEyk](https://github.com/NiclasvanEyk) made their first contribution in [#10667](https://github.com/pydantic/pydantic/pull/10667) +* [@mschoettle](https://github.com/mschoettle) made their first contribution in [#10677](https://github.com/pydantic/pydantic/pull/10677) +* [@Daraan](https://github.com/Daraan) made their first contribution in [#10713](https://github.com/pydantic/pydantic/pull/10713) +* [@k4nar](https://github.com/k4nar) made their first contribution in [#10736](https://github.com/pydantic/pydantic/pull/10736) +* [@UriyaHarpeness](https://github.com/UriyaHarpeness) made their first contribution in [#10740](https://github.com/pydantic/pydantic/pull/10740) +* [@frfahim](https://github.com/frfahim) made their first contribution in [#10727](https://github.com/pydantic/pydantic/pull/10727) + +## v2.10.0b2 (2024-11-13) + +Pre-release, see [the GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.0b2) for details. + +## v2.10.0b1 (2024-11-06) + +Pre-release, see [the GitHub release](https://github.com/pydantic/pydantic/releases/tag/v2.10.0b1) for details. + + +... see [here](https://docs.pydantic.dev/changelog/#v0322-2019-08-17) for earlier changes. diff --git a/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/RECORD b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..22201cd183128c0593766b12bcd83e79b42191c8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/RECORD @@ -0,0 +1,215 @@ +pydantic-2.11.10.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pydantic-2.11.10.dist-info/METADATA,sha256=yBe6AefUrn0Jlugz0HIuqus7JsaXVWMezLgMo74EESg,68647 +pydantic-2.11.10.dist-info/RECORD,, +pydantic-2.11.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87 +pydantic-2.11.10.dist-info/licenses/LICENSE,sha256=qeGG88oWte74QxjnpwFyE1GgDLe4rjpDlLZ7SeNSnvM,1129 +pydantic/__init__.py,sha256=D3_-0aRPoAF5EH4T4JPVOYLNEc-DeaCcDt6UzIjP_D0,15395 +pydantic/__pycache__/__init__.cpython-312.pyc,, +pydantic/__pycache__/_migration.cpython-312.pyc,, +pydantic/__pycache__/alias_generators.cpython-312.pyc,, +pydantic/__pycache__/aliases.cpython-312.pyc,, +pydantic/__pycache__/annotated_handlers.cpython-312.pyc,, +pydantic/__pycache__/class_validators.cpython-312.pyc,, +pydantic/__pycache__/color.cpython-312.pyc,, +pydantic/__pycache__/config.cpython-312.pyc,, +pydantic/__pycache__/dataclasses.cpython-312.pyc,, +pydantic/__pycache__/datetime_parse.cpython-312.pyc,, +pydantic/__pycache__/decorator.cpython-312.pyc,, +pydantic/__pycache__/env_settings.cpython-312.pyc,, +pydantic/__pycache__/error_wrappers.cpython-312.pyc,, +pydantic/__pycache__/errors.cpython-312.pyc,, +pydantic/__pycache__/fields.cpython-312.pyc,, +pydantic/__pycache__/functional_serializers.cpython-312.pyc,, +pydantic/__pycache__/functional_validators.cpython-312.pyc,, +pydantic/__pycache__/generics.cpython-312.pyc,, +pydantic/__pycache__/json.cpython-312.pyc,, +pydantic/__pycache__/json_schema.cpython-312.pyc,, +pydantic/__pycache__/main.cpython-312.pyc,, +pydantic/__pycache__/mypy.cpython-312.pyc,, +pydantic/__pycache__/networks.cpython-312.pyc,, +pydantic/__pycache__/parse.cpython-312.pyc,, +pydantic/__pycache__/root_model.cpython-312.pyc,, +pydantic/__pycache__/schema.cpython-312.pyc,, +pydantic/__pycache__/tools.cpython-312.pyc,, +pydantic/__pycache__/type_adapter.cpython-312.pyc,, +pydantic/__pycache__/types.cpython-312.pyc,, +pydantic/__pycache__/typing.cpython-312.pyc,, +pydantic/__pycache__/utils.cpython-312.pyc,, +pydantic/__pycache__/validate_call_decorator.cpython-312.pyc,, +pydantic/__pycache__/validators.cpython-312.pyc,, +pydantic/__pycache__/version.cpython-312.pyc,, +pydantic/__pycache__/warnings.cpython-312.pyc,, +pydantic/_internal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pydantic/_internal/__pycache__/__init__.cpython-312.pyc,, +pydantic/_internal/__pycache__/_config.cpython-312.pyc,, +pydantic/_internal/__pycache__/_core_metadata.cpython-312.pyc,, +pydantic/_internal/__pycache__/_core_utils.cpython-312.pyc,, +pydantic/_internal/__pycache__/_dataclasses.cpython-312.pyc,, +pydantic/_internal/__pycache__/_decorators.cpython-312.pyc,, +pydantic/_internal/__pycache__/_decorators_v1.cpython-312.pyc,, +pydantic/_internal/__pycache__/_discriminated_union.cpython-312.pyc,, +pydantic/_internal/__pycache__/_docs_extraction.cpython-312.pyc,, +pydantic/_internal/__pycache__/_fields.cpython-312.pyc,, +pydantic/_internal/__pycache__/_forward_ref.cpython-312.pyc,, +pydantic/_internal/__pycache__/_generate_schema.cpython-312.pyc,, +pydantic/_internal/__pycache__/_generics.cpython-312.pyc,, +pydantic/_internal/__pycache__/_git.cpython-312.pyc,, +pydantic/_internal/__pycache__/_import_utils.cpython-312.pyc,, +pydantic/_internal/__pycache__/_internal_dataclass.cpython-312.pyc,, +pydantic/_internal/__pycache__/_known_annotated_metadata.cpython-312.pyc,, +pydantic/_internal/__pycache__/_mock_val_ser.cpython-312.pyc,, +pydantic/_internal/__pycache__/_model_construction.cpython-312.pyc,, +pydantic/_internal/__pycache__/_namespace_utils.cpython-312.pyc,, +pydantic/_internal/__pycache__/_repr.cpython-312.pyc,, +pydantic/_internal/__pycache__/_schema_gather.cpython-312.pyc,, +pydantic/_internal/__pycache__/_schema_generation_shared.cpython-312.pyc,, +pydantic/_internal/__pycache__/_serializers.cpython-312.pyc,, +pydantic/_internal/__pycache__/_signature.cpython-312.pyc,, +pydantic/_internal/__pycache__/_typing_extra.cpython-312.pyc,, +pydantic/_internal/__pycache__/_utils.cpython-312.pyc,, +pydantic/_internal/__pycache__/_validate_call.cpython-312.pyc,, +pydantic/_internal/__pycache__/_validators.cpython-312.pyc,, +pydantic/_internal/_config.py,sha256=WV07hp8xf0Q0yP9IwMvuGLQmu34AZl5sBs2JaOgCk9I,14253 +pydantic/_internal/_core_metadata.py,sha256=Y_g2t3i7uluK-wXCZvzJfRFMPUM23aBYLfae4FzBPy0,5162 +pydantic/_internal/_core_utils.py,sha256=_-ZuXhpi_0JDpZzz8jvGr82kgS3PEritWR22fjWpw48,6746 +pydantic/_internal/_dataclasses.py,sha256=GA-NO1cgYbce0UwZP-sfPe5AujHjhvgTKbPCyg9GGP8,8990 +pydantic/_internal/_decorators.py,sha256=NS7SKQvtDgnsAd37mjqtwPh19td57FJ69LsceO5SywI,32638 +pydantic/_internal/_decorators_v1.py,sha256=tfdfdpQKY4R2XCOwqHbZeoQMur6VNigRrfhudXBHx38,6185 +pydantic/_internal/_discriminated_union.py,sha256=aMl0SRSyQyHfW4-klnMTHNvwSRoqE3H3PRV_05vRsTg,25478 +pydantic/_internal/_docs_extraction.py,sha256=p-STFvLHUzxrj6bblpaAAYWmq4INxVCAdIupDgQYSIw,3831 +pydantic/_internal/_fields.py,sha256=tFmaX47Q2z8QCCPJ4K8MrPfgKDztx9clntzPxBv0OKo,23205 +pydantic/_internal/_forward_ref.py,sha256=5n3Y7-3AKLn8_FS3Yc7KutLiPUhyXmAtkEZOaFnonwM,611 +pydantic/_internal/_generate_schema.py,sha256=LWJsmvNdWDh1QxY4WelsFSw1_nScPwEfJdpwMZH5V4k,133821 +pydantic/_internal/_generics.py,sha256=D1_0xgqnL6TJQe_fFyaSk2Ug_F-kT_jRBfLjHFLCIqQ,23849 +pydantic/_internal/_git.py,sha256=IwPh3DPfa2Xq3rBuB9Nx8luR2A1i69QdeTfWWXIuCVg,809 +pydantic/_internal/_import_utils.py,sha256=TRhxD5OuY6CUosioBdBcJUs0om7IIONiZdYAV7zQ8jM,402 +pydantic/_internal/_internal_dataclass.py,sha256=_bedc1XbuuygRGiLZqkUkwwFpQaoR1hKLlR501nyySY,144 +pydantic/_internal/_known_annotated_metadata.py,sha256=lYAPiUhfSgfpY6qH9xJPJTEMoowv27QmcyOgQzys90U,16213 +pydantic/_internal/_mock_val_ser.py,sha256=wmRRFSBvqfcLbI41PsFliB4u2AZ3mJpZeiERbD3xKTo,8885 +pydantic/_internal/_model_construction.py,sha256=2Qa5Y4EgBojkhsVHu0OjpphUIlWYuVXMg1KC2opc00s,35228 +pydantic/_internal/_namespace_utils.py,sha256=CMG7nEAXVb-Idqyd3CgdulRrM-zEXOPe3kYEDBqnSKw,12878 +pydantic/_internal/_repr.py,sha256=t7GNyaUU8xvqwlDHxVE2IyDeaNZrK7p01ojQPP0UI_o,5081 +pydantic/_internal/_schema_gather.py,sha256=VLEv51TYEeeND2czsyrmJq1MVnJqTOmnLan7VG44c8A,9114 +pydantic/_internal/_schema_generation_shared.py,sha256=F_rbQbrkoomgxsskdHpP0jUJ7TCfe0BADAEkq6CJ4nM,4842 +pydantic/_internal/_serializers.py,sha256=qQ3Rak4J6bqbnjGCRjiAY4M8poLo0s5qH46sXZSQQuA,1474 +pydantic/_internal/_signature.py,sha256=8EljPJe4pSnapuirG5DkBAgD1hggHxEAyzFPH-9H0zE,6779 +pydantic/_internal/_typing_extra.py,sha256=PO3u2JmX3JKlTFy0Ew95iyjAgYHgJsqqskev4zooB2I,28216 +pydantic/_internal/_utils.py,sha256=iRmCSO0uoFhAL_ChHaYSCKrswpSrRHYoO_YQSFfCJxU,15344 +pydantic/_internal/_validate_call.py,sha256=PfdVnSzhXOrENtaDoDw3PFWPVYD5W_gNYPe8p3Ug6Lg,5321 +pydantic/_internal/_validators.py,sha256=TJcR9bxcPXjzntN6Qgib8cyPRkFZQxHW32SoKGEcp0k,20610 +pydantic/_migration.py,sha256=_6VCCVWNYB7fDpbP2MqW4bXXqo17C5_J907u9zNJQbM,11907 +pydantic/alias_generators.py,sha256=KM1n3u4JfLSBl1UuYg3hoYHzXJD-yvgrnq8u1ccwh_A,2124 +pydantic/aliases.py,sha256=vhCHyoSWnX-EJ-wWb5qj4xyRssgGWnTQfzQp4GSZ9ug,4937 +pydantic/annotated_handlers.py,sha256=WfyFSqwoEIFXBh7T73PycKloI1DiX45GWi0-JOsCR4Y,4407 +pydantic/class_validators.py,sha256=i_V3j-PYdGLSLmj_IJZekTRjunO8SIVz8LMlquPyP7E,148 +pydantic/color.py,sha256=AzqGfVQHF92_ZctDcue0DM4yTp2P6tekkwRINTWrLIo,21481 +pydantic/config.py,sha256=roz_FbfFPoVpJVpB1G7dJ8A3swghQjdN-ozrBxbLShM,42048 +pydantic/dataclasses.py,sha256=K2e76b_Cj1yvwcwfJVR7nQnLoPdetVig5yHVMGuzkpE,16644 +pydantic/datetime_parse.py,sha256=QC-WgMxMr_wQ_mNXUS7AVf-2hLEhvvsPY1PQyhSGOdk,150 +pydantic/decorator.py,sha256=YX-jUApu5AKaVWKPoaV-n-4l7UbS69GEt9Ra3hszmKI,145 +pydantic/deprecated/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pydantic/deprecated/__pycache__/__init__.cpython-312.pyc,, +pydantic/deprecated/__pycache__/class_validators.cpython-312.pyc,, +pydantic/deprecated/__pycache__/config.cpython-312.pyc,, +pydantic/deprecated/__pycache__/copy_internals.cpython-312.pyc,, +pydantic/deprecated/__pycache__/decorator.cpython-312.pyc,, +pydantic/deprecated/__pycache__/json.cpython-312.pyc,, +pydantic/deprecated/__pycache__/parse.cpython-312.pyc,, +pydantic/deprecated/__pycache__/tools.cpython-312.pyc,, +pydantic/deprecated/class_validators.py,sha256=rwfP165xity36foy1NNCg4Jf9Sul44sJLW-A5sseahI,10245 +pydantic/deprecated/config.py,sha256=k_lsVk57paxLJOcBueH07cu1OgEgWdVBxm6lfaC3CCU,2663 +pydantic/deprecated/copy_internals.py,sha256=Ku0LHLEU0WcoIInNHls7PjuBvpLFTQ4Uus77jQ3Yi08,7616 +pydantic/deprecated/decorator.py,sha256=TBm6bJ7wJsNih_8Wq5IzDcwP32m9_vfxs96desLuk00,10845 +pydantic/deprecated/json.py,sha256=HlWCG35RRrxyzuTS6LTQiZBwRhmDZWmeqQH8rLW6wA8,4657 +pydantic/deprecated/parse.py,sha256=Gzd6b_g8zJXcuE7QRq5adhx_EMJahXfcpXCF0RgrqqI,2511 +pydantic/deprecated/tools.py,sha256=Nrm9oFRZWp8-jlfvPgJILEsywp4YzZD52XIGPDLxHcI,3330 +pydantic/env_settings.py,sha256=6IHeeWEqlUPRUv3V-AXiF_W91fg2Jw_M3O0l34J_eyA,148 +pydantic/error_wrappers.py,sha256=RK6mqATc9yMD-KBD9IJS9HpKCprWHd8wo84Bnm-3fR8,150 +pydantic/errors.py,sha256=7ctBNCtt57kZFx71Ls2H86IufQARv4wPKf8DhdsVn5w,6002 +pydantic/experimental/__init__.py,sha256=j08eROfz-xW4k_X9W4m2AW26IVdyF3Eg1OzlIGA11vk,328 +pydantic/experimental/__pycache__/__init__.cpython-312.pyc,, +pydantic/experimental/__pycache__/arguments_schema.cpython-312.pyc,, +pydantic/experimental/__pycache__/pipeline.cpython-312.pyc,, +pydantic/experimental/arguments_schema.py,sha256=EFnjX_ulp-tPyUjQX5pmQtug1OFL_Acc8bcMbLd-fVY,1866 +pydantic/experimental/pipeline.py,sha256=znbMBvir3xvPA20Xj8Moco1oJMPf1VYVrIQ8KQNtDlM,23910 +pydantic/fields.py,sha256=9Ky1nTKaMhThaNkVEkJOFHQHGq2FCKSwA6-zwUB-KWo,64416 +pydantic/functional_serializers.py,sha256=3m81unH3lYovdMi00oZywlHhn1KDz9X2CO3iTtBya6A,17102 +pydantic/functional_validators.py,sha256=-yY6uj_9_GAI4aqqfZlzyGdzs06huzy6zNWD7TJp3_0,29560 +pydantic/generics.py,sha256=0ZqZ9O9annIj_3mGBRqps4htey3b5lV1-d2tUxPMMnA,144 +pydantic/json.py,sha256=ZH8RkI7h4Bz-zp8OdTAxbJUoVvcoU-jhMdRZ0B-k0xc,140 +pydantic/json_schema.py,sha256=KhsS_MWPox0PYqklnhJcb_3uiCVrEOgyhG53cUZv6QA,115430 +pydantic/main.py,sha256=v67a4-nFooC-GJ1oHgS__Vm399Ygp_NH-1WzHXwjFM0,81012 +pydantic/mypy.py,sha256=OG7AqM_6vuTxRnBPU27eUkfX5wShU6aD0dJGmMhLaN8,59265 +pydantic/networks.py,sha256=_YpSnBR2kMfoWX76sdq34cfCH-MWr5or0ve0tow7OWo,41446 +pydantic/parse.py,sha256=wkd82dgtvWtD895U_I6E1htqMlGhBSYEV39cuBSeo3A,141 +pydantic/plugin/__init__.py,sha256=5cXMmu5xL4LVZhWPE1XD8ozHZ-qEC2-s4seLe8tbN_Y,6965 +pydantic/plugin/__pycache__/__init__.cpython-312.pyc,, +pydantic/plugin/__pycache__/_loader.cpython-312.pyc,, +pydantic/plugin/__pycache__/_schema_validator.cpython-312.pyc,, +pydantic/plugin/_loader.py,sha256=nI3SEKr0mlCB556kvbyBXjYQw9b_s8UTKE9Q6iESX6s,2167 +pydantic/plugin/_schema_validator.py,sha256=QbmqsG33MBmftNQ2nNiuN22LhbrexUA7ipDVv3J02BU,5267 +pydantic/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pydantic/root_model.py,sha256=SCXhpRCgZgfqE9AGVJTC7kMAojKffL7PV4i0qcwOMm0,6279 +pydantic/schema.py,sha256=Vqqjvq_LnapVknebUd3Bp_J1p2gXZZnZRgL48bVEG7o,142 +pydantic/tools.py,sha256=iHQpd8SJ5DCTtPV5atAV06T89bjSaMFeZZ2LX9lasZY,141 +pydantic/type_adapter.py,sha256=Y3NE0YhFwxwoqrYU9caWymLWp1Avq4sRUdb5s01RoJk,31171 +pydantic/types.py,sha256=mWTvQH_Wt_CccQcEHYjcUWpyoj1U04WOnrMsMYod_64,104781 +pydantic/typing.py,sha256=P7feA35MwTcLsR1uL7db0S-oydBxobmXa55YDoBgajQ,138 +pydantic/utils.py,sha256=15nR2QpqTBFlQV4TNtTItMyTJx_fbyV-gPmIEY1Gooc,141 +pydantic/v1/__init__.py,sha256=SxQPklgBs4XHJwE6BZ9qoewYoGiNyYUnmHzEFCZbfnI,2946 +pydantic/v1/__pycache__/__init__.cpython-312.pyc,, +pydantic/v1/__pycache__/_hypothesis_plugin.cpython-312.pyc,, +pydantic/v1/__pycache__/annotated_types.cpython-312.pyc,, +pydantic/v1/__pycache__/class_validators.cpython-312.pyc,, +pydantic/v1/__pycache__/color.cpython-312.pyc,, +pydantic/v1/__pycache__/config.cpython-312.pyc,, +pydantic/v1/__pycache__/dataclasses.cpython-312.pyc,, +pydantic/v1/__pycache__/datetime_parse.cpython-312.pyc,, +pydantic/v1/__pycache__/decorator.cpython-312.pyc,, +pydantic/v1/__pycache__/env_settings.cpython-312.pyc,, +pydantic/v1/__pycache__/error_wrappers.cpython-312.pyc,, +pydantic/v1/__pycache__/errors.cpython-312.pyc,, +pydantic/v1/__pycache__/fields.cpython-312.pyc,, +pydantic/v1/__pycache__/generics.cpython-312.pyc,, +pydantic/v1/__pycache__/json.cpython-312.pyc,, +pydantic/v1/__pycache__/main.cpython-312.pyc,, +pydantic/v1/__pycache__/mypy.cpython-312.pyc,, +pydantic/v1/__pycache__/networks.cpython-312.pyc,, +pydantic/v1/__pycache__/parse.cpython-312.pyc,, +pydantic/v1/__pycache__/schema.cpython-312.pyc,, +pydantic/v1/__pycache__/tools.cpython-312.pyc,, +pydantic/v1/__pycache__/types.cpython-312.pyc,, +pydantic/v1/__pycache__/typing.cpython-312.pyc,, +pydantic/v1/__pycache__/utils.cpython-312.pyc,, +pydantic/v1/__pycache__/validators.cpython-312.pyc,, +pydantic/v1/__pycache__/version.cpython-312.pyc,, +pydantic/v1/_hypothesis_plugin.py,sha256=5ES5xWuw1FQAsymLezy8QgnVz0ZpVfU3jkmT74H27VQ,14847 +pydantic/v1/annotated_types.py,sha256=uk2NAAxqiNELKjiHhyhxKaIOh8F1lYW_LzrW3X7oZBc,3157 +pydantic/v1/class_validators.py,sha256=ULOaIUgYUDBsHL7EEVEarcM-UubKUggoN8hSbDonsFE,14672 +pydantic/v1/color.py,sha256=iZABLYp6OVoo2AFkP9Ipri_wSc6-Kklu8YuhSartd5g,16844 +pydantic/v1/config.py,sha256=a6P0Wer9x4cbwKW7Xv8poSUqM4WP-RLWwX6YMpYq9AA,6532 +pydantic/v1/dataclasses.py,sha256=784cqvInbwIPWr9usfpX3ch7z4t3J2tTK6N067_wk1o,18172 +pydantic/v1/datetime_parse.py,sha256=4Qy1kQpq3rNVZJeIHeSPDpuS2Bvhp1KPtzJG1xu-H00,7724 +pydantic/v1/decorator.py,sha256=zaaxxxoWPCm818D1bs0yhapRjXm32V8G0ZHWCdM1uXA,10339 +pydantic/v1/env_settings.py,sha256=A9VXwtRl02AY-jH0C0ouy5VNw3fi6F_pkzuHDjgAAOM,14105 +pydantic/v1/error_wrappers.py,sha256=6625Mfw9qkC2NwitB_JFAWe8B-Xv6zBU7rL9k28tfyo,5196 +pydantic/v1/errors.py,sha256=mIwPED5vGM5Q5v4C4Z1JPldTRH-omvEylH6ksMhOmPw,17726 +pydantic/v1/fields.py,sha256=VqWJCriUNiEyptXroDVJ501JpVA0en2VANcksqXL2b8,50649 +pydantic/v1/generics.py,sha256=VzC9YUV-EbPpQ3aAfk1cNFej79_IzznkQ7WrmTTZS9E,17871 +pydantic/v1/json.py,sha256=WQ5Hy_hIpfdR3YS8k6N2E6KMJzsdbBi_ldWOPJaV81M,3390 +pydantic/v1/main.py,sha256=zuNpdN5Q0V0wG2UUTKt0HUy3XJ4OAvPSZDdiXY-FIzs,44824 +pydantic/v1/mypy.py,sha256=Cl8XRfCmIcVE3j5AEU52C8iDh8lcX__D3hz2jIWxMAs,38860 +pydantic/v1/networks.py,sha256=HYNtKAfOmOnKJpsDg1g6SIkj9WPhU_-i8l5e2JKBpG4,22124 +pydantic/v1/parse.py,sha256=BJtdqiZRtav9VRFCmOxoY-KImQmjPy-A_NoojiFUZxY,1821 +pydantic/v1/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +pydantic/v1/schema.py,sha256=aqBuA--cq8gAVkim5BJPFASHzOZ8dFtmFX_fNGr6ip4,47801 +pydantic/v1/tools.py,sha256=1lDdXHk0jL5uP3u5RCYAvUAlGClgAO-45lkq9j7fyBA,2881 +pydantic/v1/types.py,sha256=Fltx5GoP_qaUmAktlGz7nFeJa13yNy3FY1-RcMzEVt8,35455 +pydantic/v1/typing.py,sha256=HNtuKvgH4EHIeb2ytkd7VSyG6mxP9RKqEqEql-1ab14,19720 +pydantic/v1/utils.py,sha256=M5FRyfNUb1A2mk9laGgCVdfHHb3AtQgrjO5qfyBf4xA,25989 +pydantic/v1/validators.py,sha256=lyUkn1MWhHxlCX5ZfEgFj_CAHojoiPcaQeMdEM9XviU,22187 +pydantic/v1/version.py,sha256=HXnXW-1bMW5qKhlr5RgOEPohrZDCDSuyy8-gi8GCgZo,1039 +pydantic/validate_call_decorator.py,sha256=8jqLlgXTjWEj4dXDg0wI3EGQKkb0JnCsL_JSUjbU5Sg,4389 +pydantic/validators.py,sha256=pwbIJXVb1CV2mAE4w_EGfNj7DwzsKaWw_tTL6cviTus,146 +pydantic/version.py,sha256=VepAy6cBRotVjFFqWTpYJDwOkeCsiywS-1AC4IOBIbk,2711 +pydantic/warnings.py,sha256=gqDTQ2FX7wGLZJV3XboQSiRXKHknss3pfIOXL0BDXTk,3772 diff --git a/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/WHEEL b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..12228d414b6cfed7c39d3781c85c63256a1d7fb5 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: hatchling 1.27.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..488c6260c10f2e88fa1fae58a63fccec8d600cd1 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic-2.11.10.dist-info/licenses/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2017 to present Pydantic Services Inc. and individual contributors. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..1e65815cf0b3132689485874a93034ede7206bf4 --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/LICENSE @@ -0,0 +1,54 @@ +Copyright 2017- Paul Ganssle +Copyright 2017- dateutil contributors (see AUTHORS file) + + 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. + +The above license applies to all contributions after 2017-12-01, as well as +all contributions that have been re-licensed (see AUTHORS file for the list of +contributors who have re-licensed their code). +-------------------------------------------------------------------------------- +dateutil - Extensions to the standard Python datetime module. + +Copyright (c) 2003-2011 - Gustavo Niemeyer +Copyright (c) 2012-2014 - Tomi Pieviläinen +Copyright (c) 2014-2016 - Yaron de Leeuw +Copyright (c) 2015- - Paul Ganssle +Copyright (c) 2015- - dateutil contributors (see AUTHORS file) + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + * Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +The above BSD License Applies to all code, even that also covered by Apache 2.0. \ No newline at end of file diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..577f2bf2b7749e1b123b8225d610b1b257e430cc --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/METADATA @@ -0,0 +1,204 @@ +Metadata-Version: 2.1 +Name: python-dateutil +Version: 2.9.0.post0 +Summary: Extensions to the standard Python datetime module +Home-page: https://github.com/dateutil/dateutil +Author: Gustavo Niemeyer +Author-email: gustavo@niemeyer.net +Maintainer: Paul Ganssle +Maintainer-email: dateutil@python.org +License: Dual License +Project-URL: Documentation, https://dateutil.readthedocs.io/en/stable/ +Project-URL: Source, https://github.com/dateutil/dateutil +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: BSD License +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 2 +Classifier: Programming Language :: Python :: 2.7 +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.3 +Classifier: Programming Language :: Python :: 3.4 +Classifier: Programming Language :: Python :: 3.5 +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Topic :: Software Development :: Libraries +Requires-Python: !=3.0.*,!=3.1.*,!=3.2.*,>=2.7 +Description-Content-Type: text/x-rst +License-File: LICENSE +Requires-Dist: six >=1.5 + +dateutil - powerful extensions to datetime +========================================== + +|pypi| |support| |licence| + +|gitter| |readthedocs| + +|travis| |appveyor| |pipelines| |coverage| + +.. |pypi| image:: https://img.shields.io/pypi/v/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: pypi version + +.. |support| image:: https://img.shields.io/pypi/pyversions/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: supported Python version + +.. |travis| image:: https://img.shields.io/travis/dateutil/dateutil/master.svg?style=flat-square&label=Travis%20Build + :target: https://travis-ci.org/dateutil/dateutil + :alt: travis build status + +.. |appveyor| image:: https://img.shields.io/appveyor/ci/dateutil/dateutil/master.svg?style=flat-square&logo=appveyor + :target: https://ci.appveyor.com/project/dateutil/dateutil + :alt: appveyor build status + +.. |pipelines| image:: https://dev.azure.com/pythondateutilazure/dateutil/_apis/build/status/dateutil.dateutil?branchName=master + :target: https://dev.azure.com/pythondateutilazure/dateutil/_build/latest?definitionId=1&branchName=master + :alt: azure pipelines build status + +.. |coverage| image:: https://codecov.io/gh/dateutil/dateutil/branch/master/graphs/badge.svg?branch=master + :target: https://codecov.io/gh/dateutil/dateutil?branch=master + :alt: Code coverage + +.. |gitter| image:: https://badges.gitter.im/dateutil/dateutil.svg + :alt: Join the chat at https://gitter.im/dateutil/dateutil + :target: https://gitter.im/dateutil/dateutil + +.. |licence| image:: https://img.shields.io/pypi/l/python-dateutil.svg?style=flat-square + :target: https://pypi.org/project/python-dateutil/ + :alt: licence + +.. |readthedocs| image:: https://img.shields.io/readthedocs/dateutil/latest.svg?style=flat-square&label=Read%20the%20Docs + :alt: Read the documentation at https://dateutil.readthedocs.io/en/latest/ + :target: https://dateutil.readthedocs.io/en/latest/ + +The `dateutil` module provides powerful extensions to +the standard `datetime` module, available in Python. + +Installation +============ +`dateutil` can be installed from PyPI using `pip` (note that the package name is +different from the importable name):: + + pip install python-dateutil + +Download +======== +dateutil is available on PyPI +https://pypi.org/project/python-dateutil/ + +The documentation is hosted at: +https://dateutil.readthedocs.io/en/stable/ + +Code +==== +The code and issue tracker are hosted on GitHub: +https://github.com/dateutil/dateutil/ + +Features +======== + +* Computing of relative deltas (next month, next year, + next Monday, last week of month, etc); +* Computing of relative deltas between two given + date and/or datetime objects; +* Computing of dates based on very flexible recurrence rules, + using a superset of the `iCalendar `_ + specification. Parsing of RFC strings is supported as well. +* Generic parsing of dates in almost any string format; +* Timezone (tzinfo) implementations for tzfile(5) format + files (/etc/localtime, /usr/share/zoneinfo, etc), TZ + environment string (in all known formats), iCalendar + format files, given ranges (with help from relative deltas), + local machine timezone, fixed offset timezone, UTC timezone, + and Windows registry-based time zones. +* Internal up-to-date world timezone information based on + Olson's database. +* Computing of Easter Sunday dates for any given year, + using Western, Orthodox or Julian algorithms; +* A comprehensive test suite. + +Quick example +============= +Here's a snapshot, just to give an idea about the power of the +package. For more examples, look at the documentation. + +Suppose you want to know how much time is left, in +years/months/days/etc, before the next easter happening on a +year with a Friday 13th in August, and you want to get today's +date out of the "date" unix system command. Here is the code: + +.. code-block:: python3 + + >>> from dateutil.relativedelta import * + >>> from dateutil.easter import * + >>> from dateutil.rrule import * + >>> from dateutil.parser import * + >>> from datetime import * + >>> now = parse("Sat Oct 11 17:13:46 UTC 2003") + >>> today = now.date() + >>> year = rrule(YEARLY,dtstart=now,bymonth=8,bymonthday=13,byweekday=FR)[0].year + >>> rdelta = relativedelta(easter(year), today) + >>> print("Today is: %s" % today) + Today is: 2003-10-11 + >>> print("Year with next Aug 13th on a Friday is: %s" % year) + Year with next Aug 13th on a Friday is: 2004 + >>> print("How far is the Easter of that year: %s" % rdelta) + How far is the Easter of that year: relativedelta(months=+6) + >>> print("And the Easter of that year is: %s" % (today+rdelta)) + And the Easter of that year is: 2004-04-11 + +Being exactly 6 months ahead was **really** a coincidence :) + +Contributing +============ + +We welcome many types of contributions - bug reports, pull requests (code, infrastructure or documentation fixes). For more information about how to contribute to the project, see the ``CONTRIBUTING.md`` file in the repository. + + +Author +====== +The dateutil module was written by Gustavo Niemeyer +in 2003. + +It is maintained by: + +* Gustavo Niemeyer 2003-2011 +* Tomi Pieviläinen 2012-2014 +* Yaron de Leeuw 2014-2016 +* Paul Ganssle 2015- + +Starting with version 2.4.1 and running until 2.8.2, all source and binary +distributions will be signed by a PGP key that has, at the very least, been +signed by the key which made the previous release. A table of release signing +keys can be found below: + +=========== ============================ +Releases Signing key fingerprint +=========== ============================ +2.4.1-2.8.2 `6B49 ACBA DCF6 BD1C A206 67AB CD54 FCE3 D964 BEFB`_ +=========== ============================ + +New releases *may* have signed tags, but binary and source distributions +uploaded to PyPI will no longer have GPG signatures attached. + +Contact +======= +Our mailing list is available at `dateutil@python.org `_. As it is hosted by the PSF, it is subject to the `PSF code of +conduct `_. + +License +======= + +All contributions after December 1, 2017 released under dual license - either `Apache 2.0 License `_ or the `BSD 3-Clause License `_. Contributions before December 1, 2017 - except those those explicitly relicensed - are released only under the BSD 3-Clause License. + + +.. _6B49 ACBA DCF6 BD1C A206 67AB CD54 FCE3 D964 BEFB: + https://pgp.mit.edu/pks/lookup?op=vindex&search=0xCD54FCE3D964BEFB diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..8abf769beb56ed307111268106e19c87137b23df --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/RECORD @@ -0,0 +1,44 @@ +dateutil/__init__.py,sha256=Mqam67WO9IkTmUFyI66vS6IoSXTp9G388DadH2LCMLY,620 +dateutil/__pycache__/__init__.cpython-312.pyc,, +dateutil/__pycache__/_common.cpython-312.pyc,, +dateutil/__pycache__/_version.cpython-312.pyc,, +dateutil/__pycache__/easter.cpython-312.pyc,, +dateutil/__pycache__/relativedelta.cpython-312.pyc,, +dateutil/__pycache__/rrule.cpython-312.pyc,, +dateutil/__pycache__/tzwin.cpython-312.pyc,, +dateutil/__pycache__/utils.cpython-312.pyc,, +dateutil/_common.py,sha256=77w0yytkrxlYbSn--lDVPUMabUXRR9I3lBv_vQRUqUY,932 +dateutil/_version.py,sha256=BV031OxDDAmy58neUg5yyqLkLaqIw7ibK9As3jiMib0,166 +dateutil/easter.py,sha256=dyBi-lKvimH1u_k6p7Z0JJK72QhqVtVBsqByvpEPKvc,2678 +dateutil/parser/__init__.py,sha256=wWk6GFuxTpjoggCGtgkceJoti4pVjl4_fHQXpNOaSYg,1766 +dateutil/parser/__pycache__/__init__.cpython-312.pyc,, +dateutil/parser/__pycache__/_parser.cpython-312.pyc,, +dateutil/parser/__pycache__/isoparser.cpython-312.pyc,, +dateutil/parser/_parser.py,sha256=7klDdyicksQB_Xgl-3UAmBwzCYor1AIZqklIcT6dH_8,58796 +dateutil/parser/isoparser.py,sha256=8Fy999bnCd1frSdOYuOraWfJTtd5W7qQ51NwNuH_hXM,13233 +dateutil/relativedelta.py,sha256=IY_mglMjoZbYfrvloTY2ce02aiVjPIkiZfqgNTZRfuA,24903 +dateutil/rrule.py,sha256=KJzKlaCd1jEbu4A38ZltslaoAUh9nSbdbOFdjp70Kew,66557 +dateutil/tz/__init__.py,sha256=F-Mz13v6jYseklQf9Te9J6nzcLDmq47gORa61K35_FA,444 +dateutil/tz/__pycache__/__init__.cpython-312.pyc,, +dateutil/tz/__pycache__/_common.cpython-312.pyc,, +dateutil/tz/__pycache__/_factories.cpython-312.pyc,, +dateutil/tz/__pycache__/tz.cpython-312.pyc,, +dateutil/tz/__pycache__/win.cpython-312.pyc,, +dateutil/tz/_common.py,sha256=cgzDTANsOXvEc86cYF77EsliuSab8Puwpsl5-bX3_S4,12977 +dateutil/tz/_factories.py,sha256=unb6XQNXrPMveksTCU-Ag8jmVZs4SojoPUcAHpWnrvU,2569 +dateutil/tz/tz.py,sha256=EUnEdMfeThXiY6l4sh9yBabZ63_POzy01zSsh9thn1o,62855 +dateutil/tz/win.py,sha256=xJszWgSwE1xPx_HJj4ZkepyukC_hNy016WMcXhbRaB8,12935 +dateutil/tzwin.py,sha256=7Ar4vdQCnnM0mKR3MUjbIKsZrBVfHgdwsJZc_mGYRew,59 +dateutil/utils.py,sha256=dKCchEw8eObi0loGTx91unBxm_7UGlU3v_FjFMdqwYM,1965 +dateutil/zoneinfo/__init__.py,sha256=KYg0pthCMjcp5MXSEiBJn3nMjZeNZav7rlJw5-tz1S4,5889 +dateutil/zoneinfo/__pycache__/__init__.cpython-312.pyc,, +dateutil/zoneinfo/__pycache__/rebuild.cpython-312.pyc,, +dateutil/zoneinfo/dateutil-zoneinfo.tar.gz,sha256=0-pS57bpaN4NiE3xKIGTWW-pW4A9tPkqGCeac5gARHU,156400 +dateutil/zoneinfo/rebuild.py,sha256=MiqYzCIHvNbMH-LdRYLv-4T0EIA7hDKt5GLR0IRTLdI,2392 +python_dateutil-2.9.0.post0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +python_dateutil-2.9.0.post0.dist-info/LICENSE,sha256=ugD1Gg2SgjtaHN4n2LW50jIeZ-2NqbwWPv-W1eF-V34,2889 +python_dateutil-2.9.0.post0.dist-info/METADATA,sha256=qdQ22jIr6AgzL5jYgyWZjofLaTpniplp_rTPrXKabpM,8354 +python_dateutil-2.9.0.post0.dist-info/RECORD,, +python_dateutil-2.9.0.post0.dist-info/WHEEL,sha256=-G_t0oGuE7UD0DrSpVZnq1hHMBV9DD2XkS5v7XpmTnk,110 +python_dateutil-2.9.0.post0.dist-info/top_level.txt,sha256=4tjdWkhRZvF7LA_BYe_L9gB2w_p2a-z5y6ArjaRkot8,9 +python_dateutil-2.9.0.post0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4724c45738f6ac125bb3a21787855562e6870440 --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.42.0) +Root-Is-Purelib: true +Tag: py2-none-any +Tag: py3-none-any + diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..66501480ba5b63f98ee9a59c1f99e5e6917da6d9 --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/top_level.txt @@ -0,0 +1 @@ +dateutil diff --git a/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/lib/python3.12/site-packages/python_dateutil-2.9.0.post0.dist-info/zip-safe @@ -0,0 +1 @@ + diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/INSTALLER b/lib/python3.12/site-packages/pytz-2025.2.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/LICENSE.txt b/lib/python3.12/site-packages/pytz-2025.2.dist-info/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..5f1c11289f6a54cb07ebdbf31d02e8e81b18b07f --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/LICENSE.txt @@ -0,0 +1,19 @@ +Copyright (c) 2003-2019 Stuart Bishop + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the "Software"), +to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, +and/or sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/METADATA b/lib/python3.12/site-packages/pytz-2025.2.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..6ecc5a3e468af016f598e4402f6a4c6efcbe2271 --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/METADATA @@ -0,0 +1,648 @@ +Metadata-Version: 2.1 +Name: pytz +Version: 2025.2 +Summary: World timezone definitions, modern and historical +Home-page: http://pythonhosted.org/pytz +Download-URL: https://pypi.org/project/pytz/ +Author: Stuart Bishop +Author-email: stuart@stuartbishop.net +Maintainer: Stuart Bishop +Maintainer-email: stuart@stuartbishop.net +License: MIT +Keywords: timezone,tzinfo,datetime,olson,time +Platform: Independent +Classifier: Development Status :: 6 - Mature +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: MIT License +Classifier: Natural Language :: English +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 2 +Classifier: Programming Language :: Python :: 2.4 +Classifier: Programming Language :: Python :: 2.5 +Classifier: Programming Language :: Python :: 2.6 +Classifier: Programming Language :: Python :: 2.7 +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.1 +Classifier: Programming Language :: Python :: 3.2 +Classifier: Programming Language :: Python :: 3.3 +Classifier: Programming Language :: Python :: 3.4 +Classifier: Programming Language :: Python :: 3.5 +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Topic :: Software Development :: Libraries :: Python Modules +License-File: LICENSE.txt + +pytz - World Timezone Definitions for Python +============================================ + +:Author: Stuart Bishop + +Introduction +~~~~~~~~~~~~ + +pytz brings the Olson tz database into Python. This library allows +accurate and cross platform timezone calculations using Python 2.4 +or higher. It also solves the issue of ambiguous times at the end +of daylight saving time, which you can read more about in the Python +Library Reference (``datetime.tzinfo``). + +Almost all of the Olson timezones are supported. + +.. note:: + + Projects using Python 3.9 or later should be using the support + now included as part of the standard library, and third party + packages work with it such as `tzdata `_. + pytz offers no advantages beyond backwards compatibility with + code written for earlier versions of Python. + +.. note:: + + This library differs from the documented Python API for + tzinfo implementations; if you want to create local wallclock + times you need to use the ``localize()`` method documented in this + document. In addition, if you perform date arithmetic on local + times that cross DST boundaries, the result may be in an incorrect + timezone (ie. subtract 1 minute from 2002-10-27 1:00 EST and you get + 2002-10-27 0:59 EST instead of the correct 2002-10-27 1:59 EDT). A + ``normalize()`` method is provided to correct this. Unfortunately these + issues cannot be resolved without modifying the Python datetime + implementation (see PEP-431). + + +Installation +~~~~~~~~~~~~ + +This package can either be installed using ``pip`` or from a tarball using the +standard Python distutils. + +If you are installing using ``pip``, you don't need to download anything as the +latest version will be downloaded for you from PyPI:: + + pip install pytz + +If you are installing from a tarball, run the following command as an +administrative user:: + + python setup.py install + + +pytz for Enterprise +~~~~~~~~~~~~~~~~~~~ + +Available as part of the Tidelift Subscription. + +The maintainers of pytz and thousands of other packages are working with Tidelift to deliver commercial support and maintenance for the open source dependencies you use to build your applications. Save time, reduce risk, and improve code health, while paying the maintainers of the exact dependencies you use. `Learn more. `_. + + +Example & Usage +~~~~~~~~~~~~~~~ + +Localized times and date arithmetic +----------------------------------- + +>>> from datetime import datetime, timedelta +>>> from pytz import timezone +>>> import pytz +>>> utc = pytz.utc +>>> utc.zone +'UTC' +>>> eastern = timezone('US/Eastern') +>>> eastern.zone +'US/Eastern' +>>> amsterdam = timezone('Europe/Amsterdam') +>>> fmt = '%Y-%m-%d %H:%M:%S %Z%z' + +This library only supports two ways of building a localized time. The +first is to use the ``localize()`` method provided by the pytz library. +This is used to localize a naive datetime (datetime with no timezone +information): + +>>> loc_dt = eastern.localize(datetime(2002, 10, 27, 6, 0, 0)) +>>> print(loc_dt.strftime(fmt)) +2002-10-27 06:00:00 EST-0500 + +The second way of building a localized time is by converting an existing +localized time using the standard ``astimezone()`` method: + +>>> ams_dt = loc_dt.astimezone(amsterdam) +>>> ams_dt.strftime(fmt) +'2002-10-27 12:00:00 CET+0100' + +Unfortunately using the tzinfo argument of the standard datetime +constructors ''does not work'' with pytz for many timezones. + +>>> datetime(2002, 10, 27, 12, 0, 0, tzinfo=amsterdam).strftime(fmt) # /!\ Does not work this way! +'2002-10-27 12:00:00 LMT+0018' + +It is safe for timezones without daylight saving transitions though, such +as UTC: + +>>> datetime(2002, 10, 27, 12, 0, 0, tzinfo=pytz.utc).strftime(fmt) # /!\ Not recommended except for UTC +'2002-10-27 12:00:00 UTC+0000' + +The preferred way of dealing with times is to always work in UTC, +converting to localtime only when generating output to be read +by humans. + +>>> utc_dt = datetime(2002, 10, 27, 6, 0, 0, tzinfo=utc) +>>> loc_dt = utc_dt.astimezone(eastern) +>>> loc_dt.strftime(fmt) +'2002-10-27 01:00:00 EST-0500' + +This library also allows you to do date arithmetic using local +times, although it is more complicated than working in UTC as you +need to use the ``normalize()`` method to handle daylight saving time +and other timezone transitions. In this example, ``loc_dt`` is set +to the instant when daylight saving time ends in the US/Eastern +timezone. + +>>> before = loc_dt - timedelta(minutes=10) +>>> before.strftime(fmt) +'2002-10-27 00:50:00 EST-0500' +>>> eastern.normalize(before).strftime(fmt) +'2002-10-27 01:50:00 EDT-0400' +>>> after = eastern.normalize(before + timedelta(minutes=20)) +>>> after.strftime(fmt) +'2002-10-27 01:10:00 EST-0500' + +Creating local times is also tricky, and the reason why working with +local times is not recommended. Unfortunately, you cannot just pass +a ``tzinfo`` argument when constructing a datetime (see the next +section for more details) + +>>> dt = datetime(2002, 10, 27, 1, 30, 0) +>>> dt1 = eastern.localize(dt, is_dst=True) +>>> dt1.strftime(fmt) +'2002-10-27 01:30:00 EDT-0400' +>>> dt2 = eastern.localize(dt, is_dst=False) +>>> dt2.strftime(fmt) +'2002-10-27 01:30:00 EST-0500' + +Converting between timezones is more easily done, using the +standard astimezone method. + +>>> utc_dt = datetime.fromtimestamp(1143408899, tz=utc) +>>> utc_dt.strftime(fmt) +'2006-03-26 21:34:59 UTC+0000' +>>> au_tz = timezone('Australia/Sydney') +>>> au_dt = utc_dt.astimezone(au_tz) +>>> au_dt.strftime(fmt) +'2006-03-27 08:34:59 AEDT+1100' +>>> utc_dt2 = au_dt.astimezone(utc) +>>> utc_dt2.strftime(fmt) +'2006-03-26 21:34:59 UTC+0000' +>>> utc_dt == utc_dt2 +True + +You can take shortcuts when dealing with the UTC side of timezone +conversions. ``normalize()`` and ``localize()`` are not really +necessary when there are no daylight saving time transitions to +deal with. + +>>> utc_dt = datetime.fromtimestamp(1143408899, tz=utc) +>>> utc_dt.strftime(fmt) +'2006-03-26 21:34:59 UTC+0000' +>>> au_tz = timezone('Australia/Sydney') +>>> au_dt = au_tz.normalize(utc_dt.astimezone(au_tz)) +>>> au_dt.strftime(fmt) +'2006-03-27 08:34:59 AEDT+1100' +>>> utc_dt2 = au_dt.astimezone(utc) +>>> utc_dt2.strftime(fmt) +'2006-03-26 21:34:59 UTC+0000' + + +``tzinfo`` API +-------------- + +The ``tzinfo`` instances returned by the ``timezone()`` function have +been extended to cope with ambiguous times by adding an ``is_dst`` +parameter to the ``utcoffset()``, ``dst()`` && ``tzname()`` methods. + +>>> tz = timezone('America/St_Johns') + +>>> normal = datetime(2009, 9, 1) +>>> ambiguous = datetime(2009, 10, 31, 23, 30) + +The ``is_dst`` parameter is ignored for most timestamps. It is only used +during DST transition ambiguous periods to resolve that ambiguity. + +>>> print(tz.utcoffset(normal, is_dst=True)) +-1 day, 21:30:00 +>>> print(tz.dst(normal, is_dst=True)) +1:00:00 +>>> tz.tzname(normal, is_dst=True) +'NDT' + +>>> print(tz.utcoffset(ambiguous, is_dst=True)) +-1 day, 21:30:00 +>>> print(tz.dst(ambiguous, is_dst=True)) +1:00:00 +>>> tz.tzname(ambiguous, is_dst=True) +'NDT' + +>>> print(tz.utcoffset(normal, is_dst=False)) +-1 day, 21:30:00 +>>> tz.dst(normal, is_dst=False).seconds +3600 +>>> tz.tzname(normal, is_dst=False) +'NDT' + +>>> print(tz.utcoffset(ambiguous, is_dst=False)) +-1 day, 20:30:00 +>>> tz.dst(ambiguous, is_dst=False) +datetime.timedelta(0) +>>> tz.tzname(ambiguous, is_dst=False) +'NST' + +If ``is_dst`` is not specified, ambiguous timestamps will raise +an ``pytz.exceptions.AmbiguousTimeError`` exception. + +>>> print(tz.utcoffset(normal)) +-1 day, 21:30:00 +>>> print(tz.dst(normal)) +1:00:00 +>>> tz.tzname(normal) +'NDT' + +>>> import pytz.exceptions +>>> try: +... tz.utcoffset(ambiguous) +... except pytz.exceptions.AmbiguousTimeError: +... print('pytz.exceptions.AmbiguousTimeError: %s' % ambiguous) +pytz.exceptions.AmbiguousTimeError: 2009-10-31 23:30:00 +>>> try: +... tz.dst(ambiguous) +... except pytz.exceptions.AmbiguousTimeError: +... print('pytz.exceptions.AmbiguousTimeError: %s' % ambiguous) +pytz.exceptions.AmbiguousTimeError: 2009-10-31 23:30:00 +>>> try: +... tz.tzname(ambiguous) +... except pytz.exceptions.AmbiguousTimeError: +... print('pytz.exceptions.AmbiguousTimeError: %s' % ambiguous) +pytz.exceptions.AmbiguousTimeError: 2009-10-31 23:30:00 + + +Problems with Localtime +~~~~~~~~~~~~~~~~~~~~~~~ + +The major problem we have to deal with is that certain datetimes +may occur twice in a year. For example, in the US/Eastern timezone +on the last Sunday morning in October, the following sequence +happens: + + - 01:00 EDT occurs + - 1 hour later, instead of 2:00am the clock is turned back 1 hour + and 01:00 happens again (this time 01:00 EST) + +In fact, every instant between 01:00 and 02:00 occurs twice. This means +that if you try and create a time in the 'US/Eastern' timezone +the standard datetime syntax, there is no way to specify if you meant +before of after the end-of-daylight-saving-time transition. Using the +pytz custom syntax, the best you can do is make an educated guess: + +>>> loc_dt = eastern.localize(datetime(2002, 10, 27, 1, 30, 00)) +>>> loc_dt.strftime(fmt) +'2002-10-27 01:30:00 EST-0500' + +As you can see, the system has chosen one for you and there is a 50% +chance of it being out by one hour. For some applications, this does +not matter. However, if you are trying to schedule meetings with people +in different timezones or analyze log files it is not acceptable. + +The best and simplest solution is to stick with using UTC. The pytz +package encourages using UTC for internal timezone representation by +including a special UTC implementation based on the standard Python +reference implementation in the Python documentation. + +The UTC timezone unpickles to be the same instance, and pickles to a +smaller size than other pytz tzinfo instances. The UTC implementation +can be obtained as pytz.utc, pytz.UTC, or pytz.timezone('UTC'). + +>>> import pickle, pytz +>>> dt = datetime(2005, 3, 1, 14, 13, 21, tzinfo=utc) +>>> naive = dt.replace(tzinfo=None) +>>> p = pickle.dumps(dt, 1) +>>> naive_p = pickle.dumps(naive, 1) +>>> len(p) - len(naive_p) +17 +>>> new = pickle.loads(p) +>>> new == dt +True +>>> new is dt +False +>>> new.tzinfo is dt.tzinfo +True +>>> pytz.utc is pytz.UTC is pytz.timezone('UTC') +True + +Note that some other timezones are commonly thought of as the same (GMT, +Greenwich, Universal, etc.). The definition of UTC is distinct from these +other timezones, and they are not equivalent. For this reason, they will +not compare the same in Python. + +>>> utc == pytz.timezone('GMT') +False + +See the section `What is UTC`_, below. + +If you insist on working with local times, this library provides a +facility for constructing them unambiguously: + +>>> loc_dt = datetime(2002, 10, 27, 1, 30, 00) +>>> est_dt = eastern.localize(loc_dt, is_dst=True) +>>> edt_dt = eastern.localize(loc_dt, is_dst=False) +>>> print(est_dt.strftime(fmt) + ' / ' + edt_dt.strftime(fmt)) +2002-10-27 01:30:00 EDT-0400 / 2002-10-27 01:30:00 EST-0500 + +If you pass None as the is_dst flag to localize(), pytz will refuse to +guess and raise exceptions if you try to build ambiguous or non-existent +times. + +For example, 1:30am on 27th Oct 2002 happened twice in the US/Eastern +timezone when the clocks where put back at the end of Daylight Saving +Time: + +>>> dt = datetime(2002, 10, 27, 1, 30, 00) +>>> try: +... eastern.localize(dt, is_dst=None) +... except pytz.exceptions.AmbiguousTimeError: +... print('pytz.exceptions.AmbiguousTimeError: %s' % dt) +pytz.exceptions.AmbiguousTimeError: 2002-10-27 01:30:00 + +Similarly, 2:30am on 7th April 2002 never happened at all in the +US/Eastern timezone, as the clocks where put forward at 2:00am skipping +the entire hour: + +>>> dt = datetime(2002, 4, 7, 2, 30, 00) +>>> try: +... eastern.localize(dt, is_dst=None) +... except pytz.exceptions.NonExistentTimeError: +... print('pytz.exceptions.NonExistentTimeError: %s' % dt) +pytz.exceptions.NonExistentTimeError: 2002-04-07 02:30:00 + +Both of these exceptions share a common base class to make error handling +easier: + +>>> isinstance(pytz.AmbiguousTimeError(), pytz.InvalidTimeError) +True +>>> isinstance(pytz.NonExistentTimeError(), pytz.InvalidTimeError) +True + + +A special case is where countries change their timezone definitions +with no daylight savings time switch. For example, in 1915 Warsaw +switched from Warsaw time to Central European time with no daylight savings +transition. So at the stroke of midnight on August 5th 1915 the clocks +were wound back 24 minutes creating an ambiguous time period that cannot +be specified without referring to the timezone abbreviation or the +actual UTC offset. In this case midnight happened twice, neither time +during a daylight saving time period. pytz handles this transition by +treating the ambiguous period before the switch as daylight savings +time, and the ambiguous period after as standard time. + + +>>> warsaw = pytz.timezone('Europe/Warsaw') +>>> amb_dt1 = warsaw.localize(datetime(1915, 8, 4, 23, 59, 59), is_dst=True) +>>> amb_dt1.strftime(fmt) +'1915-08-04 23:59:59 WMT+0124' +>>> amb_dt2 = warsaw.localize(datetime(1915, 8, 4, 23, 59, 59), is_dst=False) +>>> amb_dt2.strftime(fmt) +'1915-08-04 23:59:59 CET+0100' +>>> switch_dt = warsaw.localize(datetime(1915, 8, 5, 00, 00, 00), is_dst=False) +>>> switch_dt.strftime(fmt) +'1915-08-05 00:00:00 CET+0100' +>>> str(switch_dt - amb_dt1) +'0:24:01' +>>> str(switch_dt - amb_dt2) +'0:00:01' + +The best way of creating a time during an ambiguous time period is +by converting from another timezone such as UTC: + +>>> utc_dt = datetime(1915, 8, 4, 22, 36, tzinfo=pytz.utc) +>>> utc_dt.astimezone(warsaw).strftime(fmt) +'1915-08-04 23:36:00 CET+0100' + +The standard Python way of handling all these ambiguities is not to +handle them, such as demonstrated in this example using the US/Eastern +timezone definition from the Python documentation (Note that this +implementation only works for dates between 1987 and 2006 - it is +included for tests only!): + +>>> from pytz.reference import Eastern # pytz.reference only for tests +>>> dt = datetime(2002, 10, 27, 0, 30, tzinfo=Eastern) +>>> str(dt) +'2002-10-27 00:30:00-04:00' +>>> str(dt + timedelta(hours=1)) +'2002-10-27 01:30:00-05:00' +>>> str(dt + timedelta(hours=2)) +'2002-10-27 02:30:00-05:00' +>>> str(dt + timedelta(hours=3)) +'2002-10-27 03:30:00-05:00' + +Notice the first two results? At first glance you might think they are +correct, but taking the UTC offset into account you find that they are +actually two hours appart instead of the 1 hour we asked for. + +>>> from pytz.reference import UTC # pytz.reference only for tests +>>> str(dt.astimezone(UTC)) +'2002-10-27 04:30:00+00:00' +>>> str((dt + timedelta(hours=1)).astimezone(UTC)) +'2002-10-27 06:30:00+00:00' + + +Country Information +~~~~~~~~~~~~~~~~~~~ + +A mechanism is provided to access the timezones commonly in use +for a particular country, looked up using the ISO 3166 country code. +It returns a list of strings that can be used to retrieve the relevant +tzinfo instance using ``pytz.timezone()``: + +>>> print(' '.join(pytz.country_timezones['nz'])) +Pacific/Auckland Pacific/Chatham + +The Olson database comes with a ISO 3166 country code to English country +name mapping that pytz exposes as a dictionary: + +>>> print(pytz.country_names['nz']) +New Zealand + + +What is UTC +~~~~~~~~~~~ + +'UTC' is `Coordinated Universal Time`_. It is a successor to, but distinct +from, Greenwich Mean Time (GMT) and the various definitions of Universal +Time. UTC is now the worldwide standard for regulating clocks and time +measurement. + +All other timezones are defined relative to UTC, and include offsets like +UTC+0800 - hours to add or subtract from UTC to derive the local time. No +daylight saving time occurs in UTC, making it a useful timezone to perform +date arithmetic without worrying about the confusion and ambiguities caused +by daylight saving time transitions, your country changing its timezone, or +mobile computers that roam through multiple timezones. + +.. _Coordinated Universal Time: https://en.wikipedia.org/wiki/Coordinated_Universal_Time + + +Helpers +~~~~~~~ + +There are two lists of timezones provided. + +``all_timezones`` is the exhaustive list of the timezone names that can +be used. + +>>> from pytz import all_timezones +>>> len(all_timezones) >= 500 +True +>>> 'Etc/Greenwich' in all_timezones +True + +``common_timezones`` is a list of useful, current timezones. It doesn't +contain deprecated zones or historical zones, except for a few I've +deemed in common usage, such as US/Eastern (open a bug report if you +think other timezones are deserving of being included here). It is also +a sequence of strings. + +>>> from pytz import common_timezones +>>> len(common_timezones) < len(all_timezones) +True +>>> 'Etc/Greenwich' in common_timezones +False +>>> 'Australia/Melbourne' in common_timezones +True +>>> 'US/Eastern' in common_timezones +True +>>> 'Canada/Eastern' in common_timezones +True +>>> 'Australia/Yancowinna' in all_timezones +True +>>> 'Australia/Yancowinna' in common_timezones +False + +Both ``common_timezones`` and ``all_timezones`` are alphabetically +sorted: + +>>> common_timezones_dupe = common_timezones[:] +>>> common_timezones_dupe.sort() +>>> common_timezones == common_timezones_dupe +True +>>> all_timezones_dupe = all_timezones[:] +>>> all_timezones_dupe.sort() +>>> all_timezones == all_timezones_dupe +True + +``all_timezones`` and ``common_timezones`` are also available as sets. + +>>> from pytz import all_timezones_set, common_timezones_set +>>> 'US/Eastern' in all_timezones_set +True +>>> 'US/Eastern' in common_timezones_set +True +>>> 'Australia/Victoria' in common_timezones_set +False + +You can also retrieve lists of timezones used by particular countries +using the ``country_timezones()`` function. It requires an ISO-3166 +two letter country code. + +>>> from pytz import country_timezones +>>> print(' '.join(country_timezones('ch'))) +Europe/Zurich +>>> print(' '.join(country_timezones('CH'))) +Europe/Zurich + + +Internationalization - i18n/l10n +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pytz is an interface to the IANA database, which uses ASCII names. The `Unicode Consortium's Unicode Locales (CLDR) `_ +project provides translations. Python packages such as +`Babel `_ +and Thomas Khyn's `l18n `_ package can be used +to access these translations from Python. + + +License +~~~~~~~ + +MIT license. + +This code is also available as part of Zope 3 under the Zope Public +License, Version 2.1 (ZPL). + +I'm happy to relicense this code if necessary for inclusion in other +open source projects. + + +Latest Versions +~~~~~~~~~~~~~~~ + +This package will be updated after releases of the Olson timezone +database. The latest version can be downloaded from the `Python Package +Index `_. The code that is used +to generate this distribution is hosted on Github and available +using git:: + + git clone https://github.com/stub42/pytz.git + +Announcements of new releases are made on +`Launchpad `_, and the +`Atom feed `_ +hosted there. + + +Bugs, Feature Requests & Patches +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Bugs should be reported on `Github `_. +Feature requests are unlikely to be considered, and efforts instead directed +to timezone support now built into Python or packages that work with it. + + +Security Issues +~~~~~~~~~~~~~~~ + +Reports about security issues can be made via `Tidelift `_. + + +Issues & Limitations +~~~~~~~~~~~~~~~~~~~~ + +- This project is in maintenance mode. Projects using Python 3.9 or later + are best served by using the timezone functionaly now included in core + Python and packages that work with it such as `tzdata `_. + +- Offsets from UTC are rounded to the nearest whole minute, so timezones + such as Europe/Amsterdam pre 1937 will be up to 30 seconds out. This + was a limitation of the Python datetime library. + +- If you think a timezone definition is incorrect, I probably can't fix + it. pytz is a direct translation of the Olson timezone database, and + changes to the timezone definitions need to be made to this source. + If you find errors they should be reported to the time zone mailing + list, linked from http://www.iana.org/time-zones. + + +Further Reading +~~~~~~~~~~~~~~~ + +More info than you want to know about timezones: +https://data.iana.org/time-zones/tz-link.html + + +Contact +~~~~~~~ + +Stuart Bishop diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/RECORD b/lib/python3.12/site-packages/pytz-2025.2.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..f9afc2a04b984a8013248b1e514ed098c0adf1a8 --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/RECORD @@ -0,0 +1,623 @@ +pytz-2025.2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pytz-2025.2.dist-info/LICENSE.txt,sha256=vosaN-vibFkqkPbA6zMQOn84POL010mMCvmlJpkKB7g,1088 +pytz-2025.2.dist-info/METADATA,sha256=5iDk4fnxyAGWGTiNTaSW6wIRmhPzWqrVXwImlKPIH2w,22374 +pytz-2025.2.dist-info/RECORD,, +pytz-2025.2.dist-info/WHEEL,sha256=-G_t0oGuE7UD0DrSpVZnq1hHMBV9DD2XkS5v7XpmTnk,110 +pytz-2025.2.dist-info/top_level.txt,sha256=6xRYlt934v1yHb1JIrXgHyGxn3cqACvd-yE8ski_kcc,5 +pytz-2025.2.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 +pytz/__init__.py,sha256=dml-pbWn-1xqv_ZWHoKoPmhruaO63GoIubnxtFnEvDc,35125 +pytz/__pycache__/__init__.cpython-312.pyc,, +pytz/__pycache__/exceptions.cpython-312.pyc,, +pytz/__pycache__/lazy.cpython-312.pyc,, +pytz/__pycache__/reference.cpython-312.pyc,, +pytz/__pycache__/tzfile.cpython-312.pyc,, +pytz/__pycache__/tzinfo.cpython-312.pyc,, +pytz/exceptions.py,sha256=434ZcuLlpLQY9mWoGq7zJMV1TyiYvVgpKBU1qZkbDjM,1571 +pytz/lazy.py,sha256=toeR5uDWKBj6ezsUZ4elNP6CEMtK7CO2jS9A30nsFbo,5404 +pytz/reference.py,sha256=zUtCki7JFEmrzrjNsfMD7YL0lWDxynKc1Ubo4iXSs74,3778 +pytz/tzfile.py,sha256=K2y7pZs4vydpZVftrfAA_-hgw17y1Szc7z_QCse6udU,4723 +pytz/tzinfo.py,sha256=XfaVOoO3KsCvtUYaCd0fvgBXWZ8tgevGYUoBh_uiE60,19340 +pytz/zoneinfo/Africa/Abidjan,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Accra,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Addis_Ababa,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Algiers,sha256=vaFpjNVCwObnbfu82rOQzdJvN6nVgmpXpQ1aqzfzsqY,735 +pytz/zoneinfo/Africa/Asmara,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Asmera,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Bamako,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Bangui,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Banjul,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Bissau,sha256=IjuxDP6EZiDHFvl_bHS6NN7sdRxLKXllooBC829poak,194 +pytz/zoneinfo/Africa/Blantyre,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Brazzaville,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Bujumbura,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Cairo,sha256=Lft-GCLQhaSJm9VqUmsEFoHIS1Vhfa7pFJn9GZCpifs,2399 +pytz/zoneinfo/Africa/Casablanca,sha256=4RqVbw_F3ZucopIC2ivAJ8WDwj5wRODAB67tBpdXcgA,2429 +pytz/zoneinfo/Africa/Ceuta,sha256=Cw-2_nFDGbN8WqIsVpcauyZooWX8j3Kmx2PnC0fHut8,2052 +pytz/zoneinfo/Africa/Conakry,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Dakar,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Dar_es_Salaam,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Djibouti,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Douala,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/El_Aaiun,sha256=UWCCqQLJxd8qsTYw82kz9W1suwW5TRgnZw31sDWDz20,2295 +pytz/zoneinfo/Africa/Freetown,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Gaborone,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Harare,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Johannesburg,sha256=bBvMdSZo53WFowiuhUO9C8zY6BOGViboCb-U8_49l34,246 +pytz/zoneinfo/Africa/Juba,sha256=UVnIqEPJwHLTMC-r5qZQHNv9opoYVsKdq-ta_5XUw_Q,679 +pytz/zoneinfo/Africa/Kampala,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Khartoum,sha256=MYWDoJ3AcCItZdApoeOgtWWDDxquwTon5v5TOGP70-o,679 +pytz/zoneinfo/Africa/Kigali,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Kinshasa,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Lagos,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Libreville,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Lome,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Luanda,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Lubumbashi,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Lusaka,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Malabo,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Maputo,sha256=RE7TpxBBS8a_Q-sn5ZHaSdO-PbFTRJpqDJRz9-Of28s,149 +pytz/zoneinfo/Africa/Maseru,sha256=bBvMdSZo53WFowiuhUO9C8zY6BOGViboCb-U8_49l34,246 +pytz/zoneinfo/Africa/Mbabane,sha256=bBvMdSZo53WFowiuhUO9C8zY6BOGViboCb-U8_49l34,246 +pytz/zoneinfo/Africa/Mogadishu,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Monrovia,sha256=-VsJW5cU4KdvfgYaQVv4lcuzmaKIVFMd42nO6RXOBdU,208 +pytz/zoneinfo/Africa/Nairobi,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Africa/Ndjamena,sha256=8T3A0Zm9Gj0Bvm6rd88t3GAXKiKdGUfHlIqYlkYI0KM,199 +pytz/zoneinfo/Africa/Niamey,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Nouakchott,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Ouagadougou,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Porto-Novo,sha256=z_6wKCzL1_ug5JP_hneh5abdUZeIUELkN_ladz-ESEY,235 +pytz/zoneinfo/Africa/Sao_Tome,sha256=MdjxpQ268uzJ7Zx1ZroFUtRUwqsJ6F_yY3AYV9FXw1I,254 +pytz/zoneinfo/Africa/Timbuktu,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Africa/Tripoli,sha256=W1dptGD70T7ppGoo0fczFQeDiIp0nultLNPV66MwB2c,625 +pytz/zoneinfo/Africa/Tunis,sha256=OFVMEM4eYT2Ez0beuhEUCTSIpcFldWxsV2uEoTZIUNI,689 +pytz/zoneinfo/Africa/Windhoek,sha256=xuhvudrMH4alnVmouSTQI8YL8F_HbgsF2EQ7AZKzuHs,955 +pytz/zoneinfo/America/Adak,sha256=IB1DhwJQAKbhPJ9jHLf8zW5Dad7HIkBS-dhv64E1OlM,2356 +pytz/zoneinfo/America/Anchorage,sha256=oZA1NSPS2BWdymYpnCHFO8BlYVS-ll5KLg2Ez9CbETs,2371 +pytz/zoneinfo/America/Anguilla,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Antigua,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Araguaina,sha256=G6v9wYFZ8EB4WQfIsqRbbiiKd2b27j7Zt5dFjBbzx2o,870 +pytz/zoneinfo/America/Argentina/Buenos_Aires,sha256=JmU8lBwmy29gR6OmeytvFdMRx6ObJKnYNHmLyMmXX2M,1062 +pytz/zoneinfo/America/Argentina/Catamarca,sha256=uMCJXXGYmNESHVvj5RYBZ0McrOdE14hwm17l25MgRW0,1062 +pytz/zoneinfo/America/Argentina/ComodRivadavia,sha256=uMCJXXGYmNESHVvj5RYBZ0McrOdE14hwm17l25MgRW0,1062 +pytz/zoneinfo/America/Argentina/Cordoba,sha256=uniNihhMHnr4XK4WpwiPUnrAT0YPmvzqB6f0hRLtXvY,1062 +pytz/zoneinfo/America/Argentina/Jujuy,sha256=PGmAehypCxj0XCenCSWqylDIPbKLK0DlrwJK_24D590,1034 +pytz/zoneinfo/America/Argentina/La_Rioja,sha256=Um6XoVXhsr62ad1mWuebe6NY0ZHauBdR9tMGDgqCOHg,1076 +pytz/zoneinfo/America/Argentina/Mendoza,sha256=xcOVtvRyVYFAU90y2QYwpyQhpMLyAp7-Fxvku4kgl0c,1062 +pytz/zoneinfo/America/Argentina/Rio_Gallegos,sha256=F9ZKR4o8gLHX7QBuIjMapGIdmzJxpqwbouPgZ5MqDpY,1062 +pytz/zoneinfo/America/Argentina/Salta,sha256=h1KYrDNIapvDkYhi1PaB8WD1qWOe4vhhgDJWDCGV4jc,1034 +pytz/zoneinfo/America/Argentina/San_Juan,sha256=AI2GltA80mPNzhHxYycuEwIbO1ANXyIqBQZMpjqKqdQ,1076 +pytz/zoneinfo/America/Argentina/San_Luis,sha256=2ItGRcLVK2wx8MyJsHbIBBeAkU4B-MN5x1ZxNyZ7UJE,1088 +pytz/zoneinfo/America/Argentina/Tucuman,sha256=twO-FqtNJV8XOzWTvFQ-xnEcWCoDUHY3gpVIG0Mzbf8,1090 +pytz/zoneinfo/America/Argentina/Ushuaia,sha256=A6IbpVlY9IIPoSKMFRR9DMROdwXUSDc2HsASueOSnqo,1062 +pytz/zoneinfo/America/Aruba,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Asuncion,sha256=9yfGiENhXNgd2_0xHVsF-hLeJkGSqRs-zD9O945KJWc,1644 +pytz/zoneinfo/America/Atikokan,sha256=kayA_pdpMcSQ0FjIzotdcf-m1JYfbKE-qcFT8LC8zqA,182 +pytz/zoneinfo/America/Atka,sha256=IB1DhwJQAKbhPJ9jHLf8zW5Dad7HIkBS-dhv64E1OlM,2356 +pytz/zoneinfo/America/Bahia,sha256=qi7dA6FofDhLxVMmd2L8bK3HeaQnc9X-jiijwyfhs3g,1010 +pytz/zoneinfo/America/Bahia_Banderas,sha256=MvrXGJ5LzaHOeguJqxszxjxMhVafGVbk-ojXEc7_YEI,1100 +pytz/zoneinfo/America/Barbados,sha256=ima-Qrrhazu4Qfvu2Z0-e6E-GTiYknuJBu6c2yVG9LE,436 +pytz/zoneinfo/America/Belem,sha256=aZMUgtFDdHNISpqyQRYbmS2IBD-BAS3CaJnhu6onLCY,562 +pytz/zoneinfo/America/Belize,sha256=pkfLY2KfPchbeJa1pWcXmWAwp4ZlRvxWLVezXnrbkws,1614 +pytz/zoneinfo/America/Blanc-Sablon,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Boa_Vista,sha256=dMtaG11kGlJrgJJgGWEDZZAmnO_HfT3L4X8pI72LLFY,618 +pytz/zoneinfo/America/Bogota,sha256=Z1ernZZGQxulE8KFWHYWcM3SV1jn2_QEc1Q0OJzHRak,232 +pytz/zoneinfo/America/Boise,sha256=7HQsNPJiUheQgFz5kVLvTnf5xhXAYaeANqDskxKz2Vs,2410 +pytz/zoneinfo/America/Buenos_Aires,sha256=JmU8lBwmy29gR6OmeytvFdMRx6ObJKnYNHmLyMmXX2M,1062 +pytz/zoneinfo/America/Cambridge_Bay,sha256=_4xRlX3WdVpEcqoT6myD7NeTCXnn9OYk_iH006bwULo,2254 +pytz/zoneinfo/America/Campo_Grande,sha256=gINiXg5i2e6Rh2Nbo2bFqhPAJL4F4cAqGnBankXTDXw,1430 +pytz/zoneinfo/America/Cancun,sha256=EdV0Nw2WjM7VnjFHoq5jsSbLuuE7eP1OE74utEyWJG4,864 +pytz/zoneinfo/America/Caracas,sha256=mUNMFdDzZLav_ePA1ocBdmqVBierkeEszTIFpNCm5J0,250 +pytz/zoneinfo/America/Catamarca,sha256=uMCJXXGYmNESHVvj5RYBZ0McrOdE14hwm17l25MgRW0,1062 +pytz/zoneinfo/America/Cayenne,sha256=4k7Iv1woX4atqePKrcvMQD2Vk9Tmma7rW_AW_R62pCc,184 +pytz/zoneinfo/America/Cayman,sha256=kayA_pdpMcSQ0FjIzotdcf-m1JYfbKE-qcFT8LC8zqA,182 +pytz/zoneinfo/America/Chicago,sha256=_roybr6I6sIAF6cYdIxGxoRpoef153Fty48dQ6bm9oY,3592 +pytz/zoneinfo/America/Chihuahua,sha256=3Ngzbedg8AzAqxsbQSG0jVRx-LxYlw1i3kx-Yzl-2Ic,1102 +pytz/zoneinfo/America/Ciudad_Juarez,sha256=ir4b27DiFrhL0H4fZQ92nEa-BBoPfLWIz3phU373dgE,1538 +pytz/zoneinfo/America/Coral_Harbour,sha256=kayA_pdpMcSQ0FjIzotdcf-m1JYfbKE-qcFT8LC8zqA,182 +pytz/zoneinfo/America/Cordoba,sha256=uniNihhMHnr4XK4WpwiPUnrAT0YPmvzqB6f0hRLtXvY,1062 +pytz/zoneinfo/America/Costa_Rica,sha256=74rYa6lrgIkyls9PkHo8SCYl9oOqiuG5S7MWdnJelP4,316 +pytz/zoneinfo/America/Coyhaique,sha256=beZXU6Lw5jx2Wp3dNRkFiofXGma2dM6DsqvJT9CSHdc,2126 +pytz/zoneinfo/America/Creston,sha256=illz0sYuLL8lIPK0Tkou6dL0Vck_D0W_3rRTOvFYRmQ,360 +pytz/zoneinfo/America/Cuiaba,sha256=GRJqkhRXNsOUcgjZddQxRIJdRYaw9pM_YLWbun88dkg,1402 +pytz/zoneinfo/America/Curacao,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Danmarkshavn,sha256=YRZAfUCoVtaL1L-MYMYMH1wyOaVQnfUo_gFnvMXSuzw,698 +pytz/zoneinfo/America/Dawson,sha256=rAHhyuMuyjf_eyA2SBG76MRBf_fj_xi5FAuiWVQgJhw,1614 +pytz/zoneinfo/America/Dawson_Creek,sha256=aJXCyP4j3ggE4wGCN-LrS9hpD_5zWHzQTeSAKTWEPUM,1050 +pytz/zoneinfo/America/Denver,sha256=MugZwApDs8NI9TnXANQlUE8guNBowWQY0m-ptpPndck,2460 +pytz/zoneinfo/America/Detroit,sha256=hecz8yqY2Cj5B61G3gLZdAVZvRgK9l0P90c_gN-uD5g,2230 +pytz/zoneinfo/America/Dominica,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Edmonton,sha256=-TkIfc3QlvaCf0p8COZ43Y1HRBAl-nARUi-JdXeK1vE,2332 +pytz/zoneinfo/America/Eirunepe,sha256=j5eExkjFaqtC-D8XK0rGzoF9yEgbSlTbPqVG9WKhEa8,642 +pytz/zoneinfo/America/El_Salvador,sha256=gvGN8Lkj-sGm2_rs8OUjAMf1oMtKp2Xes6UfWT0WqgU,224 +pytz/zoneinfo/America/Ensenada,sha256=SluV7xzZm24LgMXSUVt1cD1AlE7y_bdE65HhDIdXLcs,2458 +pytz/zoneinfo/America/Fort_Nelson,sha256=erfODr3DrSpz65kAdO7Ts2dGbZxvddEP6gx4BX3y2J0,2240 +pytz/zoneinfo/America/Fort_Wayne,sha256=kNKy9Kj9ICsiYYfCCbAggzMA7exf-GpGPMxoXocHUyw,1682 +pytz/zoneinfo/America/Fortaleza,sha256=rjiSB0q1cBuMDOM9orW_uwe5UOLBwTlfjFotwOYe1mU,702 +pytz/zoneinfo/America/Glace_Bay,sha256=G8DGLGCapH_aYCF_OhaL5Qonf7FOAgAPwelO5htCWBc,2192 +pytz/zoneinfo/America/Godthab,sha256=KGXrMN-YkYpVCgLdpcfwMFQ77EsRAGsjUCG3yAUvVfw,1889 +pytz/zoneinfo/America/Goose_Bay,sha256=JgaLueghSvX2g725FOfIgpgvsqxZGykWOhAZWGpQZRY,3210 +pytz/zoneinfo/America/Grand_Turk,sha256=4YOFEPK60Bel2_fCsY6vSZxUcMJKjiKtyOf_Q0khEwU,1834 +pytz/zoneinfo/America/Grenada,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Guadeloupe,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Guatemala,sha256=dugUgCd6QY52yHkHuUP4jRWzo5x439IQigaYCvEF46Q,280 +pytz/zoneinfo/America/Guayaquil,sha256=j2UuIo-4RgSOlTNfu77mhZ92waNTeKFSvmoVemJooT0,232 +pytz/zoneinfo/America/Guyana,sha256=R0bOvCRDC8SRIexmhsduPdHbbRPwI2GviD9otExiUrk,248 +pytz/zoneinfo/America/Halifax,sha256=TZpmc5PwWoLfTfQoQ_b3U17BE2iVKSeNkR0Ho8mbTn8,3424 +pytz/zoneinfo/America/Havana,sha256=HUQeAuKBsEkI5SLZjqynXICOUVOajkKzKH5r-Ov5Odc,2416 +pytz/zoneinfo/America/Hermosillo,sha256=ixYKestLmS7gWobk9Kq6FtLZo1yqbWActrFUKluzctw,388 +pytz/zoneinfo/America/Indiana/Indianapolis,sha256=kNKy9Kj9ICsiYYfCCbAggzMA7exf-GpGPMxoXocHUyw,1682 +pytz/zoneinfo/America/Indiana/Knox,sha256=CsvZ5BKw2qVav3x_F8CU9taJdDk7jX41Cfsqms6jXV8,2444 +pytz/zoneinfo/America/Indiana/Marengo,sha256=f3tQ-lgMSUA7nvn64pXhKtJL7mWzGajoCega5MEJSbI,1738 +pytz/zoneinfo/America/Indiana/Petersburg,sha256=A88OHuM0Rg3iMLHjKgXq_d2jZCdVSytUQs-9W0KcFyQ,1920 +pytz/zoneinfo/America/Indiana/Tell_City,sha256=4dWqAr9Y2BXfL4pAQk-81c3gGl2cNdHXOD7_wJhhhn8,1700 +pytz/zoneinfo/America/Indiana/Vevay,sha256=H7VR2G-_sD_C5Rm4P3g1iRC1FWCPg4m0MGD3P1PLzsk,1430 +pytz/zoneinfo/America/Indiana/Vincennes,sha256=62mAxT7APFCaoygflnEzdOpe-fuW1yObI6m6EUUcS7A,1710 +pytz/zoneinfo/America/Indiana/Winamac,sha256=aZGM2jR8CH9BHSUq7XygiweDd6dorXLPXg246XsbR6s,1794 +pytz/zoneinfo/America/Indianapolis,sha256=kNKy9Kj9ICsiYYfCCbAggzMA7exf-GpGPMxoXocHUyw,1682 +pytz/zoneinfo/America/Inuvik,sha256=6J-mapDnrk9A1LtswoE34tqSy_ufedcEBNxixkrEjIo,2074 +pytz/zoneinfo/America/Iqaluit,sha256=feOnxAN0N0r-M1qlkrA4JMyawoc0tqae0iiBCPDAs4k,2202 +pytz/zoneinfo/America/Jamaica,sha256=wlagieUPRf5-beie-h7QsONbNzjGsm8vMs8uf28pw28,482 +pytz/zoneinfo/America/Jujuy,sha256=PGmAehypCxj0XCenCSWqylDIPbKLK0DlrwJK_24D590,1034 +pytz/zoneinfo/America/Juneau,sha256=k7hxb0aGRnfnE-DBi3LkcjAzRPyAf0_Hw0vVFfjGeb0,2353 +pytz/zoneinfo/America/Kentucky/Louisville,sha256=tP072xV_n_vIQjxxcJ77AGeGj6yL1KPpn3fwids9g1U,2788 +pytz/zoneinfo/America/Kentucky/Monticello,sha256=LtdyCo85BrXQs6rlH61Ym-8KqWHH6PwAOjD0QxhIdzM,2368 +pytz/zoneinfo/America/Knox_IN,sha256=CsvZ5BKw2qVav3x_F8CU9taJdDk7jX41Cfsqms6jXV8,2444 +pytz/zoneinfo/America/Kralendijk,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/La_Paz,sha256=hqfD8LQHupdZhji2e93_9pOQAT-R7muzzjP0nyfbFXY,218 +pytz/zoneinfo/America/Lima,sha256=HHgTnDUnCZzibvL0MrG8qyOuvjmYYw3e3R5VbnxMZs8,392 +pytz/zoneinfo/America/Los_Angeles,sha256=aJd7ua1tGG_vxser02AQpm4wAI3LLTdgh6QcSYYecmg,2852 +pytz/zoneinfo/America/Louisville,sha256=tP072xV_n_vIQjxxcJ77AGeGj6yL1KPpn3fwids9g1U,2788 +pytz/zoneinfo/America/Lower_Princes,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Maceio,sha256=3R5DlSe32kQDmoSVIWpcyk2o7qohr-rliwqDSGFIMyQ,730 +pytz/zoneinfo/America/Managua,sha256=xBzF01AHn2E2fD8Qdy-DHFe36UqoeNpKPfChduBKWdk,430 +pytz/zoneinfo/America/Manaus,sha256=F6RLOOeOi9lymZiQmQ9pR8tFpPZ6EguNdPfOc6BhXDE,590 +pytz/zoneinfo/America/Marigot,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Martinique,sha256=fMs80kOU2YFvC0f9y2eje97JeAtTYBamXrnlTunNLzQ,232 +pytz/zoneinfo/America/Matamoros,sha256=fq-PqdmZrQ98UsFmHA9ivjBZv5GEBRTOuLQ5Cu5ajW8,1418 +pytz/zoneinfo/America/Mazatlan,sha256=BWH2NqVPA1PsyELPN_2BF8KllrsmQkqg1eujsQvnnx8,1060 +pytz/zoneinfo/America/Mendoza,sha256=xcOVtvRyVYFAU90y2QYwpyQhpMLyAp7-Fxvku4kgl0c,1062 +pytz/zoneinfo/America/Menominee,sha256=Arv9WLbfhNcpRsUjHDU757BEdwlp08Gt30AixG3gZ04,2274 +pytz/zoneinfo/America/Merida,sha256=SVNEHCazjomftnuPVBayFI-E-IQ0WmluHfTpHP0h3d0,1004 +pytz/zoneinfo/America/Metlakatla,sha256=twmieGTVY2V-U8nFxqvx7asYv8GVjeWdLtrOI7UApVI,1423 +pytz/zoneinfo/America/Mexico_City,sha256=Uog2-FMWz2o12jR6sK9vemJamLeo6OEFMQR3s0xTxkc,1222 +pytz/zoneinfo/America/Miquelon,sha256=l5txBJYe9HTRZlILcbSL_HNDYrjUb0ouecNy7QEkg9c,1652 +pytz/zoneinfo/America/Moncton,sha256=Wmv-bk9aKKcWWzOpc1UFu67HOfwaIk2Wmh3LgqGctys,3154 +pytz/zoneinfo/America/Monterrey,sha256=YixTESJubf6ZBUXy6g32hAM2gR4GXXPqOU4tv0L3kG0,1114 +pytz/zoneinfo/America/Montevideo,sha256=dQEBE4mjZPtyRjKXK6Z-bMHJdFqpwhIzxDH4x04rKYk,1496 +pytz/zoneinfo/America/Montreal,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/America/Montserrat,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Nassau,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/America/New_York,sha256=6e0H177gx2qdRC0JHvHwFmj-58TyYBTAqGixn-bBipU,3552 +pytz/zoneinfo/America/Nipigon,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/America/Nome,sha256=2izM3-P-PqJ9za6MdhzFfMvPFNq7Gim69tAvEwPeY2s,2367 +pytz/zoneinfo/America/Noronha,sha256=feeRAijQqKylZgqe84nKhsFLycT5zIBm7mLIvdyGw4w,702 +pytz/zoneinfo/America/North_Dakota/Beulah,sha256=qtgbqNu8M3AkHF2n-_oSps1pYT4SxgclbkkPKbXaBHs,2396 +pytz/zoneinfo/America/North_Dakota/Center,sha256=9ZWbK9YKkquULyBUFS3Lr_idxbt7V7y4W4EO0Kn20sw,2396 +pytz/zoneinfo/America/North_Dakota/New_Salem,sha256=DH_bsQfuUnK2obdb06KgisO4XLqht12BXdrgUsZZveg,2396 +pytz/zoneinfo/America/Nuuk,sha256=KGXrMN-YkYpVCgLdpcfwMFQ77EsRAGsjUCG3yAUvVfw,1889 +pytz/zoneinfo/America/Ojinaga,sha256=b38Q_7VdkCZzaVwb7OXuddihJAzUKPTTqXcmpBm1ntE,1524 +pytz/zoneinfo/America/Panama,sha256=kayA_pdpMcSQ0FjIzotdcf-m1JYfbKE-qcFT8LC8zqA,182 +pytz/zoneinfo/America/Pangnirtung,sha256=feOnxAN0N0r-M1qlkrA4JMyawoc0tqae0iiBCPDAs4k,2202 +pytz/zoneinfo/America/Paramaribo,sha256=Z7UZvNlgd-qEUHjEPYXIkLNTgjMcCzk9EfUUEmUyd7M,248 +pytz/zoneinfo/America/Phoenix,sha256=illz0sYuLL8lIPK0Tkou6dL0Vck_D0W_3rRTOvFYRmQ,360 +pytz/zoneinfo/America/Port-au-Prince,sha256=09ZAJd4IOiMpfdpUuF1U44R_hRt6BvpAkFXOnYO9yOM,1434 +pytz/zoneinfo/America/Port_of_Spain,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Porto_Acre,sha256=0gpJUl46hQbp0P6Xj1S0NArIWeAryuuDXjsldvB5GHE,614 +pytz/zoneinfo/America/Porto_Velho,sha256=uSMV2hZWj-VyBhFBwC950wcThfN3jq6KlycESmQTLOA,562 +pytz/zoneinfo/America/Puerto_Rico,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Punta_Arenas,sha256=tR5uIf1351AWFqrqNtmXnhQWnKREmJaZqKBzaWRVMTQ,1902 +pytz/zoneinfo/America/Rainy_River,sha256=7P-_YQrneFcon7QKSTOnkiGjEppFDn3Z48MJ1qq8VBw,2868 +pytz/zoneinfo/America/Rankin_Inlet,sha256=nXgqjL3O2BV0em-Xk8qVRRZb_X0yQmHE6vmSSvI9Kzc,2066 +pytz/zoneinfo/America/Recife,sha256=bJ_HE0-JFio4-owpZ0pLO8U3ai0fiGu8QHL0DexLiLc,702 +pytz/zoneinfo/America/Regina,sha256=yjqT08pHbICYe83H8JmtaDBvCFqRv7Tfze3Y8xuXukw,980 +pytz/zoneinfo/America/Resolute,sha256=CnMU2dBI-63vt8-J0Q1Ropx-8b9pRCLjhvrycMIedGg,2066 +pytz/zoneinfo/America/Rio_Branco,sha256=0gpJUl46hQbp0P6Xj1S0NArIWeAryuuDXjsldvB5GHE,614 +pytz/zoneinfo/America/Rosario,sha256=uniNihhMHnr4XK4WpwiPUnrAT0YPmvzqB6f0hRLtXvY,1062 +pytz/zoneinfo/America/Santa_Isabel,sha256=SluV7xzZm24LgMXSUVt1cD1AlE7y_bdE65HhDIdXLcs,2458 +pytz/zoneinfo/America/Santarem,sha256=VmZP9S5pPucFxyqAOV908EmWXQZvgCgWLmlJJTUl0LE,588 +pytz/zoneinfo/America/Santiago,sha256=0CDw13dCMUsoquMupoJgupkzAUNhDK6E0lVxURA7osA,2515 +pytz/zoneinfo/America/Santo_Domingo,sha256=DKtaEj8fQ92ybITTWU4Bm160S9pzJmUVbjaWRnenxU4,458 +pytz/zoneinfo/America/Sao_Paulo,sha256=BMBnRO4_4HjvO4t3njjrMGZr-ZPmegkvyvL8KPY6ZM4,1430 +pytz/zoneinfo/America/Scoresbysund,sha256=K-qkiMCCFgOe8ccPMABA-lDjc9vb6wpluBOCVfiBdLI,1935 +pytz/zoneinfo/America/Shiprock,sha256=MugZwApDs8NI9TnXANQlUE8guNBowWQY0m-ptpPndck,2460 +pytz/zoneinfo/America/Sitka,sha256=aiS7Fk37hZpzZ9VkeJQeF-BqTLRC1QOTCgMAJwT8UxA,2329 +pytz/zoneinfo/America/St_Barthelemy,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/St_Johns,sha256=r1-17uKv27eZ3JsVkw_DLZQbo6wvjuuVu7C2pDsmOgI,3655 +pytz/zoneinfo/America/St_Kitts,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/St_Lucia,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/St_Thomas,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/St_Vincent,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Swift_Current,sha256=RRKOF7vZC8VvYxD8PP4J1_hUPayKBP7Lu80avRkfPDY,560 +pytz/zoneinfo/America/Tegucigalpa,sha256=EzOz7ntTlreMq69JZ2CcAb8Ps98V9bUMN480tpPIyw4,252 +pytz/zoneinfo/America/Thule,sha256=8xuPRaZU8RgO5ECqFYHYmnHioc81sBOailkVu8Y02i8,1502 +pytz/zoneinfo/America/Thunder_Bay,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/America/Tijuana,sha256=SluV7xzZm24LgMXSUVt1cD1AlE7y_bdE65HhDIdXLcs,2458 +pytz/zoneinfo/America/Toronto,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/America/Tortola,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Vancouver,sha256=sknKH0jSPWam-DHfM35qXs8Nam7d5TFlkUI9Sgxryyg,2892 +pytz/zoneinfo/America/Virgin,sha256=hJHlV_-AGoMGUWuMpZRv9fLmghrzFHfrR9fRkcxaZJc,246 +pytz/zoneinfo/America/Whitehorse,sha256=TrR6PCnYG-mSClBMohqlP8qnYhXMUsydI-L-quXFxyM,1614 +pytz/zoneinfo/America/Winnipeg,sha256=7P-_YQrneFcon7QKSTOnkiGjEppFDn3Z48MJ1qq8VBw,2868 +pytz/zoneinfo/America/Yakutat,sha256=tFwnKbvwhyyn4LNTAn5ye_JWDdxjCerNDt7oOwUwO2M,2305 +pytz/zoneinfo/America/Yellowknife,sha256=-TkIfc3QlvaCf0p8COZ43Y1HRBAl-nARUi-JdXeK1vE,2332 +pytz/zoneinfo/Antarctica/Casey,sha256=VeaLOxTfDyjfGXq5Ul95JEIMXNWHSW-0N3yOoS7VK-c,423 +pytz/zoneinfo/Antarctica/Davis,sha256=XB12dEq0Q-3XkzBNTNC7G1fzH-WxxctIuZqI3zp8ypI,283 +pytz/zoneinfo/Antarctica/DumontDUrville,sha256=nB36HBWZTdh3TlP0DLFNz1KRQ0aHIfHbp7LC4Urp9fA,172 +pytz/zoneinfo/Antarctica/Macquarie,sha256=ie7RlaU8RHTorVVj-MX8StKMqx_oXf4UH2PUqpzcwe0,2260 +pytz/zoneinfo/Antarctica/Mawson,sha256=EjIFbqRdr2ZJBaI1XvoWRptnnW1LFrlhydxDDuIQjSI,185 +pytz/zoneinfo/Antarctica/McMurdo,sha256=gADjoyPo_QISQU6UJrAgcHp3HDaMoOFRdH-d23uBSyc,2437 +pytz/zoneinfo/Antarctica/Palmer,sha256=HTZY0M8td7oUx5REPgRCHuqKg5V3fjJEi4lYBNL4Etg,1404 +pytz/zoneinfo/Antarctica/Rothera,sha256=_9NY-f8vkozQYrjbUHP5YjcICg0-LuyA9PnIeK123RU,150 +pytz/zoneinfo/Antarctica/South_Pole,sha256=gADjoyPo_QISQU6UJrAgcHp3HDaMoOFRdH-d23uBSyc,2437 +pytz/zoneinfo/Antarctica/Syowa,sha256=oCKH7uafN8R1o-ijXGoT5U1JZxwvoLzJu_2Cqyi2hUM,151 +pytz/zoneinfo/Antarctica/Troll,sha256=fjcYppwr1FnjEssee-RLgGOANzoUyfjse-RGK46PR2E,1148 +pytz/zoneinfo/Antarctica/Vostok,sha256=KfftwdzK6PkMDz0d-D3z4HKIBgY9KqsqHnTnqsPMrUg,213 +pytz/zoneinfo/Arctic/Longyearbyen,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Asia/Aden,sha256=oCKH7uafN8R1o-ijXGoT5U1JZxwvoLzJu_2Cqyi2hUM,151 +pytz/zoneinfo/Asia/Almaty,sha256=lPLWXk2f1mWYRQZFkIrq_5HkhocsUBis0M-yhdDHcBQ,983 +pytz/zoneinfo/Asia/Amman,sha256=Qv4cXXw7KBQWE882cgj0kjQ3wh1vpV1orJ2v2Jjxr2U,1433 +pytz/zoneinfo/Asia/Anadyr,sha256=WqKnHo5IHSWZ08d2sS5ytHtv0MQMoczP3W9zbDDrbYU,1174 +pytz/zoneinfo/Asia/Aqtau,sha256=4n654FZtDssXSfhQszjZG5OmtbE2zo1KbiWcYrFJg00,969 +pytz/zoneinfo/Asia/Aqtobe,sha256=1oFHTb-ybcTqLXm0r1ZOVgdYMTHlGoNs-Pgvux50d3E,997 +pytz/zoneinfo/Asia/Ashgabat,sha256=-sfGnRumio7_Bs8w9YH4xRDWgjB3wBeW7c0C56Qqk64,605 +pytz/zoneinfo/Asia/Ashkhabad,sha256=-sfGnRumio7_Bs8w9YH4xRDWgjB3wBeW7c0C56Qqk64,605 +pytz/zoneinfo/Asia/Atyrau,sha256=_U8COUIE9nG_HKddZE1Q0sPuz3rMwfjwmfnVDY_vSmg,977 +pytz/zoneinfo/Asia/Baghdad,sha256=S-plKI4zCLqI0idGABEk3oRTazNyrIj2T98-EtWtZD8,969 +pytz/zoneinfo/Asia/Bahrain,sha256=wklGY3WPGp-z1OUwb_KOHzRTwBndt1RfDg9Uttt36G4,185 +pytz/zoneinfo/Asia/Baku,sha256=6_hq98SGG0j0JA8qYx96WcIMZSLW4w460QXh_OM_ccg,1213 +pytz/zoneinfo/Asia/Bangkok,sha256=hf_5PVegQcFSS60CjS80C7h-TGOrfQ4ncm83N8VmZkk,185 +pytz/zoneinfo/Asia/Barnaul,sha256=3zeUimLTMrIZE0vX6XHFvB3MoqExoVbE5CSm6GV0zf0,1207 +pytz/zoneinfo/Asia/Beirut,sha256=_Z_2ZAg_iL9vU51JDB8CB04uXBDrf1kLIis-JnXaS2o,2154 +pytz/zoneinfo/Asia/Bishkek,sha256=IOoUyjABILCkXu1rjCIqSwAufRYFklc5YAC4jdhVw6Q,969 +pytz/zoneinfo/Asia/Brunei,sha256=D5qtyWJ_SM8bTQeJJIYhqqojxlVKbrFC1EYMDU9GzXQ,469 +pytz/zoneinfo/Asia/Calcutta,sha256=6Qw0EDbLcgMgDik8s7UTJn4QSjmllPNeGVJU5rwKF88,285 +pytz/zoneinfo/Asia/Chita,sha256=LbSlS23swFkANUScg8zkNR0imANWNfOIaYd39HbLdIQ,1207 +pytz/zoneinfo/Asia/Choibalsan,sha256=qUkXRsTc_u7B90JxULSu7yzKbGtGfKcfEFIasGPC2ec,877 +pytz/zoneinfo/Asia/Chongqing,sha256=ZP_C5DqUQ1oEPAQNHTr36S0DGtx453N68YYbqk7u8-Y,561 +pytz/zoneinfo/Asia/Chungking,sha256=ZP_C5DqUQ1oEPAQNHTr36S0DGtx453N68YYbqk7u8-Y,561 +pytz/zoneinfo/Asia/Colombo,sha256=w52L7bgT4m5hcgRuevIPY83xytfkBmkLhnKMwp16KsY,358 +pytz/zoneinfo/Asia/Dacca,sha256=-xulJ2KVhvKp6rlZLMydpw7oXVirk-riEH-181xPE54,323 +pytz/zoneinfo/Asia/Damascus,sha256=EthGheaHWmy5IrLCc9NmM3jvASQFHt8TsBF07I1tgbg,1873 +pytz/zoneinfo/Asia/Dhaka,sha256=-xulJ2KVhvKp6rlZLMydpw7oXVirk-riEH-181xPE54,323 +pytz/zoneinfo/Asia/Dili,sha256=2A9uFmwSwoFA5o2ek1LA0ucohPnM42ghWvD9D5TdnJk,257 +pytz/zoneinfo/Asia/Dubai,sha256=pmdhPhaJRwKwONvxiZNGeFSICjlWzyY9JlFHv-H9upY,151 +pytz/zoneinfo/Asia/Dushanbe,sha256=koYnnYWuFsBXd1vJfZsGdpwnbFHEwvkGBmSrrx3KIss,577 +pytz/zoneinfo/Asia/Famagusta,sha256=CFrcygd8ude5x6OEtfM_Dw0KYHoxpPPzq46KoHVxjjc,2028 +pytz/zoneinfo/Asia/Gaza,sha256=t0YxcUQL53VNKnKbKijn0OE_MaryEynonabse-iTtzs,3844 +pytz/zoneinfo/Asia/Harbin,sha256=ZP_C5DqUQ1oEPAQNHTr36S0DGtx453N68YYbqk7u8-Y,561 +pytz/zoneinfo/Asia/Hebron,sha256=6Y0USHKx-xoCxCr_WpCuM3olP1vUGnzrcnGiyQFcqdQ,3872 +pytz/zoneinfo/Asia/Ho_Chi_Minh,sha256=Lnv1vpUNAXBo8v0b9d9AQpy-AEyO5Qa2Ig0PvDkjrmU,337 +pytz/zoneinfo/Asia/Hong_Kong,sha256=al_O4kPlq5JpgkLYjEaZzrcgiiLul9NC0R5B69JVWhc,1233 +pytz/zoneinfo/Asia/Hovd,sha256=Zn4PLGlD-URJDsbChor5bqWTzuAil2tbrGJW0j5TLbs,877 +pytz/zoneinfo/Asia/Irkutsk,sha256=IVuoXCwdeI-KIUfFkEt6yBjqYP3V9GTrF-_WLnffFzk,1229 +pytz/zoneinfo/Asia/Istanbul,sha256=Jk4wjndDta_uLWc8W1dWdjbavJJbsL5ROTmZboVnGKU,1933 +pytz/zoneinfo/Asia/Jakarta,sha256=TvEzBvSzfzFCdOsMAZ0QgR95JA5xf3kAZONhy5gEXRE,383 +pytz/zoneinfo/Asia/Jayapura,sha256=ihzUd-L8HUVqG-Na10MyPE-YYwjVFj-xerqjTN4EJZs,221 +pytz/zoneinfo/Asia/Jerusalem,sha256=JUuWQmW5Tha0pJjw61Q5aN7CX0z4D7ops9OOSnda6Dc,2388 +pytz/zoneinfo/Asia/Kabul,sha256=JZEbo8bSj_L7HnXUm2gAUlNlCvJlRJhFkSHCg5o3ggk,194 +pytz/zoneinfo/Asia/Kamchatka,sha256=KY1PlJvRSNkY_5hyJBxj5DDweeYVQaBK05ZgL3kdcCY,1152 +pytz/zoneinfo/Asia/Karachi,sha256=iB-mWMTXUyfBwAkZdz8_UmEw0xsgxIub-KNI7akzhkk,379 +pytz/zoneinfo/Asia/Kashgar,sha256=F1ZOdZZDsVHwDJinksR-hjcqPzqOljvdreZIWFulJxY,151 +pytz/zoneinfo/Asia/Kathmandu,sha256=_RsfeSWbCr8kM4YRJi7Xv6hAEiHW14IFhsXsfhbPjoM,198 +pytz/zoneinfo/Asia/Katmandu,sha256=_RsfeSWbCr8kM4YRJi7Xv6hAEiHW14IFhsXsfhbPjoM,198 +pytz/zoneinfo/Asia/Khandyga,sha256=bKfmw6k5qYDQsEHG3Mv-VYis3YhCeV7qijDxfxQNn_g,1257 +pytz/zoneinfo/Asia/Kolkata,sha256=6Qw0EDbLcgMgDik8s7UTJn4QSjmllPNeGVJU5rwKF88,285 +pytz/zoneinfo/Asia/Krasnoyarsk,sha256=D5KE_1wWSD2YdixDy8n3LBNaAlE1_y3TWXw6NrxFKKA,1193 +pytz/zoneinfo/Asia/Kuala_Lumpur,sha256=XmeVImeqcJ8hJzm7TjAti1nWJAxawOqq7jIzDnHX2hI,401 +pytz/zoneinfo/Asia/Kuching,sha256=D5qtyWJ_SM8bTQeJJIYhqqojxlVKbrFC1EYMDU9GzXQ,469 +pytz/zoneinfo/Asia/Kuwait,sha256=oCKH7uafN8R1o-ijXGoT5U1JZxwvoLzJu_2Cqyi2hUM,151 +pytz/zoneinfo/Asia/Macao,sha256=MvAkRyRsrA2r052ItlyF5bh2FheRjI0jPwg0uIiH2Yk,1227 +pytz/zoneinfo/Asia/Macau,sha256=MvAkRyRsrA2r052ItlyF5bh2FheRjI0jPwg0uIiH2Yk,1227 +pytz/zoneinfo/Asia/Magadan,sha256=HccEEXBQvMmLoC_JE-zP_MlLAZ1WmNLQLfM3tJt55M4,1208 +pytz/zoneinfo/Asia/Makassar,sha256=OhJtCqSTEU-u5n0opBVO5Bu-wQzcYPy9S_6aAhJXgOw,254 +pytz/zoneinfo/Asia/Manila,sha256=8xTSHFQuYVdW3ThdNqiWzVe6Fv75g_5rTQYURLvxrJ4,422 +pytz/zoneinfo/Asia/Muscat,sha256=pmdhPhaJRwKwONvxiZNGeFSICjlWzyY9JlFHv-H9upY,151 +pytz/zoneinfo/Asia/Nicosia,sha256=0Unm0IFT7HyGeQ7F3vTa_-klfysCgrulqFO6BD1plZU,2002 +pytz/zoneinfo/Asia/Novokuznetsk,sha256=pyxxtSUtYDeVmFk0Cg-F33laZS0iKtde9_GJnL9f0KM,1151 +pytz/zoneinfo/Asia/Novosibirsk,sha256=5K2-Gx15ThlHfolyW85S5zREtAcMjeHBYWK4E8x2LdY,1207 +pytz/zoneinfo/Asia/Omsk,sha256=HyXIWItJXBKVHUzWcQPi1Mmd6ZLmZk-QhRUo9Kv2XOI,1193 +pytz/zoneinfo/Asia/Oral,sha256=WQT4qRmC9RI_ll8zB9FvkAL8ezGb8qoqWd75GTlC7kQ,991 +pytz/zoneinfo/Asia/Phnom_Penh,sha256=hf_5PVegQcFSS60CjS80C7h-TGOrfQ4ncm83N8VmZkk,185 +pytz/zoneinfo/Asia/Pontianak,sha256=inOXwuKtjKv1z_eliPZSIqjSt6whtuxhPeG1YpjU_BQ,353 +pytz/zoneinfo/Asia/Pyongyang,sha256=_-g3GnDAtfDX4XAktXH9jFouLUDmOovnjoOfvRpUDsE,237 +pytz/zoneinfo/Asia/Qatar,sha256=wklGY3WPGp-z1OUwb_KOHzRTwBndt1RfDg9Uttt36G4,185 +pytz/zoneinfo/Asia/Qostanay,sha256=HIjln8QIPNRU6MkWzyPi6vDrjlmVZ4XzFxcUHtXMi7s,1025 +pytz/zoneinfo/Asia/Qyzylorda,sha256=JZLNN6NuLkqaWEeVaCZiW_gL6BrIFL9lr65iK7myVPg,1011 +pytz/zoneinfo/Asia/Rangoon,sha256=_YHASq4Z5YcUILIdhEzg27CGLzarUHPDHs1Dj0QgNGM,254 +pytz/zoneinfo/Asia/Riyadh,sha256=oCKH7uafN8R1o-ijXGoT5U1JZxwvoLzJu_2Cqyi2hUM,151 +pytz/zoneinfo/Asia/Saigon,sha256=Lnv1vpUNAXBo8v0b9d9AQpy-AEyO5Qa2Ig0PvDkjrmU,337 +pytz/zoneinfo/Asia/Sakhalin,sha256=xzAor82ihAe-yXEwC6OWiMzo9b6Z-oQl39NIkU5Hhbs,1188 +pytz/zoneinfo/Asia/Samarkand,sha256=zJKSRt3lEvd6Qvg9b49QAyO4cTJyVnTKyPYcyudpHxk,563 +pytz/zoneinfo/Asia/Seoul,sha256=LI9LsV3XcJC0l-KoQf8zI-y7rk-du57erS-N2Ptdi7Q,617 +pytz/zoneinfo/Asia/Shanghai,sha256=ZP_C5DqUQ1oEPAQNHTr36S0DGtx453N68YYbqk7u8-Y,561 +pytz/zoneinfo/Asia/Singapore,sha256=XmeVImeqcJ8hJzm7TjAti1nWJAxawOqq7jIzDnHX2hI,401 +pytz/zoneinfo/Asia/Srednekolymsk,sha256=efaaT8iFHrcccp-VZKNMvtTuPLNjG5V9JH5KKHhH3SI,1194 +pytz/zoneinfo/Asia/Taipei,sha256=DMmQwOpPql25ue3Nf8vAKKT4em06D1Z9rHbLIitxixk,761 +pytz/zoneinfo/Asia/Tashkent,sha256=apRPy251fSRy_ixsg3BOZNmUbHdO86P5-PdgC1Xws7U,577 +pytz/zoneinfo/Asia/Tbilisi,sha256=zQ-2bVq5_USUSbwN6q0qvWjD-HXkKaH4ifMVq1lEeIM,1021 +pytz/zoneinfo/Asia/Tehran,sha256=Lb2H9BCBXtz819FL6E3gBA7w2ROiIgPgx-f08XpqkVo,1248 +pytz/zoneinfo/Asia/Tel_Aviv,sha256=JUuWQmW5Tha0pJjw61Q5aN7CX0z4D7ops9OOSnda6Dc,2388 +pytz/zoneinfo/Asia/Thimbu,sha256=G2nTQVEMmKlWt0B74_fUAL7KQ3YAu__J6HciiYs2IyU,189 +pytz/zoneinfo/Asia/Thimphu,sha256=G2nTQVEMmKlWt0B74_fUAL7KQ3YAu__J6HciiYs2IyU,189 +pytz/zoneinfo/Asia/Tokyo,sha256=oCueZgRNxcNcX3ZGdif9y6Su4cyVhga4XHdwlcrYLOs,309 +pytz/zoneinfo/Asia/Tomsk,sha256=cr0ULZgWBnQfzDiJeYmqpA7Xo5QRzurvrHsrbZsnhOQ,1207 +pytz/zoneinfo/Asia/Ujung_Pandang,sha256=OhJtCqSTEU-u5n0opBVO5Bu-wQzcYPy9S_6aAhJXgOw,254 +pytz/zoneinfo/Asia/Ulaanbaatar,sha256=qUkXRsTc_u7B90JxULSu7yzKbGtGfKcfEFIasGPC2ec,877 +pytz/zoneinfo/Asia/Ulan_Bator,sha256=qUkXRsTc_u7B90JxULSu7yzKbGtGfKcfEFIasGPC2ec,877 +pytz/zoneinfo/Asia/Urumqi,sha256=F1ZOdZZDsVHwDJinksR-hjcqPzqOljvdreZIWFulJxY,151 +pytz/zoneinfo/Asia/Ust-Nera,sha256=zsG8kgnw0Fcs5N2WwNTVmvWkTlpwf7Oo8y68HcXjYyw,1238 +pytz/zoneinfo/Asia/Vientiane,sha256=hf_5PVegQcFSS60CjS80C7h-TGOrfQ4ncm83N8VmZkk,185 +pytz/zoneinfo/Asia/Vladivostok,sha256=XMQLMh5SPbI6C4R3UO4KhbnG4hWVkHNedzCQeqxFk6A,1194 +pytz/zoneinfo/Asia/Yakutsk,sha256=PPNrRGgg9jefOUM-6M8XqaIm-ElfmRZSWAtSGKLzNXQ,1193 +pytz/zoneinfo/Asia/Yangon,sha256=_YHASq4Z5YcUILIdhEzg27CGLzarUHPDHs1Dj0QgNGM,254 +pytz/zoneinfo/Asia/Yekaterinburg,sha256=4NyEW6Xjr4UsWPh63HIPI4G6GT_tVG1Xkgc2xbwGjzA,1229 +pytz/zoneinfo/Asia/Yerevan,sha256=FM0pUA4NbTWBb_CsJ5KCLVrLoNmad7njBKqFrJBDoxE,1137 +pytz/zoneinfo/Atlantic/Azores,sha256=66hDxaK8xFnktLMrpNxkD4r1gGkhS-PEpleuwzuGRA0,3442 +pytz/zoneinfo/Atlantic/Bermuda,sha256=LNGKfMsnYvwImjTyzXrLhMOHHDu7qI67RbYNKvvI15I,2396 +pytz/zoneinfo/Atlantic/Canary,sha256=ymK9ufqphvNjDK3hzikN4GfkcR3QeCBiPKyVc6FjlbA,1897 +pytz/zoneinfo/Atlantic/Cape_Verde,sha256=o92pLdLFX_b9vUiq3rNpca4tupIO3dx9rNrnPcA8474,256 +pytz/zoneinfo/Atlantic/Faeroe,sha256=NibdZPZtapnYR_myIZnMdTaSKGsOBGgujj0_T2NvAzs,1815 +pytz/zoneinfo/Atlantic/Faroe,sha256=NibdZPZtapnYR_myIZnMdTaSKGsOBGgujj0_T2NvAzs,1815 +pytz/zoneinfo/Atlantic/Jan_Mayen,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Atlantic/Madeira,sha256=lYY85MC5-GUKExm353ixwtZDxasYavTTWELvv5RXLxE,3377 +pytz/zoneinfo/Atlantic/Reykjavik,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Atlantic/South_Georgia,sha256=I9SAcPPumy6Xf9P7dg2aE16oxwDIqyKFqinJTC-XsgM,150 +pytz/zoneinfo/Atlantic/St_Helena,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Atlantic/Stanley,sha256=siEjXTAuTum_4XGtS98MBE34XW_5xgXShEX5OMnSFjo,1200 +pytz/zoneinfo/Australia/ACT,sha256=QsOFdYWxbbL4_9R7oZ-qYPRzNA3o1P6TIOp76GFgWQY,2190 +pytz/zoneinfo/Australia/Adelaide,sha256=ld2EbxU75oVgmPe703z-I6aqLg0Kmv62ZcCGzkT5R20,2208 +pytz/zoneinfo/Australia/Brisbane,sha256=eW6Qzze2t0-speJmmvt1JMzbkSadIKdE84XHc7JUtGc,419 +pytz/zoneinfo/Australia/Broken_Hill,sha256=3k_3ljTvS5GSfo7Xh6w71UgR3aAwYPBsnCJ-mlEYCqQ,2229 +pytz/zoneinfo/Australia/Canberra,sha256=QsOFdYWxbbL4_9R7oZ-qYPRzNA3o1P6TIOp76GFgWQY,2190 +pytz/zoneinfo/Australia/Currie,sha256=GLQSzgIfsWxOvmKOrhpfofWqINQf6h36NYy3mcq6gcg,2358 +pytz/zoneinfo/Australia/Darwin,sha256=fn0IZhIW98FAnzLig-_GBtW5LA54jajdeeUzg4tCGvo,325 +pytz/zoneinfo/Australia/Eucla,sha256=i1-XGG8I6E0dXIdWGF4DlkfDLWhiAxJ_3gMpt-nm_u4,456 +pytz/zoneinfo/Australia/Hobart,sha256=GLQSzgIfsWxOvmKOrhpfofWqINQf6h36NYy3mcq6gcg,2358 +pytz/zoneinfo/Australia/LHI,sha256=oyPFQzmRqWPrSXt9pNHQmEi_PvX11k2clknziOS6ud8,1846 +pytz/zoneinfo/Australia/Lindeman,sha256=xM6Udx22oLNoLR1Y7GQhHOYov8nw3xQNqgc_NVQ2JK4,475 +pytz/zoneinfo/Australia/Lord_Howe,sha256=oyPFQzmRqWPrSXt9pNHQmEi_PvX11k2clknziOS6ud8,1846 +pytz/zoneinfo/Australia/Melbourne,sha256=lvx_MQcunMc6u2smIrl8X427bLsXvjkgpCSdjYCTNBM,2190 +pytz/zoneinfo/Australia/NSW,sha256=QsOFdYWxbbL4_9R7oZ-qYPRzNA3o1P6TIOp76GFgWQY,2190 +pytz/zoneinfo/Australia/North,sha256=fn0IZhIW98FAnzLig-_GBtW5LA54jajdeeUzg4tCGvo,325 +pytz/zoneinfo/Australia/Perth,sha256=Al1DOUh4U_ofMUQSeVlzSyD3x7SUjP9dchSaBUGmeWg,446 +pytz/zoneinfo/Australia/Queensland,sha256=eW6Qzze2t0-speJmmvt1JMzbkSadIKdE84XHc7JUtGc,419 +pytz/zoneinfo/Australia/South,sha256=ld2EbxU75oVgmPe703z-I6aqLg0Kmv62ZcCGzkT5R20,2208 +pytz/zoneinfo/Australia/Sydney,sha256=QsOFdYWxbbL4_9R7oZ-qYPRzNA3o1P6TIOp76GFgWQY,2190 +pytz/zoneinfo/Australia/Tasmania,sha256=GLQSzgIfsWxOvmKOrhpfofWqINQf6h36NYy3mcq6gcg,2358 +pytz/zoneinfo/Australia/Victoria,sha256=lvx_MQcunMc6u2smIrl8X427bLsXvjkgpCSdjYCTNBM,2190 +pytz/zoneinfo/Australia/West,sha256=Al1DOUh4U_ofMUQSeVlzSyD3x7SUjP9dchSaBUGmeWg,446 +pytz/zoneinfo/Australia/Yancowinna,sha256=3k_3ljTvS5GSfo7Xh6w71UgR3aAwYPBsnCJ-mlEYCqQ,2229 +pytz/zoneinfo/Brazil/Acre,sha256=0gpJUl46hQbp0P6Xj1S0NArIWeAryuuDXjsldvB5GHE,614 +pytz/zoneinfo/Brazil/DeNoronha,sha256=feeRAijQqKylZgqe84nKhsFLycT5zIBm7mLIvdyGw4w,702 +pytz/zoneinfo/Brazil/East,sha256=BMBnRO4_4HjvO4t3njjrMGZr-ZPmegkvyvL8KPY6ZM4,1430 +pytz/zoneinfo/Brazil/West,sha256=F6RLOOeOi9lymZiQmQ9pR8tFpPZ6EguNdPfOc6BhXDE,590 +pytz/zoneinfo/CET,sha256=gS9Vrrbozend9HhuFetCVrIegs9fXSjaG60X2UVwysA,2933 +pytz/zoneinfo/CST6CDT,sha256=_roybr6I6sIAF6cYdIxGxoRpoef153Fty48dQ6bm9oY,3592 +pytz/zoneinfo/Canada/Atlantic,sha256=TZpmc5PwWoLfTfQoQ_b3U17BE2iVKSeNkR0Ho8mbTn8,3424 +pytz/zoneinfo/Canada/Central,sha256=7P-_YQrneFcon7QKSTOnkiGjEppFDn3Z48MJ1qq8VBw,2868 +pytz/zoneinfo/Canada/Eastern,sha256=pYehoWB0Ofe6woPhgV8r26-5ZJpFPRjgbC5E5pltiI8,3494 +pytz/zoneinfo/Canada/Mountain,sha256=-TkIfc3QlvaCf0p8COZ43Y1HRBAl-nARUi-JdXeK1vE,2332 +pytz/zoneinfo/Canada/Newfoundland,sha256=r1-17uKv27eZ3JsVkw_DLZQbo6wvjuuVu7C2pDsmOgI,3655 +pytz/zoneinfo/Canada/Pacific,sha256=sknKH0jSPWam-DHfM35qXs8Nam7d5TFlkUI9Sgxryyg,2892 +pytz/zoneinfo/Canada/Saskatchewan,sha256=yjqT08pHbICYe83H8JmtaDBvCFqRv7Tfze3Y8xuXukw,980 +pytz/zoneinfo/Canada/Yukon,sha256=TrR6PCnYG-mSClBMohqlP8qnYhXMUsydI-L-quXFxyM,1614 +pytz/zoneinfo/Chile/Continental,sha256=0CDw13dCMUsoquMupoJgupkzAUNhDK6E0lVxURA7osA,2515 +pytz/zoneinfo/Chile/EasterIsland,sha256=QbubBs_xQlvKweAnurhyHjIK4ji77Gh4G-usXul6XVM,2219 +pytz/zoneinfo/Cuba,sha256=HUQeAuKBsEkI5SLZjqynXICOUVOajkKzKH5r-Ov5Odc,2416 +pytz/zoneinfo/EET,sha256=XDY-FBUddRyQHN8GxQLZ4awjuOlWlzlUdjv7OdXFNzA,2262 +pytz/zoneinfo/EST,sha256=kayA_pdpMcSQ0FjIzotdcf-m1JYfbKE-qcFT8LC8zqA,182 +pytz/zoneinfo/EST5EDT,sha256=6e0H177gx2qdRC0JHvHwFmj-58TyYBTAqGixn-bBipU,3552 +pytz/zoneinfo/Egypt,sha256=Lft-GCLQhaSJm9VqUmsEFoHIS1Vhfa7pFJn9GZCpifs,2399 +pytz/zoneinfo/Eire,sha256=QOjSocO1cihNo59vQkWxvIFPRSxE9apz0KARVx1czEM,3492 +pytz/zoneinfo/Etc/GMT,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Etc/GMT+0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Etc/GMT+1,sha256=1Qzl2X9rQ_RXEf11yH09wQZCr_ph6UdFP7E0yu9s-IQ,116 +pytz/zoneinfo/Etc/GMT+10,sha256=JEQyQyQlkC0o6ZTdeVjZhCIOh6cK5TF7H00Pkls-sUI,117 +pytz/zoneinfo/Etc/GMT+11,sha256=tWvcvYMFCaE60nJVvDrrov7stJvs1KQYOyrhl3dzcUs,117 +pytz/zoneinfo/Etc/GMT+12,sha256=b70HEhErq8IJmq8x7cOZy4eR__3fq5uHHpjvPBEHqMA,117 +pytz/zoneinfo/Etc/GMT+2,sha256=T6Ep5zhslBKbYaECFUB6gUKh3iTZPyMoW1kjhonxrUo,116 +pytz/zoneinfo/Etc/GMT+3,sha256=QGoYrE04bUJ-OzL37dt2MZT5FxWNLpJDPVXgJbstYZA,116 +pytz/zoneinfo/Etc/GMT+4,sha256=RWrkNki-wV7X-coe0VvufBe6LrWVpkPJgia5QQYEnBo,116 +pytz/zoneinfo/Etc/GMT+5,sha256=oRmeC41dgYXT-zzyZIRKXN9IvdL2Da5nTuwmG2_prIA,116 +pytz/zoneinfo/Etc/GMT+6,sha256=d6dAnwiejyFI2n7AzFlFW0aFAT6zYNEjBIEG0uu0sbQ,116 +pytz/zoneinfo/Etc/GMT+7,sha256=TqjYbzd0YHpx1wisFg08J19wTpg6ztJLLongZY_lozs,116 +pytz/zoneinfo/Etc/GMT+8,sha256=th_8bIMmYgRPCesBrbmBhRr0jQO7whd70LiY9HfwJyk,116 +pytz/zoneinfo/Etc/GMT+9,sha256=Qq5E6iUS7JMJIymT7YoqlI8MtqtVy0mr9t6zWFtWc9Y,116 +pytz/zoneinfo/Etc/GMT-0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Etc/GMT-1,sha256=73F1eU8uAQGP3mcoB2q99CjfManGFHk3fefljp9pYC4,117 +pytz/zoneinfo/Etc/GMT-10,sha256=fKWWNwLBOp1OkKjtc1w9LIXJR1mTTD-JdvYflRy1IrU,118 +pytz/zoneinfo/Etc/GMT-11,sha256=D2S79n6psa9t9_2vj5wIrFpHH2OJLcCKP6vtwzFZINY,118 +pytz/zoneinfo/Etc/GMT-12,sha256=me4V6lmWI8gSr8H7N41WAD0Eww1anh_EF34Qr9UoSnI,118 +pytz/zoneinfo/Etc/GMT-13,sha256=xbmbG1BQA6Dlpa_iUwEGyJxW4a3t6lmawdPKAE8vbR8,118 +pytz/zoneinfo/Etc/GMT-14,sha256=PpXoREBh02qFpvxVMj2pV9IAzSQvBE7XPvnN9qSZ-Kc,118 +pytz/zoneinfo/Etc/GMT-2,sha256=ve6hWLdeuiLhqagaWLqMD6HNybS1chRwjudfTZ2bYBE,117 +pytz/zoneinfo/Etc/GMT-3,sha256=N77jILanuLDVkLsdujXZSu-dsHiwN5MIpwh7fMUifso,117 +pytz/zoneinfo/Etc/GMT-4,sha256=LSko5fVHqPl5zfwjGqkbMa_OFnvtpT6o_4xYxNz9n5o,117 +pytz/zoneinfo/Etc/GMT-5,sha256=uLaSR5Mb18HRTsAA5SveY9PAJ97dO8QzIWqNXe3wZb4,117 +pytz/zoneinfo/Etc/GMT-6,sha256=JSN-RUAphJ50fpIv7cYC6unrtrz9S1Wma-piDHlGe7c,117 +pytz/zoneinfo/Etc/GMT-7,sha256=vVAOF8xU9T9ESnw68c0SFXpcvkoopaiwTR0zbefHHSU,117 +pytz/zoneinfo/Etc/GMT-8,sha256=S7xFQbFMpiDZy4v5L4D9fCrjRIzzoLC5p8Se23xi7us,117 +pytz/zoneinfo/Etc/GMT-9,sha256=I5vHNmUK-Yyg_S1skFN44VGVzBgktjFgVQiDIKO4aMI,117 +pytz/zoneinfo/Etc/GMT0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Etc/Greenwich,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Etc/UCT,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/Etc/UTC,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/Etc/Universal,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/Etc/Zulu,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/Europe/Amsterdam,sha256=gS9Vrrbozend9HhuFetCVrIegs9fXSjaG60X2UVwysA,2933 +pytz/zoneinfo/Europe/Andorra,sha256=gTB5jCQmvIw3JJi1_vAcOYuhtzPBR6RXUx9gVV6p6ug,1742 +pytz/zoneinfo/Europe/Astrakhan,sha256=ZeGDZjwVVRoeR-J642zEnN26BPL58ViTJLbwnk7pLXk,1151 +pytz/zoneinfo/Europe/Athens,sha256=XDY-FBUddRyQHN8GxQLZ4awjuOlWlzlUdjv7OdXFNzA,2262 +pytz/zoneinfo/Europe/Belfast,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/Europe/Belgrade,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/Berlin,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Europe/Bratislava,sha256=G9fdhUXmzx651BnyZ6V7AOYIV9EV5aMJMm44eJaLLZw,2301 +pytz/zoneinfo/Europe/Brussels,sha256=gS9Vrrbozend9HhuFetCVrIegs9fXSjaG60X2UVwysA,2933 +pytz/zoneinfo/Europe/Bucharest,sha256=nfg6-bU2D6DMEWb9EMIBR5kxnNsbDSx0UKfHH_ZzqFc,2184 +pytz/zoneinfo/Europe/Budapest,sha256=lNwqxWciBvw9ei81VQwIKHbC_ZDJjpgHU6HFg4wCUkY,2368 +pytz/zoneinfo/Europe/Busingen,sha256=K5QY7Ujj2VUchKR4bhhb0hgdAJhmwED71ykXDQOGKe8,1909 +pytz/zoneinfo/Europe/Chisinau,sha256=p1J_rqFE13pL8cpBRrEFe-teCI8f0fKK4uTUy_4diF4,2390 +pytz/zoneinfo/Europe/Copenhagen,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Europe/Dublin,sha256=QOjSocO1cihNo59vQkWxvIFPRSxE9apz0KARVx1czEM,3492 +pytz/zoneinfo/Europe/Gibraltar,sha256=a87WpaBlvxI4gAU9OpQOkN8VUJbirVWYf-VfFLTIoS4,3068 +pytz/zoneinfo/Europe/Guernsey,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/Europe/Helsinki,sha256=GEkB7LsVhmegt7YuuWheCDvDGC7b7Nw9bTdDGS9qkJc,1900 +pytz/zoneinfo/Europe/Isle_of_Man,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/Europe/Istanbul,sha256=Jk4wjndDta_uLWc8W1dWdjbavJJbsL5ROTmZboVnGKU,1933 +pytz/zoneinfo/Europe/Jersey,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/Europe/Kaliningrad,sha256=s7GXSe1YvMcs7AiUhHNTA6I4nAOQn_Kmz_ZqJYO-LMM,1493 +pytz/zoneinfo/Europe/Kiev,sha256=-wrpG9jPuIKFP1NgBVvnxsMRf9L_h5z3J6Q3jj1AwNM,2120 +pytz/zoneinfo/Europe/Kirov,sha256=P7T2Zf5Eo6o4L4Dbg_BfiFjUgTj0dQXlrwY-QZ1eBVk,1185 +pytz/zoneinfo/Europe/Kyiv,sha256=-wrpG9jPuIKFP1NgBVvnxsMRf9L_h5z3J6Q3jj1AwNM,2120 +pytz/zoneinfo/Europe/Lisbon,sha256=krB8skaJImv5NDCNHxvTPDBqpNphDFLNW84lB3lgUCw,3527 +pytz/zoneinfo/Europe/Ljubljana,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/London,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/Europe/Luxembourg,sha256=gS9Vrrbozend9HhuFetCVrIegs9fXSjaG60X2UVwysA,2933 +pytz/zoneinfo/Europe/Madrid,sha256=mkLX03rW3t0tmzKBIPe_noUvaFDErwC6_5ZPZZsWHOo,2614 +pytz/zoneinfo/Europe/Malta,sha256=EhKcbPL47765tWAiQ57cusaK2TaIQqZCgtJoEZs3Ud0,2620 +pytz/zoneinfo/Europe/Mariehamn,sha256=GEkB7LsVhmegt7YuuWheCDvDGC7b7Nw9bTdDGS9qkJc,1900 +pytz/zoneinfo/Europe/Minsk,sha256=KgPm0fHycntgd3xbTmmDl4O13Xh_9e2zUnd8XFSU29o,1307 +pytz/zoneinfo/Europe/Monaco,sha256=q3ehSIot1GZ6TyMHIjbg0oRf4ghAXuwbSDSYVim6evg,2962 +pytz/zoneinfo/Europe/Moscow,sha256=KmkofRcj6T8Ph28PJChm8JVp13uRvef6TZ0GuPzUiDw,1535 +pytz/zoneinfo/Europe/Nicosia,sha256=0Unm0IFT7HyGeQ7F3vTa_-klfysCgrulqFO6BD1plZU,2002 +pytz/zoneinfo/Europe/Oslo,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Europe/Paris,sha256=q3ehSIot1GZ6TyMHIjbg0oRf4ghAXuwbSDSYVim6evg,2962 +pytz/zoneinfo/Europe/Podgorica,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/Prague,sha256=G9fdhUXmzx651BnyZ6V7AOYIV9EV5aMJMm44eJaLLZw,2301 +pytz/zoneinfo/Europe/Riga,sha256=hJ2_0m1taW9IuA-hMyP5n-WX7YOrR0heKszJhgljRWk,2198 +pytz/zoneinfo/Europe/Rome,sha256=1a3oLMSiMpSbh9QxV8hLLDVbZqash89iUO1urYC1AY8,2641 +pytz/zoneinfo/Europe/Samara,sha256=nXL0IxbT6qu10CNuaDHxx4W1OaAnaaKTtIJ9N9URMoU,1201 +pytz/zoneinfo/Europe/San_Marino,sha256=1a3oLMSiMpSbh9QxV8hLLDVbZqash89iUO1urYC1AY8,2641 +pytz/zoneinfo/Europe/Sarajevo,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/Saratov,sha256=ygwjvXN13TgaWxjg6ysWEnHWNxwrVtkEbrk8t9bzVVw,1169 +pytz/zoneinfo/Europe/Simferopol,sha256=tzl7xdNVSZprNCul4YE5LSpoR9JoujmOq8VbbB8wHic,1469 +pytz/zoneinfo/Europe/Skopje,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/Sofia,sha256=hCQKXfMNrnA5xHNw_uzTjKzVw4-Bvsq5oGO4yUCv5tY,2077 +pytz/zoneinfo/Europe/Stockholm,sha256=XuR19xoPwaMvrrhJ-MOcbnqmbW1B7HQrl7OnQ2s7BwE,2298 +pytz/zoneinfo/Europe/Tallinn,sha256=4a6JC0aIpMzqIV7O35zoG0LLJwkQq5AoXZ2ivkic6-w,2148 +pytz/zoneinfo/Europe/Tirane,sha256=ztlZyCS9WCXeVW8nBun3Tyi5HUY0EtFbiBbEc1gucuw,2084 +pytz/zoneinfo/Europe/Tiraspol,sha256=p1J_rqFE13pL8cpBRrEFe-teCI8f0fKK4uTUy_4diF4,2390 +pytz/zoneinfo/Europe/Ulyanovsk,sha256=c8Ad5p7CKj_1cCA7lVRpcPqbQXGYaX83cuu6uIFx-Bg,1253 +pytz/zoneinfo/Europe/Uzhgorod,sha256=-wrpG9jPuIKFP1NgBVvnxsMRf9L_h5z3J6Q3jj1AwNM,2120 +pytz/zoneinfo/Europe/Vaduz,sha256=K5QY7Ujj2VUchKR4bhhb0hgdAJhmwED71ykXDQOGKe8,1909 +pytz/zoneinfo/Europe/Vatican,sha256=1a3oLMSiMpSbh9QxV8hLLDVbZqash89iUO1urYC1AY8,2641 +pytz/zoneinfo/Europe/Vienna,sha256=ZmI3kADE6bnrJEccqh73XXBY36L1G4DkpiTQImtNrUk,2200 +pytz/zoneinfo/Europe/Vilnius,sha256=UFzRX3orCTB8d9IzlxJPy5eUA2oBPuCu1UJl-2D7C3U,2162 +pytz/zoneinfo/Europe/Volgograd,sha256=RgFvt7mzZ-TtIKL9BVHmoNZLIeLIuiDdXeY10g2_vks,1193 +pytz/zoneinfo/Europe/Warsaw,sha256=TiLDPbeVF0ckgLVEkaSeDaKZ8wctdJDOl_HE_Wd5rKs,2654 +pytz/zoneinfo/Europe/Zagreb,sha256=OpWtsGFWBE_S-mYoQcAmjCta9HwbGQANnSmVY9OHCTo,1920 +pytz/zoneinfo/Europe/Zaporozhye,sha256=-wrpG9jPuIKFP1NgBVvnxsMRf9L_h5z3J6Q3jj1AwNM,2120 +pytz/zoneinfo/Europe/Zurich,sha256=K5QY7Ujj2VUchKR4bhhb0hgdAJhmwED71ykXDQOGKe8,1909 +pytz/zoneinfo/Factory,sha256=aFFlKx93HXoJoF4SSuTlD8cZtJA-ne5oKzAa6eX2V4k,116 +pytz/zoneinfo/GB,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/GB-Eire,sha256=yFSVBw3KQmh99qHD7ngKJ8vLgvGER1Dqb2QoM6RNKbQ,3664 +pytz/zoneinfo/GMT,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/GMT+0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/GMT-0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/GMT0,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/Greenwich,sha256=bZ83iIPAefhsA4elVHqSxEmGnYBuB94QCEqwTwJJAY0,114 +pytz/zoneinfo/HST,sha256=fwPRv1Jk56sCOi75uZfd_Iy2k2aSQHx3B2K5xUlSPzM,329 +pytz/zoneinfo/Hongkong,sha256=al_O4kPlq5JpgkLYjEaZzrcgiiLul9NC0R5B69JVWhc,1233 +pytz/zoneinfo/Iceland,sha256=0u-sTl8j2IyV1ywdtCgHFw9S9D3ZiiBa9akqkbny2Zc,148 +pytz/zoneinfo/Indian/Antananarivo,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Indian/Chagos,sha256=2errXzKdFIcpU0L-XRhSHxhNabIzbI5lXV3Pq6lt40Y,185 +pytz/zoneinfo/Indian/Christmas,sha256=hf_5PVegQcFSS60CjS80C7h-TGOrfQ4ncm83N8VmZkk,185 +pytz/zoneinfo/Indian/Cocos,sha256=_YHASq4Z5YcUILIdhEzg27CGLzarUHPDHs1Dj0QgNGM,254 +pytz/zoneinfo/Indian/Comoro,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Indian/Kerguelen,sha256=F73ffVfBoUoHre0-DwsiQrYJcLpPOW-JJGk3n88lM5U,185 +pytz/zoneinfo/Indian/Mahe,sha256=pmdhPhaJRwKwONvxiZNGeFSICjlWzyY9JlFHv-H9upY,151 +pytz/zoneinfo/Indian/Maldives,sha256=F73ffVfBoUoHre0-DwsiQrYJcLpPOW-JJGk3n88lM5U,185 +pytz/zoneinfo/Indian/Mauritius,sha256=Znqrc1chimlciJsYBOl0NvIHnrNdCxncGxWczq1PBeI,227 +pytz/zoneinfo/Indian/Mayotte,sha256=yJsuJTqJJqbOz37_NOS_zbf-JNr_IthHGMMN7sDqSWg,265 +pytz/zoneinfo/Indian/Reunion,sha256=pmdhPhaJRwKwONvxiZNGeFSICjlWzyY9JlFHv-H9upY,151 +pytz/zoneinfo/Iran,sha256=Lb2H9BCBXtz819FL6E3gBA7w2ROiIgPgx-f08XpqkVo,1248 +pytz/zoneinfo/Israel,sha256=JUuWQmW5Tha0pJjw61Q5aN7CX0z4D7ops9OOSnda6Dc,2388 +pytz/zoneinfo/Jamaica,sha256=wlagieUPRf5-beie-h7QsONbNzjGsm8vMs8uf28pw28,482 +pytz/zoneinfo/Japan,sha256=oCueZgRNxcNcX3ZGdif9y6Su4cyVhga4XHdwlcrYLOs,309 +pytz/zoneinfo/Kwajalein,sha256=TmZ_0f-ySQ-saBAlRXV0f49Itwne51VBXn6rWcrWqHQ,302 +pytz/zoneinfo/Libya,sha256=W1dptGD70T7ppGoo0fczFQeDiIp0nultLNPV66MwB2c,625 +pytz/zoneinfo/MET,sha256=gS9Vrrbozend9HhuFetCVrIegs9fXSjaG60X2UVwysA,2933 +pytz/zoneinfo/MST,sha256=illz0sYuLL8lIPK0Tkou6dL0Vck_D0W_3rRTOvFYRmQ,360 +pytz/zoneinfo/MST7MDT,sha256=MugZwApDs8NI9TnXANQlUE8guNBowWQY0m-ptpPndck,2460 +pytz/zoneinfo/Mexico/BajaNorte,sha256=SluV7xzZm24LgMXSUVt1cD1AlE7y_bdE65HhDIdXLcs,2458 +pytz/zoneinfo/Mexico/BajaSur,sha256=BWH2NqVPA1PsyELPN_2BF8KllrsmQkqg1eujsQvnnx8,1060 +pytz/zoneinfo/Mexico/General,sha256=Uog2-FMWz2o12jR6sK9vemJamLeo6OEFMQR3s0xTxkc,1222 +pytz/zoneinfo/NZ,sha256=gADjoyPo_QISQU6UJrAgcHp3HDaMoOFRdH-d23uBSyc,2437 +pytz/zoneinfo/NZ-CHAT,sha256=xhexVc5lfJ_qAv2d3HrII6lfRSxKZYBAjY2zpYkCGE8,2054 +pytz/zoneinfo/Navajo,sha256=MugZwApDs8NI9TnXANQlUE8guNBowWQY0m-ptpPndck,2460 +pytz/zoneinfo/PRC,sha256=ZP_C5DqUQ1oEPAQNHTr36S0DGtx453N68YYbqk7u8-Y,561 +pytz/zoneinfo/PST8PDT,sha256=aJd7ua1tGG_vxser02AQpm4wAI3LLTdgh6QcSYYecmg,2852 +pytz/zoneinfo/Pacific/Apia,sha256=M3QKsp75Q7H1X3aeE_9ZqQli9aEkNCCQctZQ5sEKu00,598 +pytz/zoneinfo/Pacific/Auckland,sha256=gADjoyPo_QISQU6UJrAgcHp3HDaMoOFRdH-d23uBSyc,2437 +pytz/zoneinfo/Pacific/Bougainville,sha256=hWE86eXnNx-vABbp7-YSIqWyecHPMIWLftVloAoPhL8,254 +pytz/zoneinfo/Pacific/Chatham,sha256=xhexVc5lfJ_qAv2d3HrII6lfRSxKZYBAjY2zpYkCGE8,2054 +pytz/zoneinfo/Pacific/Chuuk,sha256=nB36HBWZTdh3TlP0DLFNz1KRQ0aHIfHbp7LC4Urp9fA,172 +pytz/zoneinfo/Pacific/Easter,sha256=QbubBs_xQlvKweAnurhyHjIK4ji77Gh4G-usXul6XVM,2219 +pytz/zoneinfo/Pacific/Efate,sha256=oSxNcQYx5-1FU2_yHzHI-hT-dMJcPxzy4XmdI1UxXAo,524 +pytz/zoneinfo/Pacific/Enderbury,sha256=HNTAKrsH_R2W3QRlKcmNld5KcXdP0ygXCjEovc1i-6Q,220 +pytz/zoneinfo/Pacific/Fakaofo,sha256=qOodpTMKjztvZIXVLe_f_kZ6WcHl9fCLE9ZsyvdFKLI,186 +pytz/zoneinfo/Pacific/Fiji,sha256=jB5FbOsCnHVQQ2ohPiWEQUPhG6JybB3Nog3qT6WJQ0I,564 +pytz/zoneinfo/Pacific/Funafuti,sha256=UyaKimsR8LjgL8Z2g65I0HTvr3tMZuA2wUeBB6_Zp9c,152 +pytz/zoneinfo/Pacific/Galapagos,sha256=_GJUYOjSiIjoNBO2qdq23isLMJ4NCVk3DKIRGeDc8BA,224 +pytz/zoneinfo/Pacific/Gambier,sha256=gAS7gr1HH_re0uYnL6eWo5KGJ-B5QaiM8mV2cY5mQxE,150 +pytz/zoneinfo/Pacific/Guadalcanal,sha256=M4kTWqaSQaV1AMhyLSvmwoBJF7X9icrILbvQJwp940g,152 +pytz/zoneinfo/Pacific/Guam,sha256=Ex9znmf6rNfGze6gNpZJCMr1TT4rkl2SnrhecrdJufI,494 +pytz/zoneinfo/Pacific/Honolulu,sha256=fwPRv1Jk56sCOi75uZfd_Iy2k2aSQHx3B2K5xUlSPzM,329 +pytz/zoneinfo/Pacific/Johnston,sha256=fwPRv1Jk56sCOi75uZfd_Iy2k2aSQHx3B2K5xUlSPzM,329 +pytz/zoneinfo/Pacific/Kanton,sha256=HNTAKrsH_R2W3QRlKcmNld5KcXdP0ygXCjEovc1i-6Q,220 +pytz/zoneinfo/Pacific/Kiritimati,sha256=hYk1Ooz-Lj1PuZCbNV2WJIvOLtCwSwq2u63cb1Z-3NQ,224 +pytz/zoneinfo/Pacific/Kosrae,sha256=Q0jrb4zeDrd61bU4V8TqjMc0Iep8rWZyZqJ0uqsunxs,337 +pytz/zoneinfo/Pacific/Kwajalein,sha256=TmZ_0f-ySQ-saBAlRXV0f49Itwne51VBXn6rWcrWqHQ,302 +pytz/zoneinfo/Pacific/Majuro,sha256=UyaKimsR8LjgL8Z2g65I0HTvr3tMZuA2wUeBB6_Zp9c,152 +pytz/zoneinfo/Pacific/Marquesas,sha256=FTxPJTWtk48LVb3N2U64KLpLsmvu0DQBubTCg-dvyGM,159 +pytz/zoneinfo/Pacific/Midway,sha256=fCYrYphYY6rUfxOw712y5cyRe104AC3pouqD3bCINFg,175 +pytz/zoneinfo/Pacific/Nauru,sha256=9ASKgLHB-8nsTEK1ApzfTH0yQtbNAmGX-JI7uHZiqnA,238 +pytz/zoneinfo/Pacific/Niue,sha256=OllXxukncR7a-SMmdFox5az1xpIPMhbahQhtObmpuDM,189 +pytz/zoneinfo/Pacific/Norfolk,sha256=DMdX1Bm18lzNuiCWzwfeHUMRGXPS8v5AWnh-_EX_AZw,866 +pytz/zoneinfo/Pacific/Noumea,sha256=tkHxxnxsXTOqz3YzWi0mkhTCIONzg-W7EpSRMdPjKdQ,290 +pytz/zoneinfo/Pacific/Pago_Pago,sha256=fCYrYphYY6rUfxOw712y5cyRe104AC3pouqD3bCINFg,175 +pytz/zoneinfo/Pacific/Palau,sha256=aN2HbT0reqwKrtLKDK9M2zb0d0ikdNlTrrntVxdH66o,166 +pytz/zoneinfo/Pacific/Pitcairn,sha256=U4jAUuvsRNoy8XrPa16YpcXCcqHJY0u6JvCNgPEWO1c,188 +pytz/zoneinfo/Pacific/Pohnpei,sha256=M4kTWqaSQaV1AMhyLSvmwoBJF7X9icrILbvQJwp940g,152 +pytz/zoneinfo/Pacific/Ponape,sha256=M4kTWqaSQaV1AMhyLSvmwoBJF7X9icrILbvQJwp940g,152 +pytz/zoneinfo/Pacific/Port_Moresby,sha256=nB36HBWZTdh3TlP0DLFNz1KRQ0aHIfHbp7LC4Urp9fA,172 +pytz/zoneinfo/Pacific/Rarotonga,sha256=wPEsoXbyDnuhfzkgLvUqhSzrMx_FD42uAPluSPMh3Bc,589 +pytz/zoneinfo/Pacific/Saipan,sha256=Ex9znmf6rNfGze6gNpZJCMr1TT4rkl2SnrhecrdJufI,494 +pytz/zoneinfo/Pacific/Samoa,sha256=fCYrYphYY6rUfxOw712y5cyRe104AC3pouqD3bCINFg,175 +pytz/zoneinfo/Pacific/Tahiti,sha256=BRff9G3E-iWKhOWR1Wu02Z0iMgjrwDXV-XNrqItXdTY,151 +pytz/zoneinfo/Pacific/Tarawa,sha256=UyaKimsR8LjgL8Z2g65I0HTvr3tMZuA2wUeBB6_Zp9c,152 +pytz/zoneinfo/Pacific/Tongatapu,sha256=OppBZqTAZib9HY7U9AC-JavO7m6NxPGUtUfPQAl9oBY,358 +pytz/zoneinfo/Pacific/Truk,sha256=nB36HBWZTdh3TlP0DLFNz1KRQ0aHIfHbp7LC4Urp9fA,172 +pytz/zoneinfo/Pacific/Wake,sha256=UyaKimsR8LjgL8Z2g65I0HTvr3tMZuA2wUeBB6_Zp9c,152 +pytz/zoneinfo/Pacific/Wallis,sha256=UyaKimsR8LjgL8Z2g65I0HTvr3tMZuA2wUeBB6_Zp9c,152 +pytz/zoneinfo/Pacific/Yap,sha256=nB36HBWZTdh3TlP0DLFNz1KRQ0aHIfHbp7LC4Urp9fA,172 +pytz/zoneinfo/Poland,sha256=TiLDPbeVF0ckgLVEkaSeDaKZ8wctdJDOl_HE_Wd5rKs,2654 +pytz/zoneinfo/Portugal,sha256=krB8skaJImv5NDCNHxvTPDBqpNphDFLNW84lB3lgUCw,3527 +pytz/zoneinfo/ROC,sha256=DMmQwOpPql25ue3Nf8vAKKT4em06D1Z9rHbLIitxixk,761 +pytz/zoneinfo/ROK,sha256=LI9LsV3XcJC0l-KoQf8zI-y7rk-du57erS-N2Ptdi7Q,617 +pytz/zoneinfo/Singapore,sha256=XmeVImeqcJ8hJzm7TjAti1nWJAxawOqq7jIzDnHX2hI,401 +pytz/zoneinfo/Turkey,sha256=Jk4wjndDta_uLWc8W1dWdjbavJJbsL5ROTmZboVnGKU,1933 +pytz/zoneinfo/UCT,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/US/Alaska,sha256=oZA1NSPS2BWdymYpnCHFO8BlYVS-ll5KLg2Ez9CbETs,2371 +pytz/zoneinfo/US/Aleutian,sha256=IB1DhwJQAKbhPJ9jHLf8zW5Dad7HIkBS-dhv64E1OlM,2356 +pytz/zoneinfo/US/Arizona,sha256=illz0sYuLL8lIPK0Tkou6dL0Vck_D0W_3rRTOvFYRmQ,360 +pytz/zoneinfo/US/Central,sha256=_roybr6I6sIAF6cYdIxGxoRpoef153Fty48dQ6bm9oY,3592 +pytz/zoneinfo/US/East-Indiana,sha256=kNKy9Kj9ICsiYYfCCbAggzMA7exf-GpGPMxoXocHUyw,1682 +pytz/zoneinfo/US/Eastern,sha256=6e0H177gx2qdRC0JHvHwFmj-58TyYBTAqGixn-bBipU,3552 +pytz/zoneinfo/US/Hawaii,sha256=fwPRv1Jk56sCOi75uZfd_Iy2k2aSQHx3B2K5xUlSPzM,329 +pytz/zoneinfo/US/Indiana-Starke,sha256=CsvZ5BKw2qVav3x_F8CU9taJdDk7jX41Cfsqms6jXV8,2444 +pytz/zoneinfo/US/Michigan,sha256=hecz8yqY2Cj5B61G3gLZdAVZvRgK9l0P90c_gN-uD5g,2230 +pytz/zoneinfo/US/Mountain,sha256=MugZwApDs8NI9TnXANQlUE8guNBowWQY0m-ptpPndck,2460 +pytz/zoneinfo/US/Pacific,sha256=aJd7ua1tGG_vxser02AQpm4wAI3LLTdgh6QcSYYecmg,2852 +pytz/zoneinfo/US/Samoa,sha256=fCYrYphYY6rUfxOw712y5cyRe104AC3pouqD3bCINFg,175 +pytz/zoneinfo/UTC,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/Universal,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/W-SU,sha256=KmkofRcj6T8Ph28PJChm8JVp13uRvef6TZ0GuPzUiDw,1535 +pytz/zoneinfo/WET,sha256=krB8skaJImv5NDCNHxvTPDBqpNphDFLNW84lB3lgUCw,3527 +pytz/zoneinfo/Zulu,sha256=i4WEZ5GrLIpUY8g6W-PAQ-JXDXRIQ01BOYlp7Ufj5vI,114 +pytz/zoneinfo/iso3166.tab,sha256=oBpdFY8x1GrY5vjMKgbGQYEGgqk5fUYDIPaNVCG2XnE,4791 +pytz/zoneinfo/leapseconds,sha256=gWAzwRuERloD6ADF5V6tUV26U_oVm5xh2nYC6jVwYOg,3257 +pytz/zoneinfo/tzdata.zi,sha256=-EeBhT7eWZNN3RIVLVFnf0Ekn2iff5gPQIxShpDY0lA,107471 +pytz/zoneinfo/zone.tab,sha256=WGtCB-bHZyLegq3Npr9J12H2aFF_RaZz9k2oOzM-7MQ,18822 +pytz/zoneinfo/zone1970.tab,sha256=VxlOQ7ABuPgymHshuClT2Zeu6uvrU6hSAUC8EtfYz8w,17597 +pytz/zoneinfo/zonenow.tab,sha256=JGdDvM4N0VGEYGVD2AgcuxIExbb_gje3WMtnwQkAzi8,8084 diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/WHEEL b/lib/python3.12/site-packages/pytz-2025.2.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4724c45738f6ac125bb3a21787855562e6870440 --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.42.0) +Root-Is-Purelib: true +Tag: py2-none-any +Tag: py3-none-any + diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/top_level.txt b/lib/python3.12/site-packages/pytz-2025.2.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..af44f198c687e245aada835efbab2f75ed2c9baf --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/top_level.txt @@ -0,0 +1 @@ +pytz diff --git a/lib/python3.12/site-packages/pytz-2025.2.dist-info/zip-safe b/lib/python3.12/site-packages/pytz-2025.2.dist-info/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/lib/python3.12/site-packages/pytz-2025.2.dist-info/zip-safe @@ -0,0 +1 @@ + diff --git a/lib/python3.12/site-packages/ray/util/actor_pool.py b/lib/python3.12/site-packages/ray/util/actor_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..fbdba9ce493e02ce9be16277c1605d98c5c44611 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/actor_pool.py @@ -0,0 +1,463 @@ +from typing import TYPE_CHECKING, Any, Callable, List, TypeVar + +import ray +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + import ray.actor + +V = TypeVar("V") + + +@DeveloperAPI +class ActorPool: + """Utility class to operate on a fixed pool of actors. + + Arguments: + actors: List of Ray actor handles to use in this pool. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + print(list(pool.map(lambda a, v: a.double.remote(v), + [1, 2, 3, 4]))) + + .. testoutput:: + + [2, 4, 6, 8] + """ + + def __init__(self, actors: list): + from ray._common.usage.usage_lib import record_library_usage + + record_library_usage("util.ActorPool") + + # actors to be used + self._idle_actors = list(actors) + + # get actor from future + self._future_to_actor = {} + + # get future from index + self._index_to_future = {} + + # next task to do + self._next_task_index = 0 + + # next task to return + self._next_return_index = 0 + + # next work depending when actors free + self._pending_submits = [] + + def map(self, fn: Callable[["ray.actor.ActorHandle", V], Any], values: List[V]): + """Apply the given function in parallel over the actors and values. + + This returns an ordered iterator that will return results of the map + as they finish. Note that you must iterate over the iterator to force + the computation to finish. + + Arguments: + fn: Function that takes (actor, value) as argument and + returns an ObjectRef computing the result over the value. The + actor will be considered busy until the ObjectRef completes. + values: List of values that fn(actor, value) should be + applied to. + + Returns: + Iterator over results from applying fn to the actors and values. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + print(list(pool.map(lambda a, v: a.double.remote(v), + [1, 2, 3, 4]))) + + .. testoutput:: + + [2, 4, 6, 8] + """ + # Ignore/Cancel all the previous submissions + # by calling `has_next` and `gen_next` repeteadly. + while self.has_next(): + try: + self.get_next(timeout=0, ignore_if_timedout=True) + except TimeoutError: + pass + + for v in values: + self.submit(fn, v) + + def get_generator(): + while self.has_next(): + yield self.get_next() + + return get_generator() + + def map_unordered( + self, fn: Callable[["ray.actor.ActorHandle", V], Any], values: List[V] + ): + """Similar to map(), but returning an unordered iterator. + + This returns an unordered iterator that will return results of the map + as they finish. This can be more efficient that map() if some results + take longer to compute than others. + + Arguments: + fn: Function that takes (actor, value) as argument and + returns an ObjectRef computing the result over the value. The + actor will be considered busy until the ObjectRef completes. + values: List of values that fn(actor, value) should be + applied to. + + Returns: + Iterator over results from applying fn to the actors and values. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + print(list(pool.map_unordered(lambda a, v: a.double.remote(v), + [1, 2, 3, 4]))) + + .. testoutput:: + :options: +MOCK + + [6, 8, 4, 2] + """ + # Ignore/Cancel all the previous submissions + # by calling `has_next` and `gen_next_unordered` repeteadly. + while self.has_next(): + try: + self.get_next_unordered(timeout=0) + except TimeoutError: + pass + + for v in values: + self.submit(fn, v) + + def get_generator(): + while self.has_next(): + yield self.get_next_unordered() + + return get_generator() + + def submit(self, fn, value): + """Schedule a single task to run in the pool. + + This has the same argument semantics as map(), but takes on a single + value instead of a list of values. The result can be retrieved using + get_next() / get_next_unordered(). + + Arguments: + fn: Function that takes (actor, value) as argument and + returns an ObjectRef computing the result over the value. The + actor will be considered busy until the ObjectRef completes. + value: Value to compute a result for. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + pool.submit(lambda a, v: a.double.remote(v), 1) + pool.submit(lambda a, v: a.double.remote(v), 2) + print(pool.get_next(), pool.get_next()) + + .. testoutput:: + + 2 4 + """ + if self._idle_actors: + actor = self._idle_actors.pop() + future = fn(actor, value) + future_key = tuple(future) if isinstance(future, list) else future + self._future_to_actor[future_key] = (self._next_task_index, actor) + self._index_to_future[self._next_task_index] = future + self._next_task_index += 1 + else: + self._pending_submits.append((fn, value)) + + def has_next(self): + """Returns whether there are any pending results to return. + + Returns: + True if there are any pending results not yet returned. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + pool.submit(lambda a, v: a.double.remote(v), 1) + print(pool.has_next()) + print(pool.get_next()) + print(pool.has_next()) + + .. testoutput:: + + True + 2 + False + """ + return bool(self._future_to_actor) + + def get_next(self, timeout=None, ignore_if_timedout=False): + """Returns the next pending result in order. + + This returns the next result produced by submit(), blocking for up to + the specified timeout until it is available. + + Returns: + The next result. + + Raises: + TimeoutError: if the timeout is reached. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + pool.submit(lambda a, v: a.double.remote(v), 1) + print(pool.get_next()) + + .. testoutput:: + + 2 + """ + if not self.has_next(): + raise StopIteration("No more results to get") + if self._next_return_index >= self._next_task_index: + raise ValueError( + "It is not allowed to call get_next() after get_next_unordered()." + ) + future = self._index_to_future[self._next_return_index] + timeout_msg = "Timed out waiting for result" + raise_timeout_after_ignore = False + if timeout is not None: + res, _ = ray.wait([future], timeout=timeout) + if not res: + if not ignore_if_timedout: + raise TimeoutError(timeout_msg) + else: + raise_timeout_after_ignore = True + del self._index_to_future[self._next_return_index] + self._next_return_index += 1 + + future_key = tuple(future) if isinstance(future, list) else future + i, a = self._future_to_actor.pop(future_key) + + self._return_actor(a) + if raise_timeout_after_ignore: + raise TimeoutError( + timeout_msg + ". The task {} has been ignored.".format(future) + ) + return ray.get(future) + + def get_next_unordered(self, timeout=None, ignore_if_timedout=False): + """Returns any of the next pending results. + + This returns some result produced by submit(), blocking for up to + the specified timeout until it is available. Unlike get_next(), the + results are not always returned in same order as submitted, which can + improve performance. + + Returns: + The next result. + + Raises: + TimeoutError: if the timeout is reached. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1, a2]) + pool.submit(lambda a, v: a.double.remote(v), 1) + pool.submit(lambda a, v: a.double.remote(v), 2) + print(pool.get_next_unordered()) + print(pool.get_next_unordered()) + + .. testoutput:: + :options: +MOCK + + 4 + 2 + """ + if not self.has_next(): + raise StopIteration("No more results to get") + # TODO(ekl) bulk wait for performance + res, _ = ray.wait(list(self._future_to_actor), num_returns=1, timeout=timeout) + timeout_msg = "Timed out waiting for result" + raise_timeout_after_ignore = False + if res: + [future] = res + else: + if not ignore_if_timedout: + raise TimeoutError(timeout_msg) + else: + raise_timeout_after_ignore = True + i, a = self._future_to_actor.pop(future) + self._return_actor(a) + del self._index_to_future[i] + self._next_return_index = max(self._next_return_index, i + 1) + if raise_timeout_after_ignore: + raise TimeoutError( + timeout_msg + ". The task {} has been ignored.".format(future) + ) + return ray.get(future) + + def _return_actor(self, actor): + self._idle_actors.append(actor) + if self._pending_submits: + self.submit(*self._pending_submits.pop(0)) + + def has_free(self): + """Returns whether there are any idle actors available. + + Returns: + True if there are any idle actors and no pending submits. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1 = Actor.remote() + pool = ActorPool([a1]) + pool.submit(lambda a, v: a.double.remote(v), 1) + print(pool.has_free()) + print(pool.get_next()) + print(pool.has_free()) + + .. testoutput:: + + False + 2 + True + """ + return len(self._idle_actors) > 0 and len(self._pending_submits) == 0 + + def pop_idle(self): + """Removes an idle actor from the pool. + + Returns: + An idle actor if one is available. + None if no actor was free to be removed. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1 = Actor.remote() + pool = ActorPool([a1]) + pool.submit(lambda a, v: a.double.remote(v), 1) + assert pool.pop_idle() is None + assert pool.get_next() == 2 + assert pool.pop_idle() == a1 + + """ + if self.has_free(): + return self._idle_actors.pop() + return None + + def push(self, actor): + """Pushes a new actor into the current list of idle actors. + + Examples: + .. testcode:: + + import ray + from ray.util.actor_pool import ActorPool + + @ray.remote + class Actor: + def double(self, v): + return 2 * v + + a1, a2 = Actor.remote(), Actor.remote() + pool = ActorPool([a1]) + pool.push(a2) + """ + busy_actors = [] + if self._future_to_actor.values(): + _, busy_actors = zip(*self._future_to_actor.values()) + if actor in self._idle_actors or actor in busy_actors: + raise ValueError("Actor already belongs to current ActorPool") + else: + self._return_actor(actor) diff --git a/lib/python3.12/site-packages/ray/util/dask/__init__.py b/lib/python3.12/site-packages/ray/util/dask/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9e4ac0cb1afe41031fc734bab770990e1bfbbd5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/__init__.py @@ -0,0 +1,75 @@ +import dask +from packaging.version import Version + +# Version(dask.__version__) becomes "0" during doc builds. +if Version(dask.__version__) != Version("0") and Version(dask.__version__) < Version( + "2024.11.0" +): + # Dask on Ray doesn't work if Dask version is less than 2024.11.0. + raise ImportError( + "Dask on Ray requires Dask version 2024.11.0 or later. " + "Please upgrade your Dask installation." + ) + +from .callbacks import ( + ProgressBarCallback, + RayDaskCallback, + local_ray_callbacks, + unpack_ray_callbacks, +) +from .optimizations import dataframe_optimize +from .scheduler import ( + disable_dask_on_ray, + enable_dask_on_ray, + ray_dask_get, + ray_dask_get_sync, +) + +dask_persist = dask.persist + + +def ray_dask_persist(*args, **kwargs): + kwargs["ray_persist"] = True + return dask_persist(*args, **kwargs) + + +ray_dask_persist.__doc__ = dask_persist.__doc__ + +dask_persist_mixin = dask.base.DaskMethodsMixin.persist + + +def ray_dask_persist_mixin(self, **kwargs): + kwargs["ray_persist"] = True + return dask_persist_mixin(self, **kwargs) + + +ray_dask_persist_mixin.__doc__ = dask_persist_mixin.__doc__ + + +# We patch dask in order to inject a kwarg into its `dask.persist()` calls, +# which the Dask-on-Ray scheduler needs. +# FIXME(Clark): Monkey patching is bad and we should try to avoid this. +def patch_dask(ray_dask_persist, ray_dask_persist_mixin): + dask.persist = ray_dask_persist + dask.base.DaskMethodsMixin.persist = ray_dask_persist_mixin + + +patch_dask(ray_dask_persist, ray_dask_persist_mixin) + +__all__ = [ + # Config + "enable_dask_on_ray", + "disable_dask_on_ray", + # Schedulers + "ray_dask_get", + "ray_dask_get_sync", + # Helpers + "ray_dask_persist", + # Callbacks + "RayDaskCallback", + "local_ray_callbacks", + "unpack_ray_callbacks", + # Optimizations + "dataframe_optimize", + "ProgressBarCallback", +] diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43f1f606fa4944f37535eb4c9c0b0dff95ec4fb7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/callbacks.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/callbacks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a33c19e7e00e2540c97c31002ed77ba07fdc7b44 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/callbacks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb5972d185534520ec39cb5307df83b24e5a3dc1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/optimizations.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/optimizations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..746703fc994fde47b62f64c61fc86a51c1ee9714 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/optimizations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a8252cd71a892b60dcadbdd137171ce3c0dcc3e Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e112b0dcda261b8e9f8b99003df0b74de737ad0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/dask/__pycache__/scheduler_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/dask/callbacks.py b/lib/python3.12/site-packages/ray/util/dask/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..770d2208b50432891385408b838af91e7f6088c8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/callbacks.py @@ -0,0 +1,306 @@ +import contextlib +from collections import defaultdict, namedtuple +from datetime import datetime +from typing import Any, List, Optional + +from dask.callbacks import Callback + +import ray + +# The names of the Ray-specific callbacks. These are the kwarg names that +# RayDaskCallback will accept on construction, and is considered the +# source-of-truth for what Ray-specific callbacks exist. +CBS = ( + "ray_presubmit", + "ray_postsubmit", + "ray_pretask", + "ray_posttask", + "ray_postsubmit_all", + "ray_finish", +) +# The Ray-specific callback method names for RayDaskCallback. +CB_FIELDS = tuple("_" + field for field in CBS) +# The Ray-specific callbacks that we do _not_ wish to drop from RayCallbacks +# if not given on a RayDaskCallback instance (will be filled with None +# instead). +CBS_DONT_DROP = {"ray_pretask", "ray_posttask"} + +# The Ray-specific callbacks for a single RayDaskCallback. +RayCallback = namedtuple("RayCallback", " ".join(CBS)) + +# The Ray-specific callbacks for one or more RayDaskCallbacks. +RayCallbacks = namedtuple("RayCallbacks", " ".join([field + "_cbs" for field in CBS])) + + +class RayDaskCallback(Callback): + """ + Extends Dask's `Callback` class with Ray-specific hooks. When instantiating + or subclassing this class, both the normal Dask hooks (e.g. pretask, + posttask, etc.) and the Ray-specific hooks can be provided. + + See `dask.callbacks.Callback` for usage. + + Caveats: Any Dask-Ray scheduler must bring the Ray-specific callbacks into + context using the `local_ray_callbacks` context manager, since the built-in + `local_callbacks` context manager provided by Dask isn't aware of this + class. + """ + + # Set of active Ray-specific callbacks. + ray_active = set() + + def __init__(self, **kwargs): + for cb in CBS: + cb_func = kwargs.pop(cb, None) + if cb_func is not None: + setattr(self, "_" + cb, cb_func) + + super().__init__(**kwargs) + + @property + def _ray_callback(self): + return RayCallback(*[getattr(self, field, None) for field in CB_FIELDS]) + + def __enter__(self): + self._ray_cm = add_ray_callbacks(self) + self._ray_cm.__enter__() + super().__enter__() + return self + + def __exit__(self, *args): + super().__exit__(*args) + self._ray_cm.__exit__(*args) + + def register(self): + type(self).ray_active.add(self._ray_callback) + super().register() + + def unregister(self): + type(self).ray_active.remove(self._ray_callback) + super().unregister() + + def _ray_presubmit(self, task, key, deps) -> Optional[Any]: + """Run before submitting a Ray task. + + If this callback returns a non-`None` value, Ray does _not_ create + a task and uses this value as the would-be task's result value. + + Args: + task: A Dask task, where the first tuple item is + the task function, and the remaining tuple items are + the task arguments, which are either the actual argument values, + or Dask keys into the deps dictionary whose + corresponding values are the argument values. + key: The Dask graph key for the given task. + deps: The dependencies of this task. + + Returns: + Either None, in which case Ray submits a task, or + a non-None value, in which case Ray task doesn't submit + a task and uses this return value as the + would-be task result value. + """ + pass + + def _ray_postsubmit(self, task, key, deps, object_ref: ray.ObjectRef): + """Run after submitting a Ray task. + + Args: + task: A Dask task, where the first tuple item is + the task function, and the remaining tuple items are + the task arguments, which are either the actual argument values, + or Dask keys into the deps dictionary whose + corresponding values are the argument values. + key: The Dask graph key for the given task. + deps: The dependencies of this task. + object_ref: The object reference for the + return value of the Ray task. + + """ + pass + + def _ray_pretask(self, key, object_refs: List[ray.ObjectRef]): + """Run before executing a Dask task within a Ray task. + + This method executes after Ray submits the task within a Ray + worker. Ray passes the return value of this task to the + _ray_posttask callback, if provided. + + Args: + key: The Dask graph key for the Dask task. + object_refs: The object references + for the arguments of the Ray task. + + Returns: + A value that Ray passes to the corresponding + _ray_posttask callback, if the callback is defined. + """ + pass + + def _ray_posttask(self, key, result, pre_state): + """Run after executing a Dask task within a Ray task. + + This method executes within a Ray worker. This callback receives the + return value of the _ray_pretask callback, if provided. + + Args: + key: The Dask graph key for the Dask task. + result: The task result value. + pre_state: The return value of the corresponding + _ray_pretask callback, if said callback is defined. + """ + pass + + def _ray_postsubmit_all(self, object_refs: List[ray.ObjectRef], dsk): + """Run after Ray submits all tasks. + + Args: + object_refs: The object references + for the output (leaf) Ray tasks of the task graph. + dsk: The Dask graph. + """ + pass + + def _ray_finish(self, result): + """Run after Ray finishes executing all Ray tasks and returns the final + result. + + Args: + result: The final result (output) of the Dask + computation, before any repackaging is done by + Dask collection-specific post-compute callbacks. + """ + pass + + +class add_ray_callbacks: + def __init__(self, *callbacks): + self.callbacks = [normalize_ray_callback(c) for c in callbacks] + RayDaskCallback.ray_active.update(self.callbacks) + + def __enter__(self): + return self + + def __exit__(self, *args): + for c in self.callbacks: + RayDaskCallback.ray_active.discard(c) + + +def normalize_ray_callback(cb): + if isinstance(cb, RayDaskCallback): + return cb._ray_callback + elif isinstance(cb, RayCallback): + return cb + else: + raise TypeError( + "Callbacks must be either 'RayDaskCallback' or 'RayCallback' namedtuple" + ) + + +def unpack_ray_callbacks(cbs): + """Take an iterable of callbacks, return a list of each callback.""" + if cbs: + # Only drop callback methods that aren't in CBS_DONT_DROP. + return RayCallbacks( + *( + [cb for cb in cbs_ if cb or CBS[idx] in CBS_DONT_DROP] or None + for idx, cbs_ in enumerate(zip(*cbs)) + ) + ) + else: + return RayCallbacks(*([()] * len(CBS))) + + +@contextlib.contextmanager +def local_ray_callbacks(callbacks=None): + """ + Allows Dask-Ray callbacks to work with nested schedulers. + + Callbacks will only be used by the first started scheduler they encounter. + This means that only the outermost scheduler will use global callbacks. + """ + global_callbacks = callbacks is None + if global_callbacks: + callbacks, RayDaskCallback.ray_active = (RayDaskCallback.ray_active, set()) + try: + yield callbacks or () + finally: + if global_callbacks: + RayDaskCallback.ray_active = callbacks + + +class ProgressBarCallback(RayDaskCallback): + def __init__(self): + @ray.remote + class ProgressBarActor: + def __init__(self): + self._init() + + def submit(self, key, deps, now): + for dep in deps.keys(): + self.deps[key].add(dep) + self.submitted[key] = now + self.submission_queue.append((key, now)) + + def task_scheduled(self, key, now): + self.scheduled[key] = now + + def finish(self, key, now): + self.finished[key] = now + + def result(self): + return len(self.submitted), len(self.finished) + + def report(self): + result = defaultdict(dict) + for key, finished in self.finished.items(): + submitted = self.submitted[key] + scheduled = self.scheduled[key] + # deps = self.deps[key] + result[key]["execution_time"] = ( + finished - scheduled + ).total_seconds() + # Calculate the scheduling time. + # This is inaccurate. + # We should subtract scheduled - (last dep completed). + # But currently it is not easy because + # of how getitem is implemented in dask on ray sort. + result[key]["scheduling_time"] = ( + scheduled - submitted + ).total_seconds() + result["submission_order"] = self.submission_queue + return result + + def ready(self): + pass + + def reset(self): + self._init() + + def _init(self): + self.submission_queue = [] + self.submitted = defaultdict(None) + self.scheduled = defaultdict(None) + self.finished = defaultdict(None) + self.deps = defaultdict(set) + + try: + self.pb = ray.get_actor("_dask_on_ray_pb") + ray.get(self.pb.reset.remote()) + except ValueError: + self.pb = ProgressBarActor.options(name="_dask_on_ray_pb").remote() + ray.get(self.pb.ready.remote()) + + def _ray_postsubmit(self, task, key, deps, object_ref): + # Indicate the dask task is submitted. + self.pb.submit.remote(key, deps, datetime.now()) + + def _ray_pretask(self, key, object_refs): + self.pb.task_scheduled.remote(key, datetime.now()) + + def _ray_posttask(self, key, result, pre_state): + # Indicate the dask task is finished. + self.pb.finish.remote(key, datetime.now()) + + def _ray_finish(self, result): + print("All tasks are completed.") diff --git a/lib/python3.12/site-packages/ray/util/dask/common.py b/lib/python3.12/site-packages/ray/util/dask/common.py new file mode 100644 index 0000000000000000000000000000000000000000..47ec12d79a1bd3ca37951981556e1e9fd64ea7f0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/common.py @@ -0,0 +1,87 @@ +import uuid +from collections import OrderedDict +from collections.abc import Iterator +from operator import getitem + +from dask.core import get as get_sync, quote +from dask.utils import apply + +import ray + +try: + from dataclasses import fields as dataclass_fields, is_dataclass +except ImportError: + # Python < 3.7 + def is_dataclass(x): + return False + + def dataclass_fields(x): + return [] + + +def unpack_object_refs(*args): + """ + Extract Ray object refs from a set of potentially arbitrarily nested + Python objects. + + Intended use is to find all Ray object references in a set of (possibly + nested) Python objects, do something to them (get(), wait(), etc.), then + repackage them into equivalent Python objects. + + Args: + *args: One or more (potentially nested) Python objects that contain + Ray object references. + + Returns: + A 2-tuple of a flat list of all contained Ray object references, and a + function that, when given the corresponding flat list of concrete + values, will return a set of Python objects equivalent to that which + was given in *args, but with all Ray object references replaced with + their corresponding concrete values. + """ + object_refs = [] + repack_dsk = {} + + object_refs_token = uuid.uuid4().hex + + def _unpack(expr): + if isinstance(expr, ray.ObjectRef): + token = expr.hex() + repack_dsk[token] = (getitem, object_refs_token, len(object_refs)) + object_refs.append(expr) + return token + + token = uuid.uuid4().hex + # Treat iterators like lists + typ = list if isinstance(expr, Iterator) else type(expr) + if typ in (list, tuple, set): + repack_task = (typ, [_unpack(i) for i in expr]) + elif typ in (dict, OrderedDict): + repack_task = (typ, [[_unpack(k), _unpack(v)] for k, v in expr.items()]) + elif is_dataclass(expr): + repack_task = ( + apply, + typ, + (), + ( + dict, + [ + [f.name, _unpack(getattr(expr, f.name))] + for f in dataclass_fields(expr) + ], + ), + ) + else: + return expr + repack_dsk[token] = repack_task + return token + + out = uuid.uuid4().hex + repack_dsk[out] = (tuple, [_unpack(i) for i in args]) + + def repack(results): + dsk = repack_dsk.copy() + dsk[object_refs_token] = quote(results) + return get_sync(dsk, out) + + return object_refs, repack diff --git a/lib/python3.12/site-packages/ray/util/dask/optimizations.py b/lib/python3.12/site-packages/ray/util/dask/optimizations.py new file mode 100644 index 0000000000000000000000000000000000000000..e88416774c6de57686a1f9f7b605dbe46698325e --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/optimizations.py @@ -0,0 +1,149 @@ +import warnings + +import dask +from dask import core +from dask.dataframe.core import _concat +from dask.highlevelgraph import HighLevelGraph + +from .scheduler import MultipleReturnFunc, multiple_return_get + +try: + from dask.dataframe.optimize import optimize + from dask.dataframe.shuffle import SimpleShuffleLayer, shuffle_group +except ImportError: + # SimpleShuffleLayer doesn't exist in this version of Dask. + # This is the case for dask>=2025.1.0. + SimpleShuffleLayer = None +try: + import dask_expr # noqa: F401 + + SimpleShuffleLayer = None +except ImportError: + pass + + +if SimpleShuffleLayer is not None: + + class MultipleReturnSimpleShuffleLayer(SimpleShuffleLayer): + @classmethod + def clone(cls, layer: SimpleShuffleLayer): + # TODO(Clark): Probably don't need this since SimpleShuffleLayer + # implements __copy__() and the shallow clone should be enough? + return cls( + name=layer.name, + column=layer.column, + npartitions=layer.npartitions, + npartitions_input=layer.npartitions_input, + ignore_index=layer.ignore_index, + name_input=layer.name_input, + meta_input=layer.meta_input, + parts_out=layer.parts_out, + annotations=layer.annotations, + ) + + def __repr__(self): + return ( + f"MultipleReturnSimpleShuffleLayer" + ) + + def __reduce__(self): + attrs = [ + "name", + "column", + "npartitions", + "npartitions_input", + "ignore_index", + "name_input", + "meta_input", + "parts_out", + "annotations", + ] + return ( + MultipleReturnSimpleShuffleLayer, + tuple(getattr(self, attr) for attr in attrs), + ) + + def _cull(self, parts_out): + return MultipleReturnSimpleShuffleLayer( + self.name, + self.column, + self.npartitions, + self.npartitions_input, + self.ignore_index, + self.name_input, + self.meta_input, + parts_out=parts_out, + ) + + def _construct_graph(self): + """Construct graph for a simple shuffle operation.""" + + shuffle_group_name = "group-" + self.name + shuffle_split_name = "split-" + self.name + + dsk = {} + n_parts_out = len(self.parts_out) + for part_out in self.parts_out: + # TODO(Clark): Find better pattern than in-scheduler concat. + _concat_list = [ + (shuffle_split_name, part_out, part_in) + for part_in in range(self.npartitions_input) + ] + dsk[(self.name, part_out)] = (_concat, _concat_list, self.ignore_index) + for _, _part_out, _part_in in _concat_list: + dsk[(shuffle_split_name, _part_out, _part_in)] = ( + multiple_return_get, + (shuffle_group_name, _part_in), + _part_out, + ) + if (shuffle_group_name, _part_in) not in dsk: + dsk[(shuffle_group_name, _part_in)] = ( + MultipleReturnFunc( + shuffle_group, + n_parts_out, + ), + (self.name_input, _part_in), + self.column, + 0, + self.npartitions, + self.npartitions, + self.ignore_index, + self.npartitions, + ) + + return dsk + + def rewrite_simple_shuffle_layer(dsk, keys): + if not isinstance(dsk, HighLevelGraph): + dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=()) + else: + dsk = dsk.copy() + + layers = dsk.layers.copy() + for key, layer in layers.items(): + if type(layer) is SimpleShuffleLayer: + dsk.layers[key] = MultipleReturnSimpleShuffleLayer.clone(layer) + return dsk + + def dataframe_optimize(dsk, keys, **kwargs): + if not isinstance(keys, (list, set)): + keys = [keys] + keys = list(core.flatten(keys)) + + if not isinstance(dsk, HighLevelGraph): + dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=()) + + dsk = rewrite_simple_shuffle_layer(dsk, keys=keys) + return optimize(dsk, keys, **kwargs) + +else: + + def dataframe_optimize(dsk, keys, **kwargs): + warnings.warn( + "Custom dataframe shuffle optimization only works on " + "dask>=2024.11.0,<2025.1.0, you are on version " + f"{dask.__version__}." + "Doing no additional optimization aside from the default one." + ) + return None diff --git a/lib/python3.12/site-packages/ray/util/dask/scheduler.py b/lib/python3.12/site-packages/ray/util/dask/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..0fa94706187f7062276f354b121c7d64ed01fffd --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/scheduler.py @@ -0,0 +1,678 @@ +import atexit +import threading +import time +import warnings +from collections import OrderedDict, defaultdict +from collections.abc import Mapping +from dataclasses import dataclass +from multiprocessing.pool import ThreadPool +from pprint import pprint +from typing import Optional + +import dask +from dask.core import ishashable, istask + +try: + from dask._task_spec import Alias, DataNode, Task, TaskRef, convert_legacy_graph +except ImportError: + warnings.warn( + "Dask on Ray is available only on dask>=2024.11.0, " + f"you are on version {dask.__version__}." + ) +from dask.system import CPU_COUNT +from dask.threaded import _thread_get_id, pack_exception + +import ray +from ray.util.dask.callbacks import local_ray_callbacks, unpack_ray_callbacks +from ray.util.dask.common import unpack_object_refs +from ray.util.dask.scheduler_utils import apply_sync, get_async + +main_thread = threading.current_thread() +default_pool = None +pools = defaultdict(dict) +pools_lock = threading.Lock() + +TOP_LEVEL_RESOURCES_ERR_MSG = ( + 'Use ray_remote_args={"resources": {...}} instead of resources={...} to specify ' + "required Ray task resources; see " + "https://docs.ray.io/en/master/ray-core/package-ref.html#ray-remote." +) + + +def enable_dask_on_ray( + shuffle: Optional[str] = "tasks", + use_shuffle_optimization: Optional[bool] = True, +) -> dask.config.set: + """ + Enable Dask-on-Ray scheduler. This helper sets the Dask-on-Ray scheduler + as the default Dask scheduler in the Dask config. By default, it will also + cause the task-based shuffle to be used for any Dask shuffle operations + (required for multi-node Ray clusters, not sharing a filesystem), and will + enable a Ray-specific shuffle optimization. + + >>> enable_dask_on_ray() + >>> ddf.compute() # <-- will use the Dask-on-Ray scheduler. + + If used as a context manager, the Dask-on-Ray scheduler will only be used + within the context's scope. + + >>> with enable_dask_on_ray(): + ... ddf.compute() # <-- will use the Dask-on-Ray scheduler. + >>> ddf.compute() # <-- won't use the Dask-on-Ray scheduler. + + Args: + shuffle: The shuffle method used by Dask, either "tasks" or + "disk". This should be "tasks" if using a multi-node Ray cluster. + Defaults to "tasks". + use_shuffle_optimization: Enable our custom Ray-specific shuffle + optimization. Defaults to True. + Returns: + The Dask config object, which can be used as a context manager to limit + the scope of the Dask-on-Ray scheduler to the corresponding context. + """ + if use_shuffle_optimization: + from ray.util.dask.optimizations import dataframe_optimize + else: + dataframe_optimize = None + # Manually set the global Dask scheduler config. + # We also force the task-based shuffle to be used since the disk-based + # shuffle doesn't work for a multi-node Ray cluster that doesn't share + # the filesystem. + return dask.config.set( + scheduler=ray_dask_get, shuffle=shuffle, dataframe_optimize=dataframe_optimize + ) + + +def disable_dask_on_ray(): + """ + Unsets the scheduler, shuffle method, and DataFrame optimizer. + """ + return dask.config.set(scheduler=None, shuffle=None, dataframe_optimize=None) + + +def ray_dask_get(dsk, keys, **kwargs): + """ + A Dask-Ray scheduler. This scheduler will send top-level (non-inlined) Dask + tasks to a Ray cluster for execution. The scheduler will wait for the + tasks to finish executing, fetch the results, and repackage them into the + appropriate Dask collections. This particular scheduler uses a threadpool + to submit Ray tasks. + + This can be passed directly to `dask.compute()`, as the scheduler: + + >>> dask.compute(obj, scheduler=ray_dask_get) + + You can override the currently active global Dask-Ray callbacks (e.g. + supplied via a context manager), the number of threads to use when + submitting the Ray tasks, or the threadpool used to submit Ray tasks: + + >>> dask.compute( + obj, + scheduler=ray_dask_get, + ray_callbacks=some_ray_dask_callbacks, + num_workers=8, + pool=some_cool_pool, + ) + + Args: + dsk: Dask graph, represented as a task DAG dictionary. + keys (List[str]): List of Dask graph keys whose values we wish to + compute and return. + ray_callbacks (Optional[list[callable]]): Dask-Ray callbacks. + num_workers (Optional[int]): The number of worker threads to use in + the Ray task submission traversal of the Dask graph. + pool (Optional[ThreadPool]): A multiprocessing threadpool to use to + submit Ray tasks. + + Returns: + Computed values corresponding to the provided keys. + """ + num_workers = kwargs.pop("num_workers", None) + pool = kwargs.pop("pool", None) + # We attempt to reuse any other thread pools that have been created within + # this thread and with the given number of workers. We reuse a global + # thread pool if num_workers is not given and we're in the main thread. + global default_pool + thread = threading.current_thread() + if pool is None: + with pools_lock: + if num_workers is None and thread is main_thread: + if default_pool is None: + default_pool = ThreadPool(CPU_COUNT) + atexit.register(default_pool.close) + pool = default_pool + elif thread in pools and num_workers in pools[thread]: + pool = pools[thread][num_workers] + else: + pool = ThreadPool(num_workers) + atexit.register(pool.close) + pools[thread][num_workers] = pool + + ray_callbacks = kwargs.pop("ray_callbacks", None) + persist = kwargs.pop("ray_persist", False) + enable_progress_bar = kwargs.pop("_ray_enable_progress_bar", None) + + # Handle Ray remote args and resource annotations. + if "resources" in kwargs: + raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG) + ray_remote_args = kwargs.pop("ray_remote_args", {}) + annotations = dask.get_annotations() + if "resources" in annotations: + raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG) + + # Take out the dask graph if it is an Expr for dask>=2025.4.0. + if not isinstance(dsk, Mapping): + if hasattr(dsk, "_optimized_dsk"): + # For Expr with this property + dsk = dsk._optimized_dsk + else: + # For any other Expr + dsk = dsk.__dask_graph__() + scoped_ray_remote_args = _build_key_scoped_ray_remote_args( + dsk, annotations, ray_remote_args + ) + + with local_ray_callbacks(ray_callbacks) as ray_callbacks: + # Unpack the Ray-specific callbacks. + ( + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + ray_postsubmit_all_cbs, + ray_finish_cbs, + ) = unpack_ray_callbacks(ray_callbacks) + # Make sure the graph is in the new format + dsk = convert_legacy_graph(dsk) + # NOTE: We hijack Dask's `get_async` function, injecting a different + # task executor. + object_refs = get_async( + _apply_async_wrapper( + pool.apply_async, + _rayify_task_wrapper, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + scoped_ray_remote_args, + ), + len(pool._pool), + dsk, + keys, + get_id=_thread_get_id, + pack_exception=pack_exception, + **kwargs, + ) + if ray_postsubmit_all_cbs is not None: + for cb in ray_postsubmit_all_cbs: + cb(object_refs, dsk) + # NOTE: We explicitly delete the Dask graph here so object references + # are garbage-collected before this function returns, i.e. before all + # Ray tasks are done. Otherwise, no intermediate objects will be + # cleaned up until all Ray tasks are done. + del dsk + if persist: + result = object_refs + else: + pb_actor = None + if enable_progress_bar: + pb_actor = ray.get_actor("_dask_on_ray_pb") + result = ray_get_unpack(object_refs, progress_bar_actor=pb_actor) + if ray_finish_cbs is not None: + for cb in ray_finish_cbs: + cb(result) + + # cleanup pools associated with dead threads. + with pools_lock: + active_threads = set(threading.enumerate()) + if thread is not main_thread: + for t in list(pools): + if t not in active_threads: + for p in pools.pop(t).values(): + p.close() + return result + + +def _apply_async_wrapper(apply_async, real_func, *extra_args, **extra_kwargs): + """ + Wraps the given pool `apply_async` function, hotswapping `real_func` in as + the function to be applied and adding `extra_args` and `extra_kwargs` to + `real_func`'s call. + + Args: + apply_async: The pool function to be wrapped. + real_func: The real function that we wish the pool apply + function to execute. + *extra_args: Extra positional arguments to pass to the `real_func`. + **extra_kwargs: Extra keyword arguments to pass to the `real_func`. + + Returns: + A wrapper function that will ignore it's first `func` argument and + pass `real_func` in its place. To be passed to `dask.local.get_async`. + """ + + def wrapper(func, args=(), kwds=None, callback=None): # noqa: M511 + if not kwds: + kwds = {} + return apply_async( + real_func, + args=args + extra_args, + kwds=dict(kwds, **extra_kwargs), + callback=callback, + ) + + return wrapper + + +def _rayify_task_wrapper( + key, + task_info, + dumps, + loads, + get_id, + pack_exception, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + scoped_ray_remote_args, +): + """ + The core Ray-Dask task execution wrapper, to be given to the thread pool's + `apply_async` function. Exactly the same as `execute_task`, except that it + calls `_rayify_task` on the task instead of `_execute_task`. + + Args: + key: The Dask graph key whose corresponding task we wish to + execute. + task_info: The task to execute and its dependencies. + dumps: A result serializing function. + loads: A task_info deserializing function. + get_id: An ID generating function. + pack_exception: An exception serializing function. + ray_presubmit_cbs: Pre-task submission callbacks. + ray_postsubmit_cbs: Post-task submission callbacks. + ray_pretask_cbs: Pre-task execution callbacks. + ray_posttask_cbs: Post-task execution callbacks. + scoped_ray_remote_args: Ray task options for each key. + + Returns: + A 3-tuple of the task's key, a literal or a Ray object reference for a + Ray task's result, and whether the Ray task submission failed. + """ + try: + task, deps = loads(task_info) + result = _rayify_task( + task, + key, + deps, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + scoped_ray_remote_args.get(key, {}), + ) + id = get_id() + result = dumps((result, id)) + failed = False + except BaseException as e: + result = pack_exception(e, dumps) + failed = True + return key, result, failed + + +def _rayify_task( + task, + key, + deps, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + ray_remote_args, +): + """ + Rayifies the given task, submitting it as a Ray task to the Ray cluster. + + Args: + task: A Dask graph value, being either a literal, dependency + key, Dask task, or a list thereof. + key: The Dask graph key for the given task. + deps: The dependencies of this task. + ray_presubmit_cbs: Pre-task submission callbacks. + ray_postsubmit_cbs: Post-task submission callbacks. + ray_pretask_cbs: Pre-task execution callbacks. + ray_posttask_cbs: Post-task execution callbacks. + ray_remote_args: Ray task options. See :func:`ray.remote` for details. + + Returns: + A literal, a Ray object reference representing a submitted task, or a + list thereof. + """ + if isinstance(task, list): + # Recursively rayify this list. This will still bottom out at the first + # actual task encountered, inlining any tasks in that task's arguments. + return [ + _rayify_task( + t, + key, + deps, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + ray_remote_args, + ) + for t in task + ] + elif istask(task): + # Unpacks and repacks Ray object references and submits the task to the + # Ray cluster for execution. + if ray_presubmit_cbs is not None: + alternate_returns = [cb(task, key, deps) for cb in ray_presubmit_cbs] + for alternate_return in alternate_returns: + # We don't submit a Ray task if a presubmit callback returns + # a non-`None` value, instead we return said value. + # NOTE: This returns the first non-None presubmit callback + # return value. + if alternate_return is not None: + return alternate_return + + if isinstance(task, Alias): + target = task.target + if isinstance(target, TaskRef): + # for 2024.12.0 + return deps[target.key] + else: + # for 2024.12.1+ + return deps[target] + elif isinstance(task, Task): + func = task.func + else: + raise ValueError("Invalid task type: %s" % type(task)) + + # If the function's arguments contain nested object references, we must + # unpack said object references into a flat set of arguments so that + # Ray properly tracks the object dependencies between Ray tasks. + arg_object_refs, repack = unpack_object_refs(deps) + # Submit the task using a wrapper function. + object_refs = dask_task_wrapper.options( + name=f"dask:{key!s}", + num_returns=( + 1 if not isinstance(func, MultipleReturnFunc) else func.num_returns + ), + **ray_remote_args, + ).remote( + task, + repack, + key, + ray_pretask_cbs, + ray_posttask_cbs, + *arg_object_refs, + ) + + if ray_postsubmit_cbs is not None: + for cb in ray_postsubmit_cbs: + cb(task, key, deps, object_refs) + + return object_refs + elif not ishashable(task): + return task + elif task in deps: + return deps[task] + else: + return task + + +@ray.remote +def dask_task_wrapper( + task, repack, key, ray_pretask_cbs, ray_posttask_cbs, *arg_object_refs +): + """ + A Ray remote function acting as a Dask task wrapper. This function will + repackage the given `arg_object_refs` into its original `deps` using + `repack`, and then pass it to the provided Dask Task object , `task`. + + Args: + task: The Dask Task class object to execute. + repack: A function that repackages the provided args into + the original (possibly nested) Python objects. + key: The Dask key for this task. + ray_pretask_cbs: Pre-task execution callbacks. + ray_posttask_cbs: Post-task execution callback. + *arg_object_refs (ObjectRef): Ray object references representing the dependencies' + results. + + Returns: + The output of the Dask task. In the context of Ray, a + dask_task_wrapper.remote() invocation will return a Ray object + reference representing the Ray task's result. + """ + if ray_pretask_cbs is not None: + pre_states = [ + cb(key, arg_object_refs) if cb is not None else None + for cb in ray_pretask_cbs + ] + (repacked_deps,) = repack(arg_object_refs) + # De-reference the potentially nested arguments recursively. + def _dereference_args(x): + if isinstance(x, Task): + x.args = _dereference_args(x.args) + return x + elif isinstance(x, Mapping): + return {k: _dereference_args(v) for k, v in x.items()} + elif isinstance(x, tuple): + return tuple(_dereference_args(x) for x in x) + elif isinstance(x, ray.ObjectRef): + return ray.get(x) + elif isinstance(x, DataNode): + if isinstance(x.value, ray.ObjectRef): + value = ray.get(x.value) + return DataNode(key=x.key, value=value) + return x + else: + return x + + task = _dereference_args(task) + result = task(repacked_deps) + + if ray_posttask_cbs is not None: + for cb, pre_state in zip(ray_posttask_cbs, pre_states): + if cb is not None: + cb(key, result, pre_state) + + return result + + +def render_progress_bar(tracker, object_refs): + from tqdm import tqdm + + # At this time, every task should be submitted. + total, finished = ray.get(tracker.result.remote()) + reported_finished_so_far = 0 + pb_bar = tqdm(total=total, position=0) + pb_bar.set_description("") + + ready_refs = [] + + while finished < total: + submitted, finished = ray.get(tracker.result.remote()) + pb_bar.update(finished - reported_finished_so_far) + reported_finished_so_far = finished + ready_refs, _ = ray.wait( + object_refs, timeout=0, num_returns=len(object_refs), fetch_local=False + ) + if len(ready_refs) == len(object_refs): + break + time.sleep(0.1) + pb_bar.close() + submitted, finished = ray.get(tracker.result.remote()) + if submitted != finished: + print("Completed. There was state inconsistency.") + + pprint(ray.get(tracker.report.remote())) + + +def ray_get_unpack(object_refs, progress_bar_actor=None): + """ + Unpacks object references, gets the object references, and repacks. + Traverses arbitrary data structures. + + Args: + object_refs: A (potentially nested) Python object containing Ray object + references. + + Returns: + The input Python object with all contained Ray object references + resolved with their concrete values. + """ + + def get_result(object_refs): + if progress_bar_actor: + render_progress_bar(progress_bar_actor, object_refs) + return ray.get(object_refs) + + if isinstance(object_refs, tuple): + object_refs = list(object_refs) + + if isinstance(object_refs, list) and any( + not isinstance(x, ray.ObjectRef) for x in object_refs + ): + # We flatten the object references before calling ray.get(), since Dask + # loves to nest collections in nested tuples and Ray expects a flat + # list of object references. We repack the results after ray.get() + # completes. + object_refs, repack = unpack_object_refs(*object_refs) + computed_result = get_result(object_refs) + return repack(computed_result) + else: + return get_result(object_refs) + + +def ray_dask_get_sync(dsk, keys, **kwargs): + """ + A synchronous Dask-Ray scheduler. This scheduler will send top-level + (non-inlined) Dask tasks to a Ray cluster for execution. The scheduler will + wait for the tasks to finish executing, fetch the results, and repackage + them into the appropriate Dask collections. This particular scheduler + submits Ray tasks synchronously, which can be useful for debugging. + + This can be passed directly to `dask.compute()`, as the scheduler: + + >>> dask.compute(obj, scheduler=ray_dask_get_sync) + + You can override the currently active global Dask-Ray callbacks (e.g. + supplied via a context manager): + + >>> dask.compute( + obj, + scheduler=ray_dask_get_sync, + ray_callbacks=some_ray_dask_callbacks, + ) + + Args: + dsk: Dask graph, represented as a task DAG dictionary. + keys (List[str]): List of Dask graph keys whose values we wish to + compute and return. + + Returns: + Computed values corresponding to the provided keys. + """ + + ray_callbacks = kwargs.pop("ray_callbacks", None) + persist = kwargs.pop("ray_persist", False) + + with local_ray_callbacks(ray_callbacks) as ray_callbacks: + # Unpack the Ray-specific callbacks. + ( + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + ray_postsubmit_all_cbs, + ray_finish_cbs, + ) = unpack_ray_callbacks(ray_callbacks) + # Make sure the graph is in the new format + dsk = convert_legacy_graph(dsk) + # NOTE: We hijack Dask's `get_async` function, injecting a different + # task executor. + object_refs = get_async( + _apply_async_wrapper( + apply_sync, + _rayify_task_wrapper, + ray_presubmit_cbs, + ray_postsubmit_cbs, + ray_pretask_cbs, + ray_posttask_cbs, + ), + 1, + dsk, + keys, + **kwargs, + ) + if ray_postsubmit_all_cbs is not None: + for cb in ray_postsubmit_all_cbs: + cb(object_refs, dsk) + # NOTE: We explicitly delete the Dask graph here so object references + # are garbage-collected before this function returns, i.e. before all + # Ray tasks are done. Otherwise, no intermediate objects will be + # cleaned up until all Ray tasks are done. + del dsk + if persist: + result = object_refs + else: + result = ray_get_unpack(object_refs) + if ray_finish_cbs is not None: + for cb in ray_finish_cbs: + cb(result) + + return result + + +@dataclass +class MultipleReturnFunc: + func: callable + num_returns: int + + def __call__(self, *args, **kwargs): + returns = self.func(*args, **kwargs) + if isinstance(returns, dict) or isinstance(returns, OrderedDict): + returns = [returns[k] for k in range(len(returns))] + return returns + + +def multiple_return_get(multiple_returns, idx): + return multiple_returns[idx] + + +def _build_key_scoped_ray_remote_args(dsk, annotations, ray_remote_args): + # Handle per-layer annotations. + if not isinstance(dsk, dask.highlevelgraph.HighLevelGraph): + dsk = dask.highlevelgraph.HighLevelGraph.from_collections( + id(dsk), dsk, dependencies=() + ) + # Build key-scoped annotations. + scoped_annotations = {} + layers = [(name, dsk.layers[name]) for name in dsk._toposort_layers()] + for id_, layer in layers: + layer_annotations = layer.annotations + if layer_annotations is None: + layer_annotations = annotations + elif "resources" in layer_annotations: + raise ValueError(TOP_LEVEL_RESOURCES_ERR_MSG) + for key in layer.get_output_keys(): + layer_annotations_for_key = annotations.copy() + # Layer annotations override global annotations. + layer_annotations_for_key.update(layer_annotations) + # Let same-key annotations earlier in the topological sort take precedence. + layer_annotations_for_key.update(scoped_annotations.get(key, {})) + scoped_annotations[key] = layer_annotations_for_key + # Build key-scoped Ray remote args. + scoped_ray_remote_args = {} + for key, annotations in scoped_annotations.items(): + layer_ray_remote_args = ray_remote_args.copy() + # Layer Ray remote args override global Ray remote args given in the compute + # call. + layer_ray_remote_args.update(annotations.get("ray_remote_args", {})) + scoped_ray_remote_args[key] = layer_ray_remote_args + return scoped_ray_remote_args diff --git a/lib/python3.12/site-packages/ray/util/dask/scheduler_utils.py b/lib/python3.12/site-packages/ray/util/dask/scheduler_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b4c840c6b896d1d129ca3cf0bd1c50b5baf27645 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/dask/scheduler_utils.py @@ -0,0 +1,382 @@ +""" +The following is adapted from Dask release 2021.03.1: + https://github.com/dask/dask/blob/2021.03.1/dask/local.py +""" + +import os +import warnings +from queue import Empty, Queue + +import dask +from dask import config + +try: + from dask._task_spec import DataNode, DependenciesMapping +except ImportError: + warnings.warn( + "Dask on Ray is available only on dask>=2024.11.0, " + f"you are on version {dask.__version__}." + ) +from dask.callbacks import local_callbacks, unpack_callbacks +from dask.core import flatten, get_dependencies, reverse_dict +from dask.order import order + +if os.name == "nt": + # Python 3 windows Queue.get doesn't handle interrupts properly. To + # workaround this we poll at a sufficiently large interval that it + # shouldn't affect performance, but small enough that users trying to kill + # an application shouldn't care. + def queue_get(q): + while True: + try: + return q.get(block=True, timeout=0.1) + except Empty: + pass + +else: + + def queue_get(q): + return q.get() + + +def start_state_from_dask(dsk, cache=None, sortkey=None): + """Start state from a dask + Examples + -------- + >>> dsk = { + 'x': 1, + 'y': 2, + 'z': (inc, 'x'), + 'w': (add, 'z', 'y')} # doctest: +SKIP + >>> from pprint import pprint # doctest: +SKIP + >>> pprint(start_state_from_dask(dsk)) # doctest: +SKIP + {'cache': {'x': 1, 'y': 2}, + 'dependencies': {'w': {'z', 'y'}, 'x': set(), 'y': set(), 'z': {'x'}}, + 'dependents': {'w': set(), 'x': {'z'}, 'y': {'w'}, 'z': {'w'}}, + 'finished': set(), + 'ready': ['z'], + 'released': set(), + 'running': set(), + 'waiting': {'w': {'z'}}, + 'waiting_data': {'x': {'z'}, 'y': {'w'}, 'z': {'w'}}} + """ + if sortkey is None: + sortkey = order(dsk).get + if cache is None: + cache = config.get("cache", None) + if cache is None: + cache = dict() + + data_keys = set() + for k, v in dsk.items(): + if isinstance(v, DataNode): + cache[k] = v() + data_keys.add(k) + + dsk2 = dsk.copy() + dsk2.update(cache) + + dependencies = DependenciesMapping(dsk) + waiting = {k: set(v) for k, v in dependencies.items() if k not in data_keys} + + dependents = reverse_dict(dependencies) + for a in cache: + for b in dependents.get(a, ()): + waiting[b].remove(a) + waiting_data = {k: v.copy() for k, v in dependents.items() if v} + + ready_set = {k for k, v in waiting.items() if not v} + ready = sorted(ready_set, key=sortkey, reverse=True) + waiting = {k: v for k, v in waiting.items() if v} + + state = { + "dependencies": dependencies, + "dependents": dependents, + "waiting": waiting, + "waiting_data": waiting_data, + "cache": cache, + "ready": ready, + "running": set(), + "finished": set(), + "released": set(), + } + + return state + + +def execute_task(key, task_info, dumps, loads, get_id, pack_exception): + """ + Compute task and handle all administration + See Also + -------- + _execute_task : actually execute task + """ + try: + task, data = loads(task_info) + result = task(data) + id = get_id() + result = dumps((result, id)) + failed = False + except BaseException as e: + result = pack_exception(e, dumps) + failed = True + return key, result, failed + + +def release_data(key, state, delete=True): + """Remove data from temporary storage + See Also + -------- + finish_task + """ + if key in state["waiting_data"]: + assert not state["waiting_data"][key] + del state["waiting_data"][key] + + state["released"].add(key) + + if delete: + del state["cache"][key] + + +DEBUG = False + + +def finish_task( + dsk, key, state, results, sortkey, delete=True, release_data=release_data +): + """ + Update execution state after a task finishes + Mutates. This should run atomically (with a lock). + """ + for dep in sorted(state["dependents"][key], key=sortkey, reverse=True): + s = state["waiting"][dep] + s.remove(key) + if not s: + del state["waiting"][dep] + state["ready"].append(dep) + + for dep in state["dependencies"][key]: + if dep in state["waiting_data"]: + s = state["waiting_data"][dep] + s.remove(key) + if not s and dep not in results: + if DEBUG: + from chest.core import nbytes + + print( + "Key: %s\tDep: %s\t NBytes: %.2f\t Release" + % (key, dep, sum(map(nbytes, state["cache"].values()) / 1e6)) + ) + release_data(dep, state, delete=delete) + elif delete and dep not in results: + release_data(dep, state, delete=delete) + + state["finished"].add(key) + state["running"].remove(key) + + return state + + +def nested_get(ind, coll): + """Get nested index from collection + Examples + -------- + >>> nested_get(1, 'abc') + 'b' + >>> nested_get([1, 0], 'abc') + ('b', 'a') + >>> nested_get([[1, 0], [0, 1]], 'abc') + (('b', 'a'), ('a', 'b')) + """ + if isinstance(ind, list): + return tuple(nested_get(i, coll) for i in ind) + else: + return coll[ind] + + +def default_get_id(): + """Default get_id""" + return None + + +def default_pack_exception(e, dumps): + raise + + +def reraise(exc, tb=None): + if exc.__traceback__ is not tb: + raise exc.with_traceback(tb) + raise exc + + +def identity(x): + """Identity function. Returns x. + >>> identity(3) + 3 + """ + return x + + +def get_async( + apply_async, + num_workers, + dsk, + result, + cache=None, + get_id=default_get_id, + rerun_exceptions_locally=None, + pack_exception=default_pack_exception, + raise_exception=reraise, + callbacks=None, + dumps=identity, + loads=identity, + **kwargs, +): + """Asynchronous get function + This is a general version of various asynchronous schedulers for dask. It + takes a an apply_async function as found on Pool objects to form a more + specific ``get`` method that walks through the dask array with parallel + workers, avoiding repeat computation and minimizing memory use. + Parameters + ---------- + apply_async : function + Asynchronous apply function as found on Pool or ThreadPool + num_workers : int + The number of active tasks we should have at any one time + dsk : dict + A dask dictionary specifying a workflow + result : key or list of keys + Keys corresponding to desired data + cache : dict-like, optional + Temporary storage of results + get_id : callable, optional + Function to return the worker id, takes no arguments. Examples are + `threading.current_thread` and `multiprocessing.current_process`. + rerun_exceptions_locally : bool, optional + Whether to rerun failing tasks in local process to enable debugging + (False by default) + pack_exception : callable, optional + Function to take an exception and ``dumps`` method, and return a + serialized tuple of ``(exception, traceback)`` to send back to the + scheduler. Default is to just raise the exception. + raise_exception : callable, optional + Function that takes an exception and a traceback, and raises an error. + dumps: callable, optional + Function to serialize task data and results to communicate between + worker and parent. Defaults to identity. + loads: callable, optional + Inverse function of `dumps`. Defaults to identity. + callbacks : tuple or list of tuples, optional + Callbacks are passed in as tuples of length 5. Multiple sets of + callbacks may be passed in as a list of tuples. For more information, + see the dask.diagnostics documentation. + See Also + -------- + threaded.get + """ + queue = Queue() + + if isinstance(result, list): + result_flat = set(flatten(result)) + else: + result_flat = {result} + results = set(result_flat) + + dsk = dict(dsk) + with local_callbacks(callbacks) as callbacks: + _, _, pretask_cbs, posttask_cbs, _ = unpack_callbacks(callbacks) + started_cbs = [] + succeeded = False + # if start_state_from_dask fails, we will have something + # to pass to the final block. + state = {} + try: + for cb in callbacks: + if cb[0]: + cb[0](dsk) + started_cbs.append(cb) + + keyorder = order(dsk) + + state = start_state_from_dask(dsk, cache=cache, sortkey=keyorder.get) + + for _, start_state, _, _, _ in callbacks: + if start_state: + start_state(dsk, state) + + if rerun_exceptions_locally is None: + rerun_exceptions_locally = config.get("rerun_exceptions_locally", False) + + if state["waiting"] and not state["ready"]: + raise ValueError("Found no accessible jobs in dask") + + def fire_task(): + """Fire off a task to the thread pool""" + # Choose a good task to compute + key = state["ready"].pop() + state["running"].add(key) + for f in pretask_cbs: + f(key, dsk, state) + + # Prep data to send + data = {dep: state["cache"][dep] for dep in get_dependencies(dsk, key)} + # Submit + apply_async( + execute_task, + args=( + key, + dumps((dsk[key], data)), + dumps, + loads, + get_id, + pack_exception, + ), + callback=queue.put, + ) + + # Seed initial tasks into the thread pool + while state["ready"] and len(state["running"]) < num_workers: + fire_task() + + # Main loop, wait on tasks to finish, insert new ones + while state["waiting"] or state["ready"] or state["running"]: + key, res_info, failed = queue_get(queue) + if failed: + exc, tb = loads(res_info) + if rerun_exceptions_locally: + data = { + dep: state["cache"][dep] + for dep in get_dependencies(dsk, key) + } + task = dsk[key] + task(data) # Re-execute locally + else: + raise_exception(exc, tb) + res, worker_id = loads(res_info) + state["cache"][key] = res + finish_task(dsk, key, state, results, keyorder.get) + for f in posttask_cbs: + f(key, res, dsk, state, worker_id) + + while state["ready"] and len(state["running"]) < num_workers: + fire_task() + + succeeded = True + + finally: + for _, _, _, _, finish in started_cbs: + if finish: + finish(dsk, state, not succeeded) + + return nested_get(result, state["cache"]) + + +def apply_sync(func, args=(), kwds=None, callback=None): + """A naive synchronous version of apply_async""" + if kwds is None: + kwds = {} + + res = func(*args, **kwds) + if callback is not None: + callback(res) diff --git a/lib/python3.12/site-packages/ray/util/helpers.py b/lib/python3.12/site-packages/ray/util/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..bfc400f2ffe2356f3b24de443c270be70ecf642c --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/helpers.py @@ -0,0 +1,256 @@ +from typing import TYPE_CHECKING, Any, Iterable, Iterator, Optional, Sequence, Union + +import ray +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray import ObjectRef + from ray.remote_function import RemoteFunction + + +# ray.wait() has a default num_returns of 1. +# Using a slightly larger batch until the optimization is fully implemented, see +# https://github.com/ray-project/ray/issues/49905 +DEFAULT_CHUNK_SIZE = 10 +DEFAULT_BACKPRESSURE_SIZE = 100 + + +def _wait_and_get_single_batch( + refs: "Sequence[ObjectRef]", + *, + chunk_size: int, + yield_obj_refs: bool = False, + **kwargs, +) -> tuple[list[Union[Any, "ObjectRef"]], "list[ObjectRef]"]: + """Call ray.wait and explicitly return the ready objects/results + and remaining Ray remote refs. + + Args: + refs: A list of Ray object refs. + chunk_size: The `num_returns` parameter to pass to `ray.wait()`. + yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). + **kwargs: Additional keyword arguments to pass to `ray.wait()`. + + Returns: + A tuple of two lists, ready and not ready. This is the same as the return value of `ray.wait()`. + """ + + if chunk_size < 1: + raise ValueError("`chunk_size` must be >= 1") + + kwargs = kwargs or {} + + # num_returns must be <= len(refs) + ready, refs = ray.wait( + refs, + num_returns=min(chunk_size, len(refs)), + **kwargs, + ) + + if not yield_obj_refs: + return ray.get(ready), refs + + return ready, refs + + +@PublicAPI(stability="alpha") +def as_completed( + refs: "Sequence[ObjectRef]", + *, + chunk_size: int = DEFAULT_CHUNK_SIZE, + yield_obj_refs: bool = False, + **kwargs, +) -> Iterator[Union[Any, "ObjectRef"]]: + """Given a list of Ray task references, yield results as they become available. + + Unlike calling :meth:`~ray.get` on a list of references (i.e., `ray.get(refs)`) which + waits for all results to be ready, this function begins to yield result as soon as + a batch of `chunk_size` results are ready. + + .. note:: + Generally there is no guarantee on the order of results. For example, the first result + is not necessarily the first one completed, but rather the first one submitted in the + first available batch (See :meth:`~ray.wait` for more details about + preservation of submission order). + + .. note:: + Use this function instead of calling :meth:`~ray.get` inside a for loop. See + https://docs.ray.io/en/latest/ray-core/patterns/ray-get-loop.html for more details. + + Example: + Suppose we have a function that sleeps for x seconds depending on the input. + We expect to obtain a partially sorted list of results. + + .. testcode:: python + import ray + import time + + @ray.remote + def f(x): + time.sleep(x) + return x + + refs = [f.remote(i) for i in [10, 4, 6, 8, 2]] + for x in ray.util.as_completed(refs, chunk_size=2): + print(x) + + .. testoutput:: + :options: +MOCK + + # Output: + 4 + 2 + 6 + 8 + 10 + + Args: + refs: A list of Ray object refs. + chunk_size: The number of tasks to wait for in each iteration (default 10). + The parameter is passed as `num_returns` to :meth:`~ray.wait` internally. + yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). + **kwargs: Additional keyword arguments to pass to :meth:`~ray.wait`, e.g., + `timeout` and `fetch_local`. + + Yields: + Union[Any, ObjectRef]: The results (or optionally their Ray references) of the Ray tasks as they complete. + """ + if chunk_size < 1: + raise ValueError("`chunk_size` must be >= 1") + + if "num_returns" in kwargs: + raise ValueError("Use the `chunksize` argument instead of `num_returns`.") + + while refs: + results, refs = _wait_and_get_single_batch( + refs, + chunk_size=chunk_size, + yield_obj_refs=yield_obj_refs, + **kwargs, + ) + yield from results + + +@PublicAPI(stability="alpha") +def map_unordered( + fn: "RemoteFunction", + items: Iterable[Any], + *, + backpressure_size: Optional[int] = DEFAULT_BACKPRESSURE_SIZE, + chunk_size: int = DEFAULT_CHUNK_SIZE, + yield_obj_refs: bool = False, + **kwargs, +) -> Iterator[Union[Any, "ObjectRef"]]: + """Apply a Ray remote function to a list of items and return an iterator that yields + the completed results as they become available. + + This helper function applies backpressure to control the number of pending tasks, following the + design pattern described in + https://docs.ray.io/en/latest/ray-core/patterns/limit-pending-tasks.html. + + .. note:: + There is generally no guarantee on the order of results. + + Example: + Suppose we have a function that sleeps for x seconds depending on the input. + We expect to obtain a partially sorted list of results. + + .. testcode:: python + + import ray + import time + + @ray.remote + def f(x): + time.sleep(x) + return x + + # Example 1: chunk_size=2 + for x in ray.util.map_unordered(f, [10, 4, 6, 8, 2], chunk_size=2): + print(x) + + .. testoutput:: + :options: +MOCK + + 4 + 2 + 6 + 8 + 10 + + .. testcode:: python + + # Example 2: backpressure_size=2, chunk_size=1 + for x in ray.util.map_unordered(f, [10, 4, 6, 8, 2], backpressure_size=2, chunk_size=1): + print(x) + + .. testoutput:: + :options: +MOCK + + 4 + 10 + 6 + 8 + 2 + + Args: + fn: A remote function to apply to the list of items. For more complex use cases, use Ray Data's + :meth:`~ray.data.Dataset.map` / :meth:`~ray.data.Dataset.map_batches` instead. + items: An iterable of items to apply the function to. + backpressure_size: Maximum number of in-flight tasks allowed before + calling a blocking :meth:`~ray.wait` (default 100). If None, no backpressure is applied. + chunk_size: The number of tasks to wait for when the number of in-flight tasks exceeds + `backpressure_size`. The parameter is passed as `num_returns` to :meth:`~ray.wait` internally. + yield_obj_refs: If True, return Ray remote refs instead of results (by calling :meth:`~ray.get`). + **kwargs: Additional keyword arguments to pass to :meth:`~ray.wait`, e.g., + `timeout` and `fetch_local`. + + Yields: + Union[Any, ObjectRef]: The results (or optionally their Ray references) of the Ray tasks as they complete. + + .. seealso:: + + :meth:`~ray.util.as_completed` + Call this method for an existing list of Ray object refs. + + :meth:`~ray.data.Dataset.map` + Use Ray Data APIs (e.g., :meth:`~ray.data.Dataset.map` and :meth:`~ray.data.Dataset.map_batches`) + for better control and complex use cases, e.g., functions with multiple arguments. + + .. note:: + + This is an altenative to `pool.imap_unordered()` in Ray's Actor-based `multiprocessing.Pool`. + See https://docs.ray.io/en/latest/ray-more-libs/multiprocessing.html for more details. + + """ + + if backpressure_size is None: + backpressure_size: float = float("inf") + elif backpressure_size <= 0: + raise ValueError("backpressure_size must be positive.") + + if chunk_size < 1: + raise ValueError("`chunk_size` must be >= 1") + + if "num_returns" in kwargs: + raise ValueError("Use the `chunk_size` argument instead of `num_returns`.") + + refs = [] + for item in items: + refs.append(fn.remote(item)) + + if len(refs) >= backpressure_size: + results, refs = _wait_and_get_single_batch( + refs, + chunk_size=chunk_size, + yield_obj_refs=yield_obj_refs, + **kwargs, + ) + yield from results + else: + yield from as_completed( + refs, + chunk_size=chunk_size, + yield_obj_refs=yield_obj_refs, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/ray/util/iter.py b/lib/python3.12/site-packages/ray/util/iter.py new file mode 100644 index 0000000000000000000000000000000000000000..0e3502f1d2ec8af5916a3c71e5e21c5dcdbaab33 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/iter.py @@ -0,0 +1,1286 @@ +import collections +import random +import threading +import time +from contextlib import contextmanager +from typing import Any, Callable, Generic, Iterable, List, TypeVar + +import ray +from ray.util.annotations import Deprecated +from ray.util.iter_metrics import MetricsContext, SharedMetrics + +# The type of an iterator element. +T = TypeVar("T") +U = TypeVar("U") + + +@Deprecated +def from_items( + items: List[T], num_shards: int = 2, repeat: bool = False +) -> "ParallelIterator[T]": + """Create a parallel iterator from an existing set of objects. + + The objects will be divided round-robin among the number of shards. + + Args: + items: The list of items to iterate over. + num_shards: The number of worker actors to create. + repeat: Whether to cycle over the items forever. + """ + shards = [[] for _ in range(num_shards)] + for i, item in enumerate(items): + shards[i % num_shards].append(item) + name = "from_items[{}, {}, shards={}{}]".format( + items and type(items[0]).__name__ or "None", + len(items), + num_shards, + ", repeat=True" if repeat else "", + ) + return from_iterators(shards, repeat=repeat, name=name) + + +@Deprecated +def from_range( + n: int, num_shards: int = 2, repeat: bool = False +) -> "ParallelIterator[int]": + """Create a parallel iterator over the range 0..n. + + The range will be partitioned sequentially among the number of shards. + + Args: + n: The max end of the range of numbers. + num_shards: The number of worker actors to create. + repeat: Whether to cycle over the range forever. + """ + generators = [] + shard_size = n // num_shards + for i in range(num_shards): + start = i * shard_size + if i == num_shards - 1: + end = n + else: + end = (i + 1) * shard_size + generators.append(range(start, end)) + name = ( + f"from_range[{n}, shards={num_shards}" f"{', repeat=True' if repeat else ''}]" + ) + return from_iterators( + generators, + repeat=repeat, + name=name, + ) + + +@Deprecated +def from_iterators( + generators: List[Iterable[T]], repeat: bool = False, name=None +) -> "ParallelIterator[T]": + """Create a parallel iterator from a list of iterables. + An iterable can be a conatiner (list, str, tuple, set, etc.), + a generator, or a custom class that implements __iter__ or __getitem__. + + An actor will be created for each iterable. + + Examples: + >>> # Create using a list of generators. + >>> from_iterators([range(100), range(100)]) + + >>> # Certain generators are not serializable. + >>> from_iterators([(x for x in range(100))]) + ... TypeError: can't pickle generator objects + + >>> # So use lambda functions instead. + >>> # Lambda functions are serializable. + >>> from_iterators([lambda: (x for x in range(100))]) + + Args: + generators: A list of Python iterables or lambda + functions that produce an iterable when called. We allow lambda + functions since certain generators might not be serializable, + but a lambda that returns it can be. + repeat: Whether to cycle over the iterators forever. + name: Optional name to give the iterator. + """ + worker_cls = ray.remote(ParallelIteratorWorker) + actors = [worker_cls.remote(g, repeat) for g in generators] + if not name: + name = "from_iterators[shards={}{}]".format( + len(generators), ", repeat=True" if repeat else "" + ) + return from_actors(actors, name=name) + + +@Deprecated +def from_actors( + actors: List["ray.actor.ActorHandle"], name=None +) -> "ParallelIterator[T]": + """Create a parallel iterator from an existing set of actors. + + Each actor must subclass the ParallelIteratorWorker interface. + + Args: + actors: List of actors that each implement + ParallelIteratorWorker. + name: Optional name to give the iterator. + """ + if not name: + name = f"from_actors[shards={len(actors)}]" + return ParallelIterator([_ActorSet(actors, [])], name, parent_iterators=[]) + + +@Deprecated +class ParallelIterator(Generic[T]): + """A parallel iterator over a set of remote actors. + + This can be used to iterate over a fixed set of task results + (like an actor pool), or a stream of data (e.g., a fixed range of numbers, + an infinite stream of RLlib rollout results). + + This class is **serializable** and can be passed to other remote + tasks and actors. However, each shard should be read from at most one + process at a time. + + Examples: + >>> # Applying a function over items in parallel. + >>> it = ray.util.iter.from_items([1, 2, 3], num_shards=2) + ... <__main__.ParallelIterator object> + >>> it = it.for_each(lambda x: x * 2).gather_sync() + ... <__main__.LocalIterator object> + >>> print(list(it)) + ... [2, 4, 6] + + >>> # Creating from generators. + >>> it = ray.util.iter.from_iterators([range(3), range(3)]) + ... <__main__.ParallelIterator object> + >>> print(list(it.gather_sync())) + ... [0, 0, 1, 1, 2, 2] + + >>> # Accessing the individual shards of an iterator. + >>> it = ray.util.iter.from_range(10, num_shards=2) + ... <__main__.ParallelIterator object> + >>> it0 = it.get_shard(0) + ... <__main__.LocalIterator object> + >>> print(list(it0)) + ... [0, 1, 2, 3, 4] + >>> it1 = it.get_shard(1) + ... <__main__.LocalIterator object> + >>> print(list(it1)) + ... [5, 6, 7, 8, 9] + + >>> # Gathering results from actors synchronously in parallel. + >>> it = ray.util.iter.from_actors(workers) + ... <__main__.ParallelIterator object> + >>> it = it.batch_across_shards() + ... <__main__.LocalIterator object> + >>> print(next(it)) + ... [worker_1_result_1, worker_2_result_1] + >>> print(next(it)) + ... [worker_1_result_2, worker_2_result_2] + """ + + def __init__( + self, + actor_sets: List["_ActorSet"], + name: str, + parent_iterators: List["ParallelIterator[Any]"], + ): + """Create a parallel iterator (this is an internal function).""" + + # We track multiple sets of actors to support parallel .union(). + self.actor_sets = actor_sets + self.name = name + + # keep explicit reference to parent iterator for repartition + self.parent_iterators = parent_iterators + + def __iter__(self): + raise TypeError( + "You must use it.gather_sync() or it.gather_async() to " + "iterate over the results of a ParallelIterator." + ) + + def __str__(self): + return repr(self) + + def __repr__(self): + return f"ParallelIterator[{self.name}]" + + def _with_transform(self, local_it_fn, name): + """Helper function to create new Parallel Iterator""" + return ParallelIterator( + [a.with_transform(local_it_fn) for a in self.actor_sets], + name=self.name + name, + parent_iterators=self.parent_iterators, + ) + + def transform( + self, fn: Callable[[Iterable[T]], Iterable[U]] + ) -> "ParallelIterator[U]": + """Remotely transform the iterator. + + This is advanced version of for_each that allows you to apply arbitrary + generator transformations over the iterator. Prefer to use .for_each() + when possible for simplicity. + + Args: + fn: function to use to transform the iterator. The function + should pass through instances of _NextValueNotReady that appear + in its input iterator. Note that this function is only called + **once** over the input iterator. + + Returns: + ParallelIterator[U]: a parallel iterator. + + Examples: + >>> def f(it): + ... for x in it: + ... if x % 2 == 0: + ... yield x + >>> from_range(10, 1).transform(f).gather_sync().take(5) + ... [0, 2, 4, 6, 8] + """ + return self._with_transform( + lambda local_it: local_it.transform(fn), ".transform()" + ) + + def for_each( + self, fn: Callable[[T], U], max_concurrency=1, resources=None + ) -> "ParallelIterator[U]": + """Remotely apply fn to each item in this iterator. + + If `max_concurrency` == 1 then `fn` will be executed serially by each + shards + + `max_concurrency` should be used to achieve a high degree of + parallelism without the overhead of increasing the number of shards + (which are actor based). If `max_concurrency` is not 1, this function + provides no semantic guarantees on the output order. + Results will be returned as soon as they are ready. + + A performance note: When executing concurrently, this function + maintains its own internal buffer. If `num_async` is `n` and + max_concur is `k` then the total number of buffered objects could be up + to `n + k - 1` + + Args: + fn: function to apply to each item. + max_concurrency: max number of concurrent calls to fn per + shard. If 0, then apply all operations concurrently. + resources: resources that the function requires to execute. + This has the same default as `ray.remote` and is only used + when `max_concurrency > 1`. + + Returns: + ParallelIterator[U]: a parallel iterator whose elements have `fn` + applied. + + Examples: + >>> next(from_range(4).for_each( + lambda x: x * 2, + max_concur=2, + resources={"num_cpus": 0.1}).gather_sync() + ) + ... [0, 2, 4, 8] + + """ + assert max_concurrency >= 0, "max_concurrency must be non-negative." + return self._with_transform( + lambda local_it: local_it.for_each(fn, max_concurrency, resources), + ".for_each()", + ) + + def filter(self, fn: Callable[[T], bool]) -> "ParallelIterator[T]": + """Remotely filter items from this iterator. + + Args: + fn: returns False for items to drop from the iterator. + + Examples: + >>> it = from_items([0, 1, 2]).filter(lambda x: x > 0) + >>> next(it.gather_sync()) + ... [1, 2] + """ + return self._with_transform(lambda local_it: local_it.filter(fn), ".filter()") + + def batch(self, n: int) -> "ParallelIterator[List[T]]": + """Remotely batch together items in this iterator. + + Args: + n: Number of items to batch together. + + Examples: + >>> next(from_range(10, 1).batch(4).gather_sync()) + ... [0, 1, 2, 3] + """ + return self._with_transform(lambda local_it: local_it.batch(n), f".batch({n})") + + def flatten(self) -> "ParallelIterator[T[0]]": + """Flatten batches of items into individual items. + + Examples: + >>> next(from_range(10, 1).batch(4).flatten()) + ... 0 + """ + return self._with_transform(lambda local_it: local_it.flatten(), ".flatten()") + + def combine(self, fn: Callable[[T], List[U]]) -> "ParallelIterator[U]": + """Transform and then combine items horizontally. + + This is the equivalent of for_each(fn).flatten() (flat map). + """ + it = self.for_each(fn).flatten() + it.name = self.name + ".combine()" + return it + + def local_shuffle( + self, shuffle_buffer_size: int, seed: int = None + ) -> "ParallelIterator[T]": + """Remotely shuffle items of each shard independently + + Args: + shuffle_buffer_size: The algorithm fills a buffer with + shuffle_buffer_size elements and randomly samples elements from + this buffer, replacing the selected elements with new elements. + For perfect shuffling, this argument should be greater than or + equal to the largest iterator size. + seed: Seed to use for + randomness. Default value is None. + + Returns: + A ParallelIterator with a local shuffle applied on the base + iterator + + Examples: + >>> it = from_range(10, 1).local_shuffle(shuffle_buffer_size=2) + >>> it = it.gather_sync() + >>> next(it) + 0 + >>> next(it) + 2 + >>> next(it) + 3 + >>> next(it) + 1 + """ + return self._with_transform( + lambda local_it: local_it.shuffle(shuffle_buffer_size, seed), + ".local_shuffle(shuffle_buffer_size={}, seed={})".format( + shuffle_buffer_size, str(seed) if seed is not None else "None" + ), + ) + + def repartition( + self, num_partitions: int, batch_ms: int = 0 + ) -> "ParallelIterator[T]": + """Returns a new ParallelIterator instance with num_partitions shards. + + The new iterator contains the same data in this instance except with + num_partitions shards. The data is split in round-robin fashion for + the new ParallelIterator. + + Args: + num_partitions: The number of shards to use for the new + ParallelIterator + batch_ms: Batches items for batch_ms milliseconds + on each shard before retrieving it. + Increasing batch_ms increases latency but improves throughput. + + Returns: + A ParallelIterator with num_partitions number of shards and the + data of this ParallelIterator split round-robin among the new + number of shards. + + Examples: + >>> it = from_range(8, 2) + >>> it = it.repartition(3) + >>> list(it.get_shard(0)) + [0, 4, 3, 7] + >>> list(it.get_shard(1)) + [1, 5] + >>> list(it.get_shard(2)) + [2, 6] + """ + + # initialize the local iterators for all the actors + all_actors = [] + for actor_set in self.actor_sets: + actor_set.init_actors() + all_actors.extend(actor_set.actors) + + def base_iterator(num_partitions, partition_index, timeout=None): + futures = {} + for a in all_actors: + futures[ + a.par_iter_slice_batch.remote( + step=num_partitions, start=partition_index, batch_ms=batch_ms + ) + ] = a + while futures: + pending = list(futures) + if timeout is None: + # First try to do a batch wait for efficiency. + ready, _ = ray.wait(pending, num_returns=len(pending), timeout=0) + # Fall back to a blocking wait. + if not ready: + ready, _ = ray.wait(pending, num_returns=1) + else: + ready, _ = ray.wait( + pending, num_returns=len(pending), timeout=timeout + ) + for obj_ref in ready: + actor = futures.pop(obj_ref) + try: + batch = ray.get(obj_ref) + futures[ + actor.par_iter_slice_batch.remote( + step=num_partitions, + start=partition_index, + batch_ms=batch_ms, + ) + ] = actor + for item in batch: + yield item + except StopIteration: + pass + # Always yield after each round of wait with timeout. + if timeout is not None: + yield _NextValueNotReady() + + def make_gen_i(i): + return lambda: base_iterator(num_partitions, i) + + name = self.name + f".repartition[num_partitions={num_partitions}]" + + generators = [make_gen_i(s) for s in range(num_partitions)] + worker_cls = ray.remote(ParallelIteratorWorker) + actors = [worker_cls.remote(g, repeat=False) for g in generators] + # need explicit reference to self so actors in this instance do not die + return ParallelIterator([_ActorSet(actors, [])], name, parent_iterators=[self]) + + def gather_sync(self) -> "LocalIterator[T]": + """Returns a local iterable for synchronous iteration. + + New items will be fetched from the shards on-demand as the iterator + is stepped through. + + This is the equivalent of batch_across_shards().flatten(). + + Examples: + >>> it = from_range(100, 1).gather_sync() + >>> next(it) + ... 0 + >>> next(it) + ... 1 + >>> next(it) + ... 2 + """ + it = self.batch_across_shards().flatten() + it.name = f"{self}.gather_sync()" + return it + + def batch_across_shards(self) -> "LocalIterator[List[T]]": + """Iterate over the results of multiple shards in parallel. + + Examples: + >>> it = from_iterators([range(3), range(3)]) + >>> next(it.batch_across_shards()) + ... [0, 0] + """ + + def base_iterator(timeout=None): + active = [] + for actor_set in self.actor_sets: + actor_set.init_actors() + active.extend(actor_set.actors) + futures = [a.par_iter_next.remote() for a in active] + while active: + try: + yield ray.get(futures, timeout=timeout) + futures = [a.par_iter_next.remote() for a in active] + # Always yield after each round of gets with timeout. + if timeout is not None: + yield _NextValueNotReady() + except TimeoutError: + yield _NextValueNotReady() + except StopIteration: + # Find and remove the actor that produced StopIteration. + results = [] + for a, f in zip(list(active), futures): + try: + results.append(ray.get(f)) + except StopIteration: + active.remove(a) + if results: + yield results + futures = [a.par_iter_next.remote() for a in active] + + name = f"{self}.batch_across_shards()" + return LocalIterator(base_iterator, SharedMetrics(), name=name) + + def gather_async(self, batch_ms=0, num_async=1) -> "LocalIterator[T]": + """Returns a local iterable for asynchronous iteration. + + New items will be fetched from the shards asynchronously as soon as + the previous one is computed. Items arrive in non-deterministic order. + + Arguments: + batch_ms: Batches items for batch_ms milliseconds + on each shard before retrieving it. + Increasing batch_ms increases latency but improves throughput. + If this value is 0, then items are returned immediately. + num_async: The max number of async requests in flight + per actor. Increasing this improves the amount of pipeline + parallelism in the iterator. + + Examples: + >>> it = from_range(100, 1).gather_async() + >>> next(it) + ... 3 + >>> next(it) + ... 0 + >>> next(it) + ... 1 + """ + + if num_async < 1: + raise ValueError("queue depth must be positive") + if batch_ms < 0: + raise ValueError("batch time must be positive") + + # Forward reference to the returned iterator. + local_iter = None + + def base_iterator(timeout=None): + all_actors = [] + for actor_set in self.actor_sets: + actor_set.init_actors() + all_actors.extend(actor_set.actors) + futures = {} + for _ in range(num_async): + for a in all_actors: + futures[a.par_iter_next_batch.remote(batch_ms)] = a + while futures: + pending = list(futures) + if timeout is None: + # First try to do a batch wait for efficiency. + ready, _ = ray.wait(pending, num_returns=len(pending), timeout=0) + # Fall back to a blocking wait. + if not ready: + ready, _ = ray.wait(pending, num_returns=1) + else: + ready, _ = ray.wait( + pending, num_returns=len(pending), timeout=timeout + ) + for obj_ref in ready: + actor = futures.pop(obj_ref) + try: + local_iter.shared_metrics.get().current_actor = actor + batch = ray.get(obj_ref) + futures[actor.par_iter_next_batch.remote(batch_ms)] = actor + for item in batch: + yield item + except StopIteration: + pass + # Always yield after each round of wait with timeout. + if timeout is not None: + yield _NextValueNotReady() + + name = f"{self}.gather_async()" + local_iter = LocalIterator(base_iterator, SharedMetrics(), name=name) + return local_iter + + def take(self, n: int) -> List[T]: + """Return up to the first n items from this iterator.""" + return self.gather_sync().take(n) + + def show(self, n: int = 20): + """Print up to the first n items from this iterator.""" + return self.gather_sync().show(n) + + def union(self, other: "ParallelIterator[T]") -> "ParallelIterator[T]": + """Return an iterator that is the union of this and the other.""" + if not isinstance(other, ParallelIterator): + raise TypeError( + f"other must be of type ParallelIterator, got {type(other)}" + ) + actor_sets = [] + actor_sets.extend(self.actor_sets) + actor_sets.extend(other.actor_sets) + # if one of these iterators is a result of a repartition, we need to + # keep an explicit reference to its parent iterator + return ParallelIterator( + actor_sets, + f"ParallelUnion[{self}, {other}]", + parent_iterators=self.parent_iterators + other.parent_iterators, + ) + + def select_shards(self, shards_to_keep: List[int]) -> "ParallelIterator[T]": + """Return a child iterator that only iterates over given shards. + + It is the user's responsibility to ensure child iterators are operating + over disjoint sub-sets of this iterator's shards. + """ + if len(self.actor_sets) > 1: + raise ValueError("select_shards() is not allowed after union()") + if len(shards_to_keep) == 0: + raise ValueError("at least one shard must be selected") + old_actor_set = self.actor_sets[0] + new_actors = [ + a for (i, a) in enumerate(old_actor_set.actors) if i in shards_to_keep + ] + assert len(new_actors) == len(shards_to_keep), "Invalid actor index" + new_actor_set = _ActorSet(new_actors, old_actor_set.transforms) + return ParallelIterator( + [new_actor_set], + f"{self}.select_shards({len(shards_to_keep)} total)", + parent_iterators=self.parent_iterators, + ) + + def num_shards(self) -> int: + """Return the number of worker actors backing this iterator.""" + return sum(len(a.actors) for a in self.actor_sets) + + def shards(self) -> List["LocalIterator[T]"]: + """Return the list of all shards.""" + return [self.get_shard(i) for i in range(self.num_shards())] + + def get_shard( + self, shard_index: int, batch_ms: int = 0, num_async: int = 1 + ) -> "LocalIterator[T]": + """Return a local iterator for the given shard. + + The iterator is guaranteed to be serializable and can be passed to + remote tasks or actors. + + Arguments: + shard_index: Index of the shard to gather. + batch_ms: Batches items for batch_ms milliseconds + before retrieving it. + Increasing batch_ms increases latency but improves throughput. + If this value is 0, then items are returned immediately. + num_async: The max number of requests in flight. + Increasing this improves the amount of pipeline + parallelism in the iterator. + """ + if num_async < 1: + raise ValueError("num async must be positive") + if batch_ms < 0: + raise ValueError("batch time must be positive") + a, t = None, None + i = shard_index + for actor_set in self.actor_sets: + if i < len(actor_set.actors): + a = actor_set.actors[i] + t = actor_set.transforms + break + else: + i -= len(actor_set.actors) + if a is None: + raise ValueError("Shard index out of range", shard_index, self.num_shards()) + + def base_iterator(timeout=None): + queue = collections.deque() + ray.get(a.par_iter_init.remote(t)) + for _ in range(num_async): + queue.append(a.par_iter_next_batch.remote(batch_ms)) + while True: + try: + batch = ray.get(queue.popleft(), timeout=timeout) + queue.append(a.par_iter_next_batch.remote(batch_ms)) + for item in batch: + yield item + # Always yield after each round of gets with timeout. + if timeout is not None: + yield _NextValueNotReady() + except TimeoutError: + yield _NextValueNotReady() + except StopIteration: + break + + name = self.name + f".shard[{shard_index}]" + return LocalIterator(base_iterator, SharedMetrics(), name=name) + + +@Deprecated +class LocalIterator(Generic[T]): + """An iterator over a single shard of data. + + It implements similar transformations as ParallelIterator[T], but the + transforms will be applied locally and not remotely in parallel. + + This class is **serializable** and can be passed to other remote + tasks and actors. However, it should be read from at most one process at + a time.""" + + # If a function passed to LocalIterator.for_each() has this method, + # we will call it at the beginning of each data fetch call. This can be + # used to measure the underlying wait latency for measurement purposes. + ON_FETCH_START_HOOK_NAME = "_on_fetch_start" + + thread_local = threading.local() + + def __init__( + self, + base_iterator: Callable[[], Iterable[T]], + shared_metrics: SharedMetrics, + local_transforms: List[Callable[[Iterable], Any]] = None, + timeout: int = None, + name=None, + ): + """Create a local iterator (this is an internal function). + + Args: + base_iterator: A function that produces the base iterator. + This is a function so that we can ensure LocalIterator is + serializable. + shared_metrics: Existing metrics context or a new + context. Should be the same for each chained iterator. + local_transforms: A list of transformation functions to be + applied on top of the base iterator. When iteration begins, we + create the base iterator and apply these functions. This lazy + creation ensures LocalIterator is serializable until you start + iterating over it. + timeout: Optional timeout in seconds for this iterator, after + which _NextValueNotReady will be returned. This avoids + blocking. + name: Optional name for this iterator. + """ + assert isinstance(shared_metrics, SharedMetrics) + self.base_iterator = base_iterator + self.built_iterator = None + self.local_transforms = local_transforms or [] + self.shared_metrics = shared_metrics + self.timeout = timeout + self.name = name or "unknown" + + @staticmethod + def get_metrics() -> MetricsContext: + """Return the current metrics context. + + This can only be called within an iterator function.""" + if ( + not hasattr(LocalIterator.thread_local, "metrics") + or LocalIterator.thread_local.metrics is None + ): + raise ValueError("Cannot access context outside an iterator.") + return LocalIterator.thread_local.metrics + + def _build_once(self): + if self.built_iterator is None: + it = iter(self.base_iterator(self.timeout)) + for fn in self.local_transforms: + it = fn(it) + self.built_iterator = it + + @contextmanager + def _metrics_context(self): + self.thread_local.metrics = self.shared_metrics.get() + yield + + def __iter__(self): + self._build_once() + return self.built_iterator + + def __next__(self): + self._build_once() + return next(self.built_iterator) + + def __str__(self): + return repr(self) + + def __repr__(self): + return f"LocalIterator[{self.name}]" + + def transform(self, fn: Callable[[Iterable[T]], Iterable[U]]) -> "LocalIterator[U]": + + # TODO(ekl) can we automatically handle NextValueNotReady here? + def apply_transform(it): + for item in fn(it): + yield item + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_transform], + name=self.name + ".transform()", + ) + + def for_each( + self, fn: Callable[[T], U], max_concurrency=1, resources=None + ) -> "LocalIterator[U]": + if max_concurrency == 1: + + def apply_foreach(it): + for item in it: + if isinstance(item, _NextValueNotReady): + yield item + else: + # Keep retrying the function until it returns a valid + # value. This allows for non-blocking functions. + while True: + with self._metrics_context(): + result = fn(item) + yield result + if not isinstance(result, _NextValueNotReady): + break + + else: + if resources is None: + resources = {} + + def apply_foreach(it): + cur = [] + remote = ray.remote(fn).options(**resources) + remote_fn = remote.remote + for item in it: + if isinstance(item, _NextValueNotReady): + yield item + else: + if max_concurrency and len(cur) >= max_concurrency: + finished, cur = ray.wait(cur) + yield from ray.get(finished) + cur.append(remote_fn(item)) + while cur: + finished, cur = ray.wait(cur) + yield from ray.get(finished) + + if hasattr(fn, LocalIterator.ON_FETCH_START_HOOK_NAME): + unwrapped = apply_foreach + + def add_wait_hooks(it): + it = unwrapped(it) + new_item = True + while True: + # Avoids calling on_fetch_start repeatedly if we are + # yielding _NextValueNotReady. + if new_item: + with self._metrics_context(): + fn._on_fetch_start() + new_item = False + item = next(it) + if not isinstance(item, _NextValueNotReady): + new_item = True + yield item + + apply_foreach = add_wait_hooks + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_foreach], + name=self.name + ".for_each()", + ) + + def filter(self, fn: Callable[[T], bool]) -> "LocalIterator[T]": + def apply_filter(it): + for item in it: + with self._metrics_context(): + if isinstance(item, _NextValueNotReady) or fn(item): + yield item + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_filter], + name=self.name + ".filter()", + ) + + def batch(self, n: int) -> "LocalIterator[List[T]]": + def apply_batch(it): + batch = [] + for item in it: + if isinstance(item, _NextValueNotReady): + yield item + else: + batch.append(item) + if len(batch) >= n: + yield batch + batch = [] + if batch: + yield batch + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_batch], + name=self.name + f".batch({n})", + ) + + def flatten(self) -> "LocalIterator[T[0]]": + def apply_flatten(it): + for item in it: + if isinstance(item, _NextValueNotReady): + yield item + else: + for subitem in item: + yield subitem + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_flatten], + name=self.name + ".flatten()", + ) + + def shuffle(self, shuffle_buffer_size: int, seed: int = None) -> "LocalIterator[T]": + """Shuffle items of this iterator + + Args: + shuffle_buffer_size: The algorithm fills a buffer with + shuffle_buffer_size elements and randomly samples elements from + this buffer, replacing the selected elements with new elements. + For perfect shuffling, this argument should be greater than or + equal to the largest iterator size. + seed: Seed to use for + randomness. Default value is None. + + Returns: + A new LocalIterator with shuffling applied + """ + shuffle_random = random.Random(seed) + + def apply_shuffle(it): + buffer = [] + for item in it: + if isinstance(item, _NextValueNotReady): + yield item + else: + buffer.append(item) + if len(buffer) >= shuffle_buffer_size: + yield buffer.pop(shuffle_random.randint(0, len(buffer) - 1)) + while len(buffer) > 0: + yield buffer.pop(shuffle_random.randint(0, len(buffer) - 1)) + + return LocalIterator( + self.base_iterator, + self.shared_metrics, + self.local_transforms + [apply_shuffle], + name=self.name + + ".shuffle(shuffle_buffer_size={}, seed={})".format( + shuffle_buffer_size, str(seed) if seed is not None else "None" + ), + ) + + def combine(self, fn: Callable[[T], List[U]]) -> "LocalIterator[U]": + it = self.for_each(fn).flatten() + it.name = self.name + ".combine()" + return it + + def zip_with_source_actor(self): + def zip_with_source(item): + metrics = LocalIterator.get_metrics() + if metrics.current_actor is None: + raise ValueError("Could not identify source actor of item") + return metrics.current_actor, item + + it = self.for_each(zip_with_source) + it.name = self.name + ".zip_with_source_actor()" + return it + + def take(self, n: int) -> List[T]: + """Return up to the first n items from this iterator.""" + out = [] + for item in self: + out.append(item) + if len(out) >= n: + break + return out + + def show(self, n: int = 20): + """Print up to the first n items from this iterator.""" + i = 0 + for item in self: + print(item) + i += 1 + if i >= n: + break + + def duplicate(self, n) -> List["LocalIterator[T]"]: + """Copy this iterator `n` times, duplicating the data. + + The child iterators will be prioritized by how much of the parent + stream they have consumed. That is, we will not allow children to fall + behind, since that can cause infinite memory buildup in this operator. + + Returns: + List[LocalIterator[T]]: child iterators that each have a copy + of the data of this iterator. + """ + + if n < 2: + raise ValueError("Number of copies must be >= 2") + + queues = [] + for _ in range(n): + queues.append(collections.deque()) + + def fill_next(timeout): + self.timeout = timeout + item = next(self) + for q in queues: + q.append(item) + + def make_next(i): + def gen(timeout): + while True: + my_len = len(queues[i]) + max_len = max(len(q) for q in queues) + # Yield to let other iterators that have fallen behind + # process more items. + if my_len < max_len: + yield _NextValueNotReady() + else: + if len(queues[i]) == 0: + try: + fill_next(timeout) + except StopIteration: + return + yield queues[i].popleft() + + return gen + + iterators = [] + for i in range(n): + iterators.append( + LocalIterator( + make_next(i), + self.shared_metrics, + [], + name=self.name + f".duplicate[{i}]", + ) + ) + + return iterators + + def union( + self, + *others: "LocalIterator[T]", + deterministic: bool = False, + round_robin_weights: List[float] = None, + ) -> "LocalIterator[T]": + """Return an iterator that is the union of this and the others. + + Args: + deterministic: If deterministic=True, we alternate between + reading from one iterator and the others. Otherwise we return + items from iterators as they become ready. + round_robin_weights: List of weights to use for round robin + mode. For example, [2, 1] will cause the iterator to pull twice + as many items from the first iterator as the second. + [2, 1, "*"] will cause as many items to be pulled as possible + from the third iterator without blocking. This overrides the + deterministic flag. + """ + + for it in others: + if not isinstance(it, LocalIterator): + raise ValueError(f"other must be of type LocalIterator, got {type(it)}") + + active = [] + parent_iters = [self] + list(others) + shared_metrics = SharedMetrics(parents=[p.shared_metrics for p in parent_iters]) + + timeout = None if deterministic else 0 + if round_robin_weights: + if len(round_robin_weights) != len(parent_iters): + raise ValueError( + "Length of round robin weights must equal number of " + "iterators total." + ) + timeouts = [0 if w == "*" else None for w in round_robin_weights] + else: + timeouts = [timeout] * len(parent_iters) + round_robin_weights = [1] * len(parent_iters) + + for i, it in enumerate(parent_iters): + active.append( + LocalIterator( + it.base_iterator, + shared_metrics, + it.local_transforms, + timeout=timeouts[i], + ) + ) + active = list(zip(round_robin_weights, active)) + + def build_union(timeout=None): + while True: + for weight, it in list(active): + if weight == "*": + max_pull = 100 # TOOD(ekl) how to best bound this? + else: + max_pull = _randomized_int_cast(weight) + try: + for _ in range(max_pull): + item = next(it) + if isinstance(item, _NextValueNotReady): + if timeout is not None: + yield item + break + else: + yield item + except StopIteration: + active.remove((weight, it)) + if not active: + break + + return LocalIterator( + build_union, + shared_metrics, + [], + name=f"LocalUnion[{self}, {', '.join(map(str, others))}]", + ) + + +@Deprecated +class ParallelIteratorWorker(object): + """Worker actor for a ParallelIterator. + + Actors that are passed to iter.from_actors() must subclass this interface. + """ + + def __init__(self, item_generator: Any, repeat: bool): + """Create an iterator worker. + + Subclasses must call this init function. + + Args: + item_generator: A Python iterable or lambda function + that produces a generator when called. We allow lambda + functions since the generator itself might not be serializable, + but a lambda that returns it can be. + repeat: Whether to loop over the iterator forever. + """ + + def make_iterator(): + if callable(item_generator): + return item_generator() + else: + return item_generator + + if repeat: + + def cycle(): + while True: + it = iter(make_iterator()) + if it is item_generator: + raise ValueError( + "Cannot iterate over {0} multiple times." + + "Please pass in the base iterable or" + + "lambda: {0} instead.".format(item_generator) + ) + for item in it: + yield item + + self.item_generator = cycle() + else: + self.item_generator = make_iterator() + + self.transforms = [] + self.local_it = None + self.next_ith_buffer = None + + def par_iter_init(self, transforms): + """Implements ParallelIterator worker init.""" + it = LocalIterator(lambda timeout: self.item_generator, SharedMetrics()) + for fn in transforms: + it = fn(it) + assert it is not None, fn + self.local_it = iter(it) + + def par_iter_next(self): + """Implements ParallelIterator worker item fetch.""" + assert self.local_it is not None, "must call par_iter_init()" + return next(self.local_it) + + def par_iter_next_batch(self, batch_ms: int): + """Batches par_iter_next.""" + batch = [] + if batch_ms == 0: + batch.append(self.par_iter_next()) + return batch + t_end = time.time() + (0.001 * batch_ms) + while time.time() < t_end: + try: + batch.append(self.par_iter_next()) + except StopIteration: + if len(batch) == 0: + raise StopIteration + else: + pass + return batch + + def par_iter_slice(self, step: int, start: int): + """Iterates in increments of step starting from start.""" + assert self.local_it is not None, "must call par_iter_init()" + + if self.next_ith_buffer is None: + self.next_ith_buffer = collections.defaultdict(list) + + index_buffer = self.next_ith_buffer[start] + if len(index_buffer) > 0: + return index_buffer.pop(0) + else: + for j in range(step): + try: + val = next(self.local_it) + self.next_ith_buffer[j].append(val) + except StopIteration: + pass + + if not self.next_ith_buffer[start]: + raise StopIteration + + return self.next_ith_buffer[start].pop(0) + + def par_iter_slice_batch(self, step: int, start: int, batch_ms: int): + """Batches par_iter_slice.""" + batch = [] + if batch_ms == 0: + batch.append(self.par_iter_slice(step, start)) + return batch + t_end = time.time() + (0.001 * batch_ms) + while time.time() < t_end: + try: + batch.append(self.par_iter_slice(step, start)) + except StopIteration: + if len(batch) == 0: + raise StopIteration + else: + pass + return batch + + +def _randomized_int_cast(float_value): + base = int(float_value) + remainder = float_value - base + if random.random() < remainder: + base += 1 + return base + + +class _NextValueNotReady(Exception): + """Indicates that a local iterator has no value currently available. + + This is used internally to implement the union() of multiple blocking + local generators.""" + + pass + + +class _ActorSet(object): + """Helper class that represents a set of actors and transforms.""" + + def __init__( + self, + actors: List["ray.actor.ActorHandle"], + transforms: List[Callable[["LocalIterator"], "LocalIterator"]], + ): + self.actors = actors + self.transforms = transforms + + def init_actors(self): + ray.get([a.par_iter_init.remote(self.transforms) for a in self.actors]) + + def with_transform(self, fn): + return _ActorSet(self.actors, self.transforms + [fn]) diff --git a/lib/python3.12/site-packages/ray/util/joblib/__init__.py b/lib/python3.12/site-packages/ray/util/joblib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1b02709aa53b5ba1d52709b6f80334b8e2308594 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/joblib/__init__.py @@ -0,0 +1,20 @@ +from joblib.parallel import register_parallel_backend + + +def register_ray(): + """Register Ray Backend to be called with parallel_backend("ray").""" + try: + from ray.util.joblib.ray_backend import RayBackend + + register_parallel_backend("ray", RayBackend) + except ImportError: + msg = ( + "To use the ray backend you must install ray." + "Try running 'pip install ray'." + "See https://docs.ray.io/en/master/installation.html" + "for more information." + ) + raise ImportError(msg) + + +__all__ = ["register_ray"] diff --git a/lib/python3.12/site-packages/ray/util/joblib/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/joblib/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63eaed6020a9c0d7be0cc8ac31674669e44c2542 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/joblib/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/joblib/__pycache__/ray_backend.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/joblib/__pycache__/ray_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4de003246f37bb4623cc05f02ed710f36617e7c Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/joblib/__pycache__/ray_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/joblib/ray_backend.py b/lib/python3.12/site-packages/ray/util/joblib/ray_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..72cb7032556abd890f31d85f377ab76cff3c6076 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/joblib/ray_backend.py @@ -0,0 +1,96 @@ +import logging +from typing import Any, Dict, Optional + +from joblib import Parallel +from joblib._parallel_backends import MultiprocessingBackend +from joblib.pool import PicklingPool + +import ray +from ray._common.usage import usage_lib +from ray.util.multiprocessing.pool import Pool + +logger = logging.getLogger(__name__) + + +class RayBackend(MultiprocessingBackend): + """Ray backend uses ray, a system for scalable distributed computing. + More info about Ray is available here: https://docs.ray.io. + """ + + def __init__( + self, + nesting_level: Optional[int] = None, + inner_max_num_threads: Optional[int] = None, + ray_remote_args: Optional[Dict[str, Any]] = None, + **kwargs + ): + """``ray_remote_args`` will be used to configure Ray Actors + making up the pool.""" + usage_lib.record_library_usage("util.joblib") + + self.ray_remote_args = ray_remote_args + super().__init__( + nesting_level=nesting_level, + inner_max_num_threads=inner_max_num_threads, + **kwargs + ) + + # ray_remote_args is used both in __init__ and configure to allow for it to be + # set in both `parallel_backend` and `Parallel` respectively + + def configure( + self, + n_jobs: int = 1, + parallel: Optional[Parallel] = None, + prefer: Optional[str] = None, + require: Optional[str] = None, + ray_remote_args: Optional[Dict[str, Any]] = None, + **memmappingpool_args + ): + """Make Ray Pool the father class of PicklingPool. PicklingPool is a + father class that inherits Pool from multiprocessing.pool. The next + line is a patch, which changes the inheritance of Pool to be from + ray.util.multiprocessing.pool. + + ``ray_remote_args`` will be used to configure Ray Actors making up the pool. + This will override ``ray_remote_args`` set during initialization. + """ + PicklingPool.__bases__ = (Pool,) + """Use all available resources when n_jobs == -1. Must set RAY_ADDRESS + variable in the environment or run ray.init(address=..) to run on + multiple nodes. + """ + if n_jobs == -1: + if not ray.is_initialized(): + import os + + if "RAY_ADDRESS" in os.environ: + logger.info( + "Connecting to ray cluster at address='{}'".format( + os.environ["RAY_ADDRESS"] + ) + ) + else: + logger.info("Starting local ray cluster") + ray.init() + ray_cpus = int(ray._private.state.cluster_resources()["CPU"]) + n_jobs = ray_cpus + + eff_n_jobs = super(RayBackend, self).configure( + n_jobs, + parallel, + prefer, + require, + ray_remote_args=ray_remote_args + if ray_remote_args is not None + else self.ray_remote_args, + **memmappingpool_args + ) + return eff_n_jobs + + def effective_n_jobs(self, n_jobs): + eff_n_jobs = super(RayBackend, self).effective_n_jobs(n_jobs) + if n_jobs == -1: + ray_cpus = int(ray._private.state.cluster_resources()["CPU"]) + eff_n_jobs = ray_cpus + return eff_n_jobs diff --git a/lib/python3.12/site-packages/ray/util/metrics.py b/lib/python3.12/site-packages/ray/util/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..51c8082f6e2fe05abb45c87233e9a7ba1fbd2330 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/metrics.py @@ -0,0 +1,372 @@ +import logging +import re +import warnings +from typing import Any, Dict, List, Optional, Tuple, Union + +from ray._private.ray_constants import env_bool +from ray._raylet import ( + Count as CythonCount, + Gauge as CythonGauge, + Histogram as CythonHistogram, + Sum as CythonSum, +) # noqa: E402 + +# Sum is used for CythonCount because it allows incrementing by positive +# values that are different from one. +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + +# Copied from Prometheus Python Client. While the regex is not part of the public API +# for Prometheus, it's not expected to change. +# https://github.com/prometheus/client_python/blob/46eae7bae88f76951f7246d9f359f2dd5eeff110/prometheus_client/validation.py#L4 +_VALID_METRIC_NAME_RE = re.compile(r"^[a-zA-Z_:][a-zA-Z0-9_:]*$") + + +def _is_invalid_metric_name(name: str) -> bool: + if len(name) == 0: + raise ValueError("Empty name is not allowed. Please provide a metric name.") + if not _VALID_METRIC_NAME_RE.match(name): + warnings.warn( + f"Invalid metric name: {name}. Metric will be discarded " + "and data will not be collected or published. " + "Metric names can only contain letters, numbers, _, and :. " + "Metric names cannot start with numbers.", + UserWarning, + ) + return True + return False + + +@DeveloperAPI +class Metric: + """The parent class of custom metrics. + + Ray's custom metrics APIs are rooted from this class and share + the same public methods. + """ + + def __init__( + self, + name: str, + description: str = "", + tag_keys: Optional[Tuple[str, ...]] = None, + ): + # Metrics with invalid names will be discarded and will not be collected + # by Prometheus. + self._discard_metric = _is_invalid_metric_name(name) + self._name = name + self._description = description + # The default tags key-value pair. + self._default_tags = {} + # Keys of tags. + self._tag_keys = tag_keys or tuple() + # The Cython metric class. This should be set in the child class. + self._metric = None + + if not isinstance(self._tag_keys, tuple): + raise TypeError( + "tag_keys should be a tuple type, got: " f"{type(self._tag_keys)}" + ) + + for key in self._tag_keys: + if not isinstance(key, str): + raise TypeError(f"Tag keys must be str, got {type(key)}.") + + if ":" in self._name: + warnings.warn( + f"Metric name {self._name} contains a : character, which is no longer allowed. " + f"Please migrate to the new metric name format. " + f"This will be an error in the future.", + FutureWarning, + ) + + def set_default_tags(self, default_tags: Dict[str, str]): + """Set default tags of metrics. + + Example: + >>> from ray.util.metrics import Counter + >>> # Note that set_default_tags returns the instance itself. + >>> counter = Counter("name", tag_keys=("a",)) + >>> counter2 = counter.set_default_tags({"a": "b"}) + >>> assert counter is counter2 + >>> # this means you can instantiate it in this way. + >>> counter = Counter("name", tag_keys=("a",)).set_default_tags({"a": "b"}) + + Args: + default_tags: Default tags that are + used for every record method. + + Returns: + Metric: it returns the instance itself. + """ + for key, val in default_tags.items(): + if key not in self._tag_keys: + raise ValueError(f"Unrecognized tag key {key}.") + if not isinstance(val, str): + raise TypeError(f"Tag values must be str, got {type(val)}.") + + self._default_tags = default_tags + return self + + def _record( + self, + value: Union[int, float], + tags: Optional[Dict[str, str]] = None, + ) -> None: + """Record the metric point of the metric. + + Tags passed in will take precedence over the metric's default tags. + + Args: + value: The value to be recorded as a metric point. + """ + if self._discard_metric: + return + + assert self._metric is not None + + final_tags = self._get_final_tags(tags) + self._validate_tags(final_tags) + self._metric.record(value, tags=final_tags) + + def _get_final_tags(self, tags): + if not tags: + return self._default_tags + + for val in tags.values(): + if not isinstance(val, str): + raise TypeError(f"Tag values must be str, got {type(val)}.") + + return {**self._default_tags, **tags} + + def _validate_tags(self, final_tags): + missing_tags = [] + for tag_key in self._tag_keys: + # Prefer passed tags over default tags. + if tag_key not in final_tags: + missing_tags.append(tag_key) + + # Strict validation: if any required tag_keys are missing, raise error + if missing_tags: + raise ValueError(f"Missing value for tag key(s): {','.join(missing_tags)}.") + + @property + def info(self) -> Dict[str, Any]: + """Return the information of this metric. + + Example: + >>> from ray.util.metrics import Counter + >>> counter = Counter("name", description="desc") + >>> print(counter.info) + {'name': 'name', 'description': 'desc', 'tag_keys': (), 'default_tags': {}} + """ + return { + "name": self._name, + "description": self._description, + "tag_keys": self._tag_keys, + "default_tags": self._default_tags, + } + + +@DeveloperAPI +class Counter(Metric): + """A cumulative metric that is monotonically increasing. + + This corresponds to Prometheus' counter metric: + https://prometheus.io/docs/concepts/metric_types/#counter + + Before Ray 2.10, this exports a Prometheus gauge metric instead of + a counter metric, which is wrong. + Since 2.10, this exports both counter (with a suffix "_total") and + gauge metrics (for bug compatibility). + Use `RAY_EXPORT_COUNTER_AS_GAUGE=0` to disable exporting the gauge metric. + + Args: + name: Name of the metric. + description: Description of the metric. + tag_keys: Tag keys of the metric. + """ + + def __init__( + self, + name: str, + description: str = "", + tag_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__(name, description, tag_keys) + if self._discard_metric: + self._metric = None + else: + if env_bool("RAY_enable_open_telemetry", False): + """ + For the previous opencensus implementation, we used Sum to support + exporting Counter as a gauge metric. We'll drop that feature in the + new opentelemetry implementation. + """ + self._metric = CythonSum(self._name, self._description, self._tag_keys) + else: + """ + For the new opentelemetry implementation, we'll correctly use Counter + rather than Sum. + """ + self._metric = CythonCount( + self._name, self._description, self._tag_keys + ) + + def __reduce__(self): + deserializer = self.__class__ + serialized_data = (self._name, self._description, self._tag_keys) + return deserializer, serialized_data + + def inc(self, value: Union[int, float] = 1.0, tags: Dict[str, str] = None): + """Increment the counter by `value` (defaults to 1). + + Tags passed in will take precedence over the metric's default tags. + + Args: + value(int, float): Value to increment the counter by (default=1). + tags(Dict[str, str]): Tags to set or override for this counter. + """ + if not isinstance(value, (int, float)): + raise TypeError(f"value must be int or float, got {type(value)}.") + if value <= 0: + raise ValueError(f"value must be >0, got {value}") + + self._record(value, tags=tags) + + +@DeveloperAPI +class Histogram(Metric): + """Tracks the size and number of events in buckets. + + Histograms allow you to calculate aggregate quantiles + such as 25, 50, 95, 99 percentile latency for an RPC. + + This corresponds to Prometheus' histogram metric: + https://prometheus.io/docs/concepts/metric_types/#histogram + + Args: + name: Name of the metric. + description: Description of the metric. + boundaries: Boundaries of histogram buckets. + tag_keys: Tag keys of the metric. + """ + + def __init__( + self, + name: str, + description: str = "", + boundaries: List[float] = None, + tag_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__(name, description, tag_keys) + if boundaries is None or len(boundaries) == 0: + raise ValueError( + "boundaries argument should be provided when using " + "the Histogram class. e.g., " + 'Histogram("name", boundaries=[1.0, 2.0])' + ) + for i, boundary in enumerate(boundaries): + if boundary <= 0: + raise ValueError( + "Invalid `boundaries` argument at index " + f"{i}, {boundaries}. Use positive values for the arguments." + ) + + self.boundaries = boundaries + if self._discard_metric: + self._metric = None + else: + self._metric = CythonHistogram( + self._name, self._description, self.boundaries, self._tag_keys + ) + + def observe(self, value: Union[int, float], tags: Dict[str, str] = None): + """Observe a given `value` and add it to the appropriate bucket. + + Tags passed in will take precedence over the metric's default tags. + + Args: + value(int, float): Value to set the gauge to. + tags(Dict[str, str]): Tags to set or override for this gauge. + """ + if not isinstance(value, (int, float)): + raise TypeError(f"value must be int or float, got {type(value)}.") + + self._record(value, tags) + + def __reduce__(self): + deserializer = Histogram + serialized_data = ( + self._name, + self._description, + self.boundaries, + self._tag_keys, + ) + return deserializer, serialized_data + + @property + def info(self): + """Return information about histogram metric.""" + info = super().info + info.update({"boundaries": self.boundaries}) + return info + + +@DeveloperAPI +class Gauge(Metric): + """Gauges keep the last recorded value and drop everything before. + + Unlike counters, gauges can go up or down over time. + + This corresponds to Prometheus' gauge metric: + https://prometheus.io/docs/concepts/metric_types/#gauge + + Args: + name: Name of the metric. + description: Description of the metric. + tag_keys: Tag keys of the metric. + """ + + def __init__( + self, + name: str, + description: str = "", + tag_keys: Optional[Tuple[str, ...]] = None, + ): + super().__init__(name, description, tag_keys) + if self._discard_metric: + self._metric = None + else: + self._metric = CythonGauge(self._name, self._description, self._tag_keys) + + def set(self, value: Optional[Union[int, float]], tags: Dict[str, str] = None): + """Set the gauge to the given `value`. + + Tags passed in will take precedence over the metric's default tags. + + Args: + value(int, float): Value to set the gauge to. If `None`, this method is a + no-op. + tags(Dict[str, str]): Tags to set or override for this gauge. + """ + if value is None: + return + + if not isinstance(value, (int, float)): + raise TypeError(f"value must be int or float, got {type(value)}.") + + self._record(value, tags) + + def __reduce__(self): + deserializer = Gauge + serialized_data = (self._name, self._description, self._tag_keys) + return deserializer, serialized_data + + +__all__ = [ + "Counter", + "Histogram", + "Gauge", +] diff --git a/lib/python3.12/site-packages/ray/util/rpdb.py b/lib/python3.12/site-packages/ray/util/rpdb.py new file mode 100644 index 0000000000000000000000000000000000000000..9610e070cbf9f55bbcac79f7e4f92223cef88121 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/rpdb.py @@ -0,0 +1,384 @@ +# Some code in this file is from +# https://github.com/ionelmc/python-remote-pdb/blob/07d563331c4ab9eb45731bb272b158816d98236e/src/remote_pdb.py +# (BSD 2-Clause "Simplified" License) + +import errno +import inspect +import json +import logging +import os +import re +import select +import socket +import sys +import time +import traceback +import uuid +from pdb import Pdb +from typing import Callable + +import ray +from ray._common.network_utils import build_address, is_ipv6 +from ray._private import ray_constants +from ray.experimental.internal_kv import _internal_kv_del, _internal_kv_put +from ray.util.annotations import DeveloperAPI + +log = logging.getLogger(__name__) + + +def _cry(message, stderr=sys.__stderr__): + print(message, file=stderr) + stderr.flush() + + +class _LF2CRLF_FileWrapper(object): + def __init__(self, connection): + self.connection = connection + self.stream = fh = connection.makefile("rw") + self.read = fh.read + self.readline = fh.readline + self.readlines = fh.readlines + self.close = fh.close + self.flush = fh.flush + self.fileno = fh.fileno + if hasattr(fh, "encoding"): + self._send = lambda data: connection.sendall( + data.encode(fh.encoding, errors="replace") + ) + else: + self._send = connection.sendall + + @property + def encoding(self): + return self.stream.encoding + + def __iter__(self): + return self.stream.__iter__() + + def write(self, data, nl_rex=re.compile("\r?\n")): + data = nl_rex.sub("\r\n", data) + self._send(data) + + def writelines(self, lines, nl_rex=re.compile("\r?\n")): + for line in lines: + self.write(line, nl_rex) + + +class _PdbWrap(Pdb): + """Wrap PDB to run a custom exit hook on continue.""" + + def __init__(self, exit_hook: Callable[[], None]): + self._exit_hook = exit_hook + Pdb.__init__(self) + + def do_continue(self, arg): + self._exit_hook() + return Pdb.do_continue(self, arg) + + do_c = do_cont = do_continue + + +class _RemotePdb(Pdb): + """ + This will run pdb as a ephemeral telnet service. Once you connect no one + else can connect. On construction this object will block execution till a + client has connected. + Based on https://github.com/tamentis/rpdb I think ... + To use this:: + RemotePdb(host="0.0.0.0", port=4444).set_trace() + Then run: telnet 127.0.0.1 4444 + """ + + active_instance = None + + def __init__( + self, + breakpoint_uuid, + host, + port, + ip_address, + patch_stdstreams=False, + quiet=False, + ): + self._breakpoint_uuid = breakpoint_uuid + self._quiet = quiet + self._patch_stdstreams = patch_stdstreams + self._listen_socket = socket.socket( + socket.AF_INET6 if is_ipv6(host) else socket.AF_INET, socket.SOCK_STREAM + ) + self._listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True) + self._listen_socket.bind((host, port)) + self._ip_address = ip_address + + def listen(self): + if not self._quiet: + _cry( + "RemotePdb session open at %s, " + "use 'ray debug' to connect..." + % build_address(self._ip_address, self._listen_socket.getsockname()[1]) + ) + self._listen_socket.listen(1) + connection, address = self._listen_socket.accept() + if not self._quiet: + _cry(f"RemotePdb accepted connection from {address}") + self.handle = _LF2CRLF_FileWrapper(connection) + Pdb.__init__( + self, + completekey="tab", + stdin=self.handle, + stdout=self.handle, + skip=["ray.*"], + ) + self.backup = [] + if self._patch_stdstreams: + for name in ( + "stderr", + "stdout", + "__stderr__", + "__stdout__", + "stdin", + "__stdin__", + ): + self.backup.append((name, getattr(sys, name))) + setattr(sys, name, self.handle) + _RemotePdb.active_instance = self + + def __restore(self): + if self.backup and not self._quiet: + _cry("Restoring streams: %s ..." % self.backup) + for name, fh in self.backup: + setattr(sys, name, fh) + self.handle.close() + _RemotePdb.active_instance = None + + def do_quit(self, arg): + self.__restore() + return Pdb.do_quit(self, arg) + + do_q = do_exit = do_quit + + def do_continue(self, arg): + self.__restore() + self.handle.connection.close() + return Pdb.do_continue(self, arg) + + do_c = do_cont = do_continue + + def set_trace(self, frame=None): + if frame is None: + frame = sys._getframe().f_back + try: + Pdb.set_trace(self, frame) + except IOError as exc: + if exc.errno != errno.ECONNRESET: + raise + + def post_mortem(self, traceback=None): + # See https://github.com/python/cpython/blob/ + # 022bc7572f061e1d1132a4db9d085b29707701e7/Lib/pdb.py#L1617 + try: + t = sys.exc_info()[2] + self.reset() + Pdb.interaction(self, None, t) + except IOError as exc: + if exc.errno != errno.ECONNRESET: + raise + + def do_remote(self, arg): + """remote + Skip into the next remote call. + """ + # Tell the next task to drop into the debugger. + ray._private.worker.global_worker.debugger_breakpoint = self._breakpoint_uuid + # Tell the debug loop to connect to the next task. + data = json.dumps( + { + "job_id": ray.get_runtime_context().get_job_id(), + } + ) + _internal_kv_put( + "RAY_PDB_CONTINUE_{}".format(self._breakpoint_uuid), + data, + namespace=ray_constants.KV_NAMESPACE_PDB, + ) + self.__restore() + self.handle.connection.close() + return Pdb.do_continue(self, arg) + + def do_get(self, arg): + """get + Skip to where the current task returns to. + """ + ray._private.worker.global_worker.debugger_get_breakpoint = ( + self._breakpoint_uuid + ) + self.__restore() + self.handle.connection.close() + return Pdb.do_continue(self, arg) + + +def _connect_ray_pdb( + host=None, + port=None, + patch_stdstreams=False, + quiet=None, + breakpoint_uuid=None, + debugger_external=False, +): + """ + Opens a remote PDB on first available port. + """ + if debugger_external: + assert not host, "Cannot specify both host and debugger_external" + host = "0.0.0.0" + elif host is None: + host = os.environ.get("REMOTE_PDB_HOST", "127.0.0.1") + if port is None: + port = int(os.environ.get("REMOTE_PDB_PORT", "0")) + if quiet is None: + quiet = bool(os.environ.get("REMOTE_PDB_QUIET", "")) + if not breakpoint_uuid: + breakpoint_uuid = uuid.uuid4().hex + if debugger_external: + ip_address = ray._private.worker.global_worker.node_ip_address + else: + ip_address = "localhost" + rdb = _RemotePdb( + breakpoint_uuid=breakpoint_uuid, + host=host, + port=port, + ip_address=ip_address, + patch_stdstreams=patch_stdstreams, + quiet=quiet, + ) + sockname = rdb._listen_socket.getsockname() + pdb_address = build_address(ip_address, sockname[1]) + parentframeinfo = inspect.getouterframes(inspect.currentframe())[2] + data = { + "proctitle": ray._raylet.getproctitle(), + "pdb_address": pdb_address, + "filename": parentframeinfo.filename, + "lineno": parentframeinfo.lineno, + "traceback": "\n".join(traceback.format_exception(*sys.exc_info())), + "timestamp": time.time(), + "job_id": ray.get_runtime_context().get_job_id(), + "node_id": ray.get_runtime_context().get_node_id(), + "worker_id": ray.get_runtime_context().get_worker_id(), + "actor_id": ray.get_runtime_context().get_actor_id(), + "task_id": ray.get_runtime_context().get_task_id(), + } + _internal_kv_put( + "RAY_PDB_{}".format(breakpoint_uuid), + json.dumps(data), + overwrite=True, + namespace=ray_constants.KV_NAMESPACE_PDB, + ) + rdb.listen() + _internal_kv_del( + "RAY_PDB_{}".format(breakpoint_uuid), namespace=ray_constants.KV_NAMESPACE_PDB + ) + + return rdb + + +@DeveloperAPI +def set_trace(breakpoint_uuid=None): + """Interrupt the flow of the program and drop into the Ray debugger. + + Can be used within a Ray task or actor. + """ + if os.environ.get("RAY_DEBUG", "1") == "1": + return ray.util.ray_debugpy.set_trace(breakpoint_uuid) + if os.environ.get("RAY_DEBUG", "1") == "legacy": + # If there is an active debugger already, we do not want to + # start another one, so "set_trace" is just a no-op in that case. + if ray._private.worker.global_worker.debugger_breakpoint == b"": + frame = sys._getframe().f_back + rdb = _connect_ray_pdb( + host=None, + port=None, + patch_stdstreams=False, + quiet=None, + breakpoint_uuid=breakpoint_uuid.decode() if breakpoint_uuid else None, + debugger_external=ray._private.worker.global_worker.ray_debugger_external, # noqa: E501 + ) + rdb.set_trace(frame=frame) + + +def _driver_set_trace(): + """The breakpoint hook to use for the driver. + + This disables Ray driver logs temporarily so that the PDB console is not + spammed: https://github.com/ray-project/ray/issues/18172 + """ + if os.environ.get("RAY_DEBUG", "1") == "1": + return ray.util.ray_debugpy.set_trace() + if os.environ.get("RAY_DEBUG", "1") == "legacy": + print("*** Temporarily disabling Ray worker logs ***") + ray._private.worker._worker_logs_enabled = False + + def enable_logging(): + print("*** Re-enabling Ray worker logs ***") + ray._private.worker._worker_logs_enabled = True + + pdb = _PdbWrap(enable_logging) + frame = sys._getframe().f_back + pdb.set_trace(frame) + + +def _is_ray_debugger_post_mortem_enabled(): + return os.environ.get("RAY_DEBUG_POST_MORTEM", "0") == "1" + + +def _post_mortem(): + if os.environ.get("RAY_DEBUG", "1") == "1": + return ray.util.ray_debugpy._post_mortem() + + rdb = _connect_ray_pdb( + host=None, + port=None, + patch_stdstreams=False, + quiet=None, + debugger_external=ray._private.worker.global_worker.ray_debugger_external, + ) + rdb.post_mortem() + + +def _connect_pdb_client(host, port): + if sys.platform == "win32": + import msvcrt + + s = socket.socket( + socket.AF_INET6 if is_ipv6(host) else socket.AF_INET, socket.SOCK_STREAM + ) + s.connect((host, port)) + + while True: + # Get the list of sockets which are readable. + if sys.platform == "win32": + ready_to_read = select.select([s], [], [], 1)[0] + if msvcrt.kbhit(): + ready_to_read.append(sys.stdin) + if not ready_to_read and not sys.stdin.isatty(): + # in tests, when using pexpect, the pipe makes + # the msvcrt.kbhit() trick fail. Assume we are waiting + # for stdin, since this will block waiting for input + ready_to_read.append(sys.stdin) + else: + ready_to_read, write_sockets, error_sockets = select.select( + [sys.stdin, s], [], [] + ) + + for sock in ready_to_read: + if sock == s: + # Incoming message from remote debugger. + data = sock.recv(4096) + if not data: + return + else: + sys.stdout.write(data.decode()) + sys.stdout.flush() + else: + # User entered a message. + msg = sys.stdin.readline() + s.send(msg.encode()) diff --git a/lib/python3.12/site-packages/ray/util/sgd/__init__.py b/lib/python3.12/site-packages/ray/util/sgd/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6b55229eba31a7c4779f324771ff7c90e584c3a --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/sgd/__init__.py @@ -0,0 +1,4 @@ +raise DeprecationWarning( + "Ray SGD has been deprecated as of Ray 1.13. For distributed " + "deep learning on Ray please use Ray Train instead." +) diff --git a/lib/python3.12/site-packages/ray/util/sgd/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/sgd/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91c06cc210a1d8dfe5cae1006ffee4adb0def2e7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/sgd/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/xgboost/__init__.py b/lib/python3.12/site-packages/ray/util/xgboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ce2f204c260fcef980e6505804e72977d5a5a1cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/xgboost/__init__.py @@ -0,0 +1,4 @@ +raise DeprecationWarning( + "ray.util.xgboost has been removed as of Ray 2.0. Instead, use the `xgboost-ray` " + "library directly or the `XGBoostTrainer` in Ray Train." +) diff --git a/lib/python3.12/site-packages/ray/util/xgboost/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/xgboost/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..abedb52dd65124238fa9c10e32a21619713858ff Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/xgboost/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/_C/__init__.pyi b/lib/python3.12/site-packages/tensordict/_C/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..7ca445588fe7f8c99eb875ba77ee636350afcf85 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_C/__init__.pyi @@ -0,0 +1,9 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +def unravel_key_list(list_of_keys): ... +def unravel_keys(keys): ... +def unravel_key(key): ... +def _unravel_key_to_tuple(key): ... diff --git a/lib/python3.12/site-packages/tensordict/__init__.py b/lib/python3.12/site-packages/tensordict/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b2d2423380387ceca217d11f29e5752b4c3786c9 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/__init__.py @@ -0,0 +1,173 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import tensordict._reductions +from tensordict._lazy import LazyStackedTensorDict +from tensordict._nestedkey import NestedKey +from tensordict._td import ( + cat, + from_consolidated, + from_module, + from_modules, + from_pytree, + fromkeys, + is_tensor_collection, + lazy_stack, + load, + load_memmap, + maybe_dense_stack, + memmap, + save, + stack, + TensorDict, +) +from tensordict._unbatched import UnbatchedTensor + +from tensordict.base import ( + _default_is_leaf as default_is_leaf, + _is_leaf_nontensor as is_leaf_nontensor, + from_any, + from_dict, + from_h5, + from_namedtuple, + from_struct_array, + from_tuple, + get_defaults_to_none, + set_get_defaults_to_none, + TensorDictBase, +) +from tensordict.functional import ( + dense_stack_tds, + make_tensordict, + merge_tensordicts, + pad, + pad_sequence, +) +from tensordict.memmap import MemoryMappedTensor +from tensordict.persistent import PersistentTensorDict +from tensordict.tensorclass import ( + from_dataclass, + MetaData, + NonTensorData, + NonTensorDataBase, + NonTensorStack, + tensorclass, + TensorClass, +) +from tensordict.utils import ( + assert_allclose_td, + assert_close, + capture_non_tensor_stack, + is_batchedtensor, + is_non_tensor, + is_tensorclass, + lazy_legacy, + list_to_stack, + parse_tensor_dict_string, + set_capture_non_tensor_stack, + set_lazy_legacy, + set_list_to_stack, + unravel_key, + unravel_key_list, +) +from tensordict._pytree import * +from tensordict.nn import ( + as_tensordict_module, + TensorClassModuleBase, + TensorClassModuleWrapper, + TensorDictParams, +) + +__version__ = None # type: ignore +try: + try: + from importlib.metadata import version as _dist_version + except ImportError: # pragma: no cover + from importlib_metadata import version as _dist_version # type: ignore + + __version__ = _dist_version("tensordict") +except Exception: + try: + from tensordict._version import ( + __version__, + ) # @manual=//pytorch/tensordict:version + except ImportError: + __version__ = None # type: ignore + +__all__ = [ + # Core classes + "TensorDict", + "TensorDictBase", + "LazyStackedTensorDict", + "UnbatchedTensor", + "TensorClass", + "MemoryMappedTensor", + "PersistentTensorDict", + "NestedKey", + # Factory functions + "from_dict", + "from_any", + "from_h5", + "from_namedtuple", + "from_struct_array", + "from_tuple", + "from_dataclass", + "fromkeys", + "from_module", + "from_modules", + "from_pytree", + "from_consolidated", + "make_tensordict", + # Stacking and concatenation + "stack", + "cat", + "lazy_stack", + "maybe_dense_stack", + "dense_stack_tds", + # Memory mapping + "memmap", + "load_memmap", + # Saving and loading + "save", + "load", + # Merging and padding + "merge_tensordicts", + "pad", + "pad_sequence", + # Utility functions + "is_tensor_collection", + "is_batchedtensor", + "is_non_tensor", + "is_tensorclass", + "assert_close", + "assert_allclose_td", + "unravel_key", + "unravel_key_list", + "parse_tensor_dict_string", + # Configuration + "default_is_leaf", + "is_leaf_nontensor", + "get_defaults_to_none", + "set_get_defaults_to_none", + "capture_non_tensor_stack", + "set_capture_non_tensor_stack", + "lazy_legacy", + "set_lazy_legacy", + "list_to_stack", + "set_list_to_stack", + # TensorClass components + "tensorclass", + "MetaData", + "NonTensorData", + "NonTensorDataBase", + "NonTensorStack", + # NN imports + "as_tensordict_module", + "TensorClassModuleBase", + "TensorClassModuleWrapper", + "TensorDictParams", + # Version + "__version__", +] diff --git a/lib/python3.12/site-packages/tensordict/__init__.pyi b/lib/python3.12/site-packages/tensordict/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..933f390b55a5c9f2609770d6c9aa14673a625dea --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/__init__.pyi @@ -0,0 +1,147 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from tensordict._lazy import LazyStackedTensorDict +from tensordict._nestedkey import NestedKey +from tensordict._td import ( + cat, + from_consolidated, + from_module, + from_modules, + from_pytree, + fromkeys, + is_tensor_collection, + lazy_stack, + load, + load_memmap, + maybe_dense_stack, + memmap, + save, + stack, + TensorDict, +) +from tensordict._unbatched import UnbatchedTensor +from tensordict.base import ( + _default_is_leaf as default_is_leaf, + _is_leaf_nontensor as is_leaf_nontensor, + from_any, + from_dict, + from_h5, + from_namedtuple, + from_struct_array, + from_tuple, + get_defaults_to_none, + set_get_defaults_to_none, + TensorDictBase, +) +from tensordict.functional import ( + dense_stack_tds, + make_tensordict, + merge_tensordicts, + pad, + pad_sequence, +) +from tensordict.memmap import MemoryMappedTensor +from tensordict.nn import as_tensordict_module, TensorDictParams +from tensordict.persistent import PersistentTensorDict +from tensordict.tensorclass import ( + from_dataclass, + MetaData, + NonTensorData, + NonTensorDataBase, + NonTensorStack, + TensorClass, + tensorclass, +) +from tensordict.utils import ( + assert_allclose_td, + assert_close, + capture_non_tensor_stack, + is_batchedtensor, + is_non_tensor, + is_tensorclass, + lazy_legacy, + list_to_stack, + parse_tensor_dict_string, + set_capture_non_tensor_stack, + set_lazy_legacy, + set_list_to_stack, + unravel_key, + unravel_key_list, +) + +__version__: str | None + +__all__ = [ + # Core classes + "TensorDict", + "TensorDictBase", + "LazyStackedTensorDict", + "UnbatchedTensor", + "TensorClass", + "MemoryMappedTensor", + "PersistentTensorDict", + "NestedKey", + # Factory functions + "from_dict", + "from_any", + "from_h5", + "from_namedtuple", + "from_struct_array", + "from_tuple", + "from_dataclass", + "fromkeys", + "from_module", + "from_modules", + "from_pytree", + "from_consolidated", + "make_tensordict", + # Stacking and concatenation + "stack", + "cat", + "lazy_stack", + "maybe_dense_stack", + "dense_stack_tds", + # Memory mapping + "memmap", + "load_memmap", + # Saving and loading + "save", + "load", + # Merging and padding + "merge_tensordicts", + "pad", + "pad_sequence", + # Utility functions + "is_tensor_collection", + "is_batchedtensor", + "is_non_tensor", + "is_tensorclass", + "assert_close", + "assert_allclose_td", + "unravel_key", + "unravel_key_list", + "parse_tensor_dict_string", + # Configuration + "default_is_leaf", + "is_leaf_nontensor", + "get_defaults_to_none", + "set_get_defaults_to_none", + "capture_non_tensor_stack", + "set_capture_non_tensor_stack", + "lazy_legacy", + "set_lazy_legacy", + "list_to_stack", + "set_list_to_stack", + # TensorClass components + "tensorclass", + "MetaData", + "NonTensorData", + "NonTensorDataBase", + "NonTensorStack", + # NN imports + "as_tensordict_module", + "TensorDictParams", +] diff --git a/lib/python3.12/site-packages/tensordict/_contextlib.py b/lib/python3.12/site-packages/tensordict/_contextlib.py new file mode 100644 index 0000000000000000000000000000000000000000..e1d119e56b1dec368c10a20b9b3d36c44541bc9d --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_contextlib.py @@ -0,0 +1,461 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import contextlib + +# This is a copy from https://github.com/pytorch/pytorch/blob/main/torch/utils/_contextlib.py#L120 +# We use it for compatibility with torch >= 1.10 where the implementation fails +# for some tests in torchrl. + +# Extra utilities for working with context managers that should have been +# in the standard library but are not + +import functools +import inspect +import sys +from typing import Any, Callable, cast, TypeVar + +import numpy as np + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + + +# Used for annotating the decorator usage of _DecoratorContextManager (e.g., +# 'no_grad' and 'enable_grad'). +# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators +FuncType = Callable[..., Any] +F = TypeVar("F", bound=FuncType) + + +def _wrap_generator(ctx_factory, func): + """Wrap each generator invocation with the context manager factory. + + The input should be a function that returns a context manager, + not a context manager itself, to handle one-shot context managers. + """ + + @functools.wraps(func) + def generator_context(*args, **kwargs): + gen = func(*args, **kwargs) + + # Generators are suspended and unsuspended at `yield`, hence we + # make sure the grad mode is properly set every time the execution + # flow returns into the wrapped generator and restored when it + # returns through our `yield` to our caller (see PR #49017). + try: + # Issuing `None` to a generator fires it up + with ctx_factory(): + response = gen.send(None) + + while True: + try: + # Forward the response to our caller and get its next request + request = yield response + + except GeneratorExit: + # Inform the still active generator about its imminent closure + with ctx_factory(): + gen.close() + raise + + except BaseException: + # Propagate the exception thrown at us by the caller + with ctx_factory(): + response = gen.throw(*sys.exc_info()) + + else: + # Pass the last request to the generator and get its response + with ctx_factory(): + response = gen.send(request) + + # We let the exceptions raised above by the generator's `.throw` or + # `.send` methods bubble up to our caller, except for StopIteration + except StopIteration as e: + # The generator informed us that it is done: take whatever its + # returned value (if any) was and indicate that we're done too + # by returning it (see docs for python's return-statement). + return e.value + + return generator_context + + +def context_decorator(ctx, func): + """Like contextlib.ContextDecorator. + + Except: + + 1. Is done by wrapping, rather than inheritance, so it works with context + managers that are implemented from C and thus cannot easily inherit from + Python classes + 2. Wraps generators in the intuitive way (c.f. https://bugs.python.org/issue37743) + 3. Errors out if you try to wrap a class, because it is ambiguous whether + or not you intended to wrap only the constructor + + The input argument can either be a context manager (in which case it must + be a multi-shot context manager that can be directly invoked multiple times) + or a callable that produces a context manager. + """ + if callable(ctx) and hasattr(ctx, "__enter__"): + raise RuntimeError( + f"Passed in {ctx} is both callable and also a valid context manager " + "(has __enter__), making it ambiguous which interface to use. If you " + "intended to pass a context manager factory, rewrite your call as " + "context_decorator(lambda: ctx()); if you intended to pass a context " + "manager directly, rewrite your call as context_decorator(lambda: ctx)" + ) + + if not callable(ctx): + + def ctx_factory(): + return ctx + + else: + ctx_factory = ctx + + if inspect.isclass(func): + raise RuntimeError( + "Cannot decorate classes; it is ambiguous whether or not only the " + "constructor or all methods should have the context manager applied; " + "additionally, decorating a class at definition-site will prevent " + "use of the identifier as a conventional type. " + "To specify which methods to decorate, decorate each of them " + "individually." + ) + + if inspect.isgeneratorfunction(func): + return _wrap_generator(ctx_factory, func) + + @functools.wraps(func) + def decorate_context(*args, **kwargs): + with ctx_factory(): + return func(*args, **kwargs) + + return decorate_context + + +class _DecoratorContextManager: + """Allows a context manager to be used as a decorator.""" + + def __call__(self, orig_func: F) -> F: + if inspect.isclass(orig_func): + raise RuntimeError( + "Decorating classes is no longer supported. " + "You should only decorate functions or methods. " + "To preserve the current behavior of class decoration, you can " + "directly decorate the `__init__` method and nothing else." + ) + func = orig_func + + return cast(F, context_decorator(self.clone, func)) + + def __enter__(self) -> None: + raise NotImplementedError + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + raise NotImplementedError + + def clone(self): + # override this method if your children class takes __init__ parameters + return type(self)() + + +# TD cm functions +LAST_OP_MAPS = {} + + +def _reverse_lock(self, args, kwargs, out): + return self.unlock_() + + +LAST_OP_MAPS["lock_"] = _reverse_lock + + +def _reverse_unlock(self, args, kwargs, out): + return self.lock_() + + +LAST_OP_MAPS["unlock_"] = _reverse_unlock + + +def _reverse_transpose(self, args, kwargs, out): + dim0, dim1 = args + if not out.is_locked: + return out.update(self.transpose(dim0, dim1), inplace=False) + else: + return out.update_(self.transpose(dim0, dim1)) + + +LAST_OP_MAPS["transpose"] = _reverse_transpose +LAST_OP_MAPS["swapaxes"] = _reverse_transpose +LAST_OP_MAPS["swapdims"] = _reverse_transpose + + +def _reverse_flatten_keys(self, args, kwargs, out): + sep = args[0] if args else "." + if not out.is_locked: + return out.update(self.unflatten_keys(sep), inplace=False) + else: + return out.update_(self.unflatten_keys(sep)) + + +LAST_OP_MAPS["flatten_keys"] = _reverse_flatten_keys + + +def _reverse_unflatten_keys(self, args, kwargs, out): + sep = args[0] if args else "." + if not out.is_locked: + return out.update(self.flatten_keys(sep), inplace=False) + else: + return out.update_(self.flatten_keys(sep)) + + +LAST_OP_MAPS["unflatten_keys"] = _reverse_unflatten_keys + + +def _reverse_flatten(self, args, kwargs, out): + if len(args) == 2: + dim0, dim1 = args + elif len(args) == 1: + dim0 = args[0] + dim1 = kwargs.get("end_dim", -1) + else: + dim0 = kwargs.get("start_dim", 0) + dim1 = kwargs.get("end_dim", -1) + if dim1 < 0: + dim1 = out.ndim + dim1 + if dim0 < 0: + dim0 = out.ndim + dim0 + + if not out.is_locked: + return out.update( + self.unflatten(dim0, out.shape[dim0 : dim1 + 1]), inplace=False + ) + else: + return out.update_(self.unflatten(dim0, out.shape[dim0 : dim1 + 1])) + + +LAST_OP_MAPS["flatten"] = _reverse_flatten + + +def _reverse_unflatten(self, args, kwargs, out): + if args: + dim0 = args[0] + if len(args) > 1: + unflattened_size = args[1] + else: + unflattened_size = kwargs.get("unflattened_size") + else: + dim0 = kwargs.get("dim") + unflattened_size = kwargs.get("unflattened_size") + if dim0 < 0: + dim0 = out.ndim + dim0 + dim1 = dim0 + len(unflattened_size) - 1 + if not out.is_locked: + unflattened = self.flatten(dim0, dim1) + return out.update(unflattened, inplace=False) + else: + unflattened = self.flatten(dim0, dim1) + return out.update_(unflattened) + + +LAST_OP_MAPS["unflatten"] = _reverse_unflatten + + +def _reverse_permute(self, args, kwargs, out): + from tensordict.utils import _get_shape_from_args + + dims_list = _get_shape_from_args(*args, kwarg_name="dims", **kwargs) + dims_list = [dim if dim >= 0 else self.ndim + dim for dim in dims_list] + # inverse map + inv_dims_list = np.argsort(dims_list) + if not out.is_locked: + return out.update(self.permute(inv_dims_list), inplace=False) + else: + return out.update_(self.permute(inv_dims_list)) + + +LAST_OP_MAPS["permute"] = _reverse_permute + + +def _reverse_movedim(self, args, kwargs, out): + # Reverse of movedim(source, destination) is movedim(destination, source) + if len(args) >= 2: + source, destination = args[0], args[1] + else: + source = args[0] if args else kwargs.get("source") + destination = kwargs.get("destination") + if not out.is_locked: + return out.update(self.movedim(destination, source), inplace=False) + else: + return out.update_(self.movedim(destination, source)) + + +LAST_OP_MAPS["movedim"] = _reverse_movedim +LAST_OP_MAPS["moveaxis"] = _reverse_movedim + + +def _reverse_flip(self, args, kwargs, out): + # Flip is its own inverse + dims = args[0] if args else kwargs.get("dims") + if not out.is_locked: + return out.update(self.flip(dims), inplace=False) + else: + return out.update_(self.flip(dims)) + + +LAST_OP_MAPS["flip"] = _reverse_flip + + +def _reverse_fliplr(self, args, kwargs, out): + # fliplr is its own inverse + if not out.is_locked: + return out.update(self.fliplr(), inplace=False) + else: + return out.update_(self.fliplr()) + + +LAST_OP_MAPS["fliplr"] = _reverse_fliplr + + +def _reverse_flipud(self, args, kwargs, out): + # flipud is its own inverse + if not out.is_locked: + return out.update(self.flipud(), inplace=False) + else: + return out.update_(self.flipud()) + + +LAST_OP_MAPS["flipud"] = _reverse_flipud + + +def _reverse_roll(self, args, kwargs, out): + # Reverse of roll(shifts, dims) is roll(-shifts, dims) + shifts = args[0] if args else kwargs.get("shifts") + dims = args[1] if len(args) > 1 else kwargs.get("dims") + if isinstance(shifts, int): + neg_shifts = -shifts + else: + neg_shifts = tuple(-s for s in shifts) + if not out.is_locked: + return out.update(self.roll(neg_shifts, dims), inplace=False) + else: + return out.update_(self.roll(neg_shifts, dims)) + + +LAST_OP_MAPS["roll"] = _reverse_roll + + +def _reverse_rot90(self, args, kwargs, out): + # Reverse of rot90(k, dims) is rot90(-k, dims) or rot90(4-k, dims) + k = args[0] if args else kwargs.get("k", 1) + dims = args[1] if len(args) > 1 else kwargs.get("dims", (0, 1)) + if not out.is_locked: + return out.update(self.rot90(-k, dims), inplace=False) + else: + return out.update_(self.rot90(-k, dims)) + + +LAST_OP_MAPS["rot90"] = _reverse_rot90 + + +def _reverse_atleast_1d(self, args, kwargs, out): + # Remove added dimensions + out_ndim = out.ndim + self_ndim = self.ndim + if out_ndim > self_ndim: + for _ in range(out_ndim - self_ndim): + self = self.squeeze(0) + if not out.is_locked: + return out.update(self, inplace=False) + else: + return out.update_(self) + + +LAST_OP_MAPS["atleast_1d"] = _reverse_atleast_1d +LAST_OP_MAPS["atleast_2d"] = _reverse_atleast_1d +LAST_OP_MAPS["atleast_3d"] = _reverse_atleast_1d + + +def _reverse_view(self, args, kwargs, out): + if not out.is_locked: + return out.update(self.view(out.shape), inplace=False) + else: + return out.update_(self.view(out.shape)) + + +LAST_OP_MAPS["view"] = _reverse_view + + +def _reverse_unsqueeze(self, args, kwargs, out): + if args: + (dim,) = args + elif kwargs: + dim = kwargs["dim"] + else: + raise RuntimeError( + "Cannot use td.unsqueeze() as a decorator if the dimension is implicit." + ) + if not out.is_locked: + return out.update(self.squeeze(dim), inplace=False) + else: + return out.update_(self.squeeze(dim)) + + +LAST_OP_MAPS["unsqueeze"] = _reverse_unsqueeze + + +def _reverse_squeeze(self, args, kwargs, out): + if args: + (dim,) = args + elif kwargs: + dim = kwargs["dim"] + else: + raise RuntimeError( + "Cannot use td.squeeze() as a decorator if the dimension is implicit." + ) + if not out.is_locked: + return out.update(self.unsqueeze(dim), inplace=False) + else: + return out.update_(self.unsqueeze(dim)) + + +LAST_OP_MAPS["squeeze"] = _reverse_squeeze + + +def _reverse_to_module(self, args, kwargs, out): + try: + with ( + out.unlock_() + if not is_compiling() and out is not None + else contextlib.nullcontext() + ): + return self.to_module(*args, **kwargs, swap_dest=out) + except AttributeError: + # This is a bit unsafe but we assume that out won't have an unlock_() if it's not a TD + raise RuntimeError( + "to_module cannot be used as a decorator when return_swap=False." + ) + + +LAST_OP_MAPS["to_module"] = _reverse_to_module + + +def _reverse_to(self, args, kwargs, out): + """Reverse the to() operation by restoring the original device. + + Uses the input tensordict (self) to determine the original device of each tensor + and restores the output tensordict (out) to those original devices. + """ + if out is None: + return self + # Restore each tensor to its original device/dtype from the input tensordict + return self.apply( + lambda x, y: x.to(y.device) if y is not None else x, out, default=None, out=out + ) + + +LAST_OP_MAPS["to"] = _reverse_to diff --git a/lib/python3.12/site-packages/tensordict/_datasets.py b/lib/python3.12/site-packages/tensordict/_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..d65037cf9ad659e4dee12b0e3e84ed05e0db3ced --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_datasets.py @@ -0,0 +1,219 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +from typing import Literal + +import numpy as np +import torch +from tensordict._tensorcollection import TensorCollection +from tensordict.utils import is_non_tensor + + +def to_mds( + self: TensorCollection, + *, + columns: dict[str, str] | None = None, + out: str | tuple[str, str], + keep_local: bool = False, + compression: str | None = None, + hashes: list[str] | None = None, + size_limit: int | str | None = 1 << 26, + writer: "streaming.MDSWriter" | None = None, # noqa # type-ignore + **kwargs, +) -> None: + """Writes the content of a TensorCollection to a streaming dataset. + + Keyword Args: + out (str | Tuple[str, str]): Output dataset directory to save shard files. + + 1. If ``out`` is a local directory, shard files are saved locally. + 2. If ``out`` is a remote directory, a local temporary directory is created to + cache the shard files and then the shard files are uploaded to a remote + location. At the end, the temp directory is deleted once shards are uploaded. + 3. If ``out`` is a tuple of ``(local_dir, remote_dir)``, shard files are saved in the + `local_dir` and also uploaded to a remote location. + columns (dict[str, str]): an optional dict of columns. Will be automatically inferred from + the tensor collection if not provided. + keep_local (bool): If the dataset is uploaded, whether to keep the local dataset directory + or remove it after uploading. Defaults to ``False``. + compression (str, optional): Optional compression or compression:level. Defaults to + ``None``. + hashes (List[str], optional): Optional list of hash algorithms to apply to shard files. + Defaults to ``None``. + size_limit (Union[int, str], optional): Optional shard size limit, after which point to + start a new shard. If ``None``, puts everything in one shard. Can specify bytes + human-readable format as well, for example ``"100kb"`` for 100 kilobyte + (100*1024) and so on. Defaults to ``1 << 26``. + Ignored if + writer: (MDSWriter, optional): the write to use. Will be created from the `out` kwarg as well as other + input kwargs. + **kwargs (Any): Additional settings for the Writer. + + .. note:: The MDSWriter has limited support for nested dictionaries. The proper way to handle nested tensordicts + is to use the :meth:`~tensordict.TensorDictBase.flatten_keys` method before writing, and :meth:`~tensordict.TensorDictBase.unflatten_keys` method after reading. + + .. warning:: + This method requires `mosaicml-streaming` to be installed. + + .. warning:: + For non-tensor data, the type of the data must be fixed. The way tensordict recovers the data type is by + looking at the first element of the list. If it's `None` an error will be thrown. Otherwise all the + data in the list must have the same type (or be None if missing). + + .. seealso:: See the Mosaic streaming library API at ``_ + + The following example shows and end-to-end example of how to create a dataset and load it in a + PyTorch dataloader. + + Examples: + >>> import tempfile + >>> from typing import Any + >>> from tensordict import TensorDict, LazyStackedTensorDict + >>> import torch + >>> + >>> + >>> td = LazyStackedTensorDict( + ... TensorDict(a=0, b=1, c=torch.randn(2), d="a string"), + ... TensorDict(a=0, b=1, c=torch.randn(3), d="another string"), + ... TensorDict(a=0, b=1, c=torch.randn(3), d="yet another string"), + ... ) + >>> + >>> with tempfile.TemporaryDirectory() as tempdir: + ... # Create a dataset on one process / thread / node... + ... td.to_mds(out=tempdir) + ... + ... # Create a dataloader + ... from streaming import StreamingDataset + ... + ... # Load the dataset on another thread / process / node... + ... dataset = StreamingDataset(local=tempdir, remote=None, batch_size=2) + ... + ... # Use the class `from_list` method as a collate_fn + ... dl = torch.utils.data.DataLoader(dataset=dataset, batch_size=2, collate_fn=LazyStackedTensorDict.from_list) + ... for batch in dl: + ... print("batch", batch) + + """ + try: + from streaming import MDSWriter + except ImportError: + raise ImportError( + "Failed to load MDSWriter from streaming. Check that mosaicml-streaming is installed." + ) + + if not self.numel(): + raise ValueError("Cannot write an empty TensorDict to a streaming dataset.") + if self.ndim: + if writer is None: + if columns is None: + columns = _columns(self[0]) + writer = MDSWriter( + out=out, + columns=columns, + keep_local=keep_local, + compression=compression, + hashes=hashes, + size_limit=size_limit, + **kwargs, + ) + td_dict = self.tolist(convert_tensors="numpy") + if not isinstance(td_dict, list): + raise ValueError(f"Expected a list of dictionarie, got {type(td_dict)}.") + with writer as w: + for td_i in td_dict: + for k, v in td_i.items(): + if isinstance(v, np.ndarray) and v.shape == (): + td_i[k] = v.item() + w.write(td_i) + return + + if writer is None: + if columns is None: + columns = _columns(self) + writer = MDSWriter( + out=out, + columns=columns, + keep_local=keep_local, + compression=compression, + hashes=hashes, + size_limit=size_limit, + ) + writer.write(self) + + +def _columns(self: TensorCollection) -> dict[str, str]: + return {k: _get_elt_type(v) for k, v in self.items()} + + +def _get_elt_type( + elt, +) -> Literal[ + "bytes", + "float", + "float16", + "float32", + "float64", + "int", + "int16", + "int32", + "int64", + "int8", + "jpeg", + "jpeg_array", + "jpegarray", + "json", + "list[jpeg]", + "list[pil]", + "list[png]", + "ndarray", + "pil", + "pkl", + "png", + "str", + "str_decimal", + "str_float", + "str_int", + "uint16", + "uint32", + "uint64", + "uint8", +]: + if is_non_tensor(elt): + from tensordict._lazy import LazyStackedTensorDict + + if isinstance(elt, LazyStackedTensorDict): + return _get_elt_type(elt.tolist()) + return _get_elt_type(elt.data) + if isinstance(elt, list): + return "json" + if isinstance(elt, bytes): + return "bytes" + elif isinstance(elt, (torch.Tensor, np.ndarray)): + if not elt.shape: + # get the dtype and return + if isinstance(elt, torch.Tensor): + dtype = elt.dtype + if dtype.is_floating_point: + return "float" + else: + return "int" + else: + dtype = elt.dtype + if dtype.kind == "f": + return "float" + elif dtype.kind in ("i", "b"): + return "int" + else: + return "str" + return "ndarray" + elif isinstance(elt, float): + return "float" + elif isinstance(elt, (int, bool)): + return "int" + elif isinstance(elt, str): + return "str" + else: + raise ValueError(f"Unknown type {type(elt)}") diff --git a/lib/python3.12/site-packages/tensordict/_lazy.py b/lib/python3.12/site-packages/tensordict/_lazy.py new file mode 100644 index 0000000000000000000000000000000000000000..6608d281ddacb154a5c2bca25efec8903342d0b5 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_lazy.py @@ -0,0 +1,4922 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import numbers +import os +import re +import textwrap +import weakref +from collections import defaultdict +from concurrent.futures import Future, ThreadPoolExecutor +from copy import copy, deepcopy +from functools import wraps +from pathlib import Path +from textwrap import indent + +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Literal, + Mapping, + OrderedDict, + Sequence, + Tuple, + Type, + TYPE_CHECKING, +) + +import numpy as np + +import torch +from tensordict._td import _SubTensorDict, _TensorDictKeysView, TensorDict +from tensordict._tensorcollection import TensorCollection +from tensordict.base import ( + _is_leaf_nontensor, + _is_tensor_collection, + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + _register_tensor_class, + BEST_ATTEMPT_INPLACE, + CompatibleType, + is_tensor_collection, + NO_DEFAULT, + T, + TensorDictBase, +) + +from tensordict.memmap import MemoryMappedTensor +from tensordict.utils import ( + _as_context_manager, + _broadcast_tensors, + _check_is_flatten, + _check_is_unflatten, + _get_shape_from_args, + _getitem_batch_size, + _infer_size_impl, + _is_number, + _maybe_correct_neg_dim, + _parse_to, + _recursive_unbind_list, + _renamed_inplace_method, + _shape, + _td_fields, + _unravel_key_to_tuple, + _zip_strict, + cache, + convert_ellipsis_to_idx, + DeviceType, + erase_cache, + expand_right, + IndexType, + infer_size_impl, + is_non_tensor, + is_tensorclass, + list_to_stack, + lock_blocked, + NestedKey, + unravel_key_list, +) +from torch import Tensor +from torch.nn.utils.rnn import pad_sequence + + +try: + from functorch import dim as ftdim + + _has_funcdim = True +except ImportError: + from tensordict.utils import _ftdim_mock as ftdim + + _has_funcdim = False + +try: + from tensordict.utils import _import_and_wrap_functorch + + _add_batch_dim, _remove_batch_dim = _import_and_wrap_functorch( + "_add_batch_dim", + "_remove_batch_dim", + ) +except ImportError: + + def _add_batch_dim(*args, **kwargs) -> Tensor: + raise NotImplementedError + + def _remove_batch_dim(*args, **kwargs) -> Tensor: + raise NotImplementedError + + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +class _LazyStackedTensorDictKeysView(_TensorDictKeysView): + tensordict: LazyStackedTensorDict + + def __len__(self) -> int: + return len(self._keys()) + + def _keys(self) -> list[str]: + result = self.tensordict._key_list() + if self.is_leaf in ( + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + ): + return [ + (key, str(i)) + for key in result + for i in range(len(self.tensordict.tensordicts)) + ] + return result + + def __contains__(self, item): + item = _unravel_key_to_tuple(item) + if item[0] in self.tensordict._iterate_over_keys(): + if self.leaves_only: + return not _is_tensor_collection(self.tensordict.entry_class(item[0])) + has_first_key = True + else: + has_first_key = False + if not has_first_key or len(item) == 1: + return has_first_key + # otherwise take the long way + return all( + item[1:] + in tensordict.get(item[0]).keys(self.include_nested, self.leaves_only) + for tensordict in self.tensordict.tensordicts + ) + + def __repr__(self): + return f"{type(self).__name__}({tuple(self)})" + + +def _fails_exclusive_keys(func): + @wraps(func) + def newfunc(self, *args, **kwargs): + if self._has_exclusive_keys: + raise RuntimeError( + f"the method {func.__name__} cannot complete when there are exclusive keys." + ) + parent_func = getattr(TensorDictBase, func.__name__, None) + if parent_func is None: + parent_func = getattr(TensorDict, func.__name__) + return parent_func(self, *args, **kwargs) + + return newfunc + + +class LazyStackedTensorDict(TensorDictBase): + """A Lazy stack of TensorDicts. + + When stacking TensorDicts together, the default behaviour is to put them + in a stack that is not instantiated. + This allows to seamlessly work with stacks of tensordicts with operations + that will affect the original tensordicts. + + Args: + *tensordicts (TensorDict instances): a list of tensordict with + same batch size. + stack_dim (int): a dimension (between `-td.ndimension()` and + `td.ndimension()-1` along which the stack should be performed. + hook_out (callable, optional): a callable to execute after :meth:`~.get`. + hook_in (callable, optional): a callable to execute before :meth:`~.set`. + stack_dim_name (str, optional): the name of the stack dimension. + Defaults to ``None``. + strict_shape (bool, optional): if ``True``, every tensordict's shapes must match. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> tds = [TensorDict({'a': torch.randn(3, 4)}, batch_size=[3]) + ... for _ in range(10)] + >>> td_stack = torch.stack(tds, -1) + >>> print(td_stack.shape) + torch.Size([3, 10]) + >>> print(td_stack.get("a").shape) + torch.Size([3, 10, 4]) + >>> print(td_stack[:, 0] is tds[0]) + True + + .. note:: Lazy stacks support assignment via lists. For consistency, the lists should be + presented as `tensor.tolist()` data structure. This means that the length of the first + level of the nested lists should match the first dimension of the lazy stack (whether or + not this is the stack dimension). + + >>> td = LazyStackedTensorDict(TensorDict(), TensorDict(), stack_dim=0) + >>> td["a"] = [torch.ones(2), torch.zeros(1)] + >>> assert td[1]["a"] == torch.zeros(1) + >>> td["b"] = ["a string", "another string"] + >>> assert td[1]["b"] == "another string" + + .. note:: When using the :meth:`~.get` method, one can pass `as_nested_tensor`, `as_padded_tensor` + or the `as_list` arguments to control how the data should be presented if the dimensions of the + tensors mismatch. When passed, the nesting/padding will occur regardless of whether the + dimensions mismatch or not. + + """ + + _is_vmapped: bool = False + _device: torch.device | None = None + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + from tensordict._torch_func import LAZY_TD_HANDLED_FUNCTIONS + + if func in LAZY_TD_HANDLED_FUNCTIONS: + if kwargs is None: + kwargs = {} + if func not in LAZY_TD_HANDLED_FUNCTIONS or not all( + issubclass(t, (Tensor, TensorDictBase)) for t in types + ): + return NotImplemented + return LAZY_TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + else: + return super().__torch_function__(func, types, args, kwargs) + + _td_dim_name = None + _safe = False + _lazy = True + + def __init__( + self, + *tensordicts: T, + stack_dim: int = 0, + hook_out: callable | None = None, + hook_in: callable | None = None, + batch_size: Sequence[int] | None = None, + device: torch.device | None = None, + names: Sequence[str] | None = None, + stack_dim_name: str | None = None, + strict_shape: bool = False, + ) -> None: + self._is_locked = None + + # sanity check + num_tds = len(tensordicts) + batch_size = torch.Size(batch_size) if batch_size is not None else None + if not num_tds: + # create an empty tensor + td0 = TensorDict(batch_size=batch_size, device=device, names=names) + self._device = torch.device(device) if device is not None else None + else: + td0 = tensordicts[0] + device = td0.device + if stack_dim < 0: + ndim = td0.ndim + try: + stack_dim = _maybe_correct_neg_dim(stack_dim, ndim=ndim + 1, shape=None) + except Exception: + raise RuntimeError( + f"Couldn't infer stack dim from negative value, got stack_dim={stack_dim}" + ) + self.stack_dim = stack_dim + self._reset_batch_size(td0, tensordicts, device, num_tds, strict_shape) + if stack_dim > len(self.batch_size): + raise RuntimeError( + f"Stack dim {stack_dim} is too big for batch size {self.batch_size}." + ) + + self.tensordicts: list[TensorDictBase] = list(tensordicts) + self.hook_out = hook_out + self.hook_in = hook_in + if batch_size is not None and batch_size != self.batch_size and num_tds != 0: + raise RuntimeError( + f"batch_size does not match self.batch_size: {batch_size} vs {self.batch_size}." + ) + if stack_dim_name is not None: + self._td_dim_name = stack_dim_name + + def _new_impl( + self, + size: torch.Size, + *, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + funcname: str, + empty_lazy: bool = False, + **kwargs, + ): + # The tricky thing here is to figure out if (1) we want to return a lazy stack or + # a plain TD and (2) if it's a lazy stack, what to do if the ndim the user is asking for + # differs from the ndims of the lazy stack (ie, where should the lazy stack dim lie in the + # new td)? Currently, we use the same stack dim if possible, and only when empty_lazy is True. + if not empty_lazy: + return TensorDictBase._new_impl( + self, + size=size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + layout=layout, + funcname=funcname, + empty_lazy=empty_lazy, + **kwargs, + ) + if isinstance(size, int): + size = (size,) + elif len(size) == 1 and not isinstance(size[0], int): + size = size[0] + + ndim = self.ndim + if device is not NO_DEFAULT: + kwargs["device"] = device + if ( + empty_lazy and len(size) > self.stack_dim + ): # eg, stack_dim is 1 and size is (3, 4) + sub_size = list(size) + # TODO: consider starting from the end? + stack_dim = self.stack_dim # - self.ndim + sub_size.pop(stack_dim) + td = ( + self.tensordicts[0] + .empty(recurse=True) + .new_empty(sub_size, empty_lazy=empty_lazy) + ) + # reproduce td as many times as needed + return self._new_lazy_unsafe( + *[td.copy() for _ in range(size[self.stack_dim])], + stack_dim=self.stack_dim, + hook_out=self.hook_out, + hook_in=self.hook_in, + stack_dim_name=self._td_dim_name, + ) + + names = self._maybe_names() + if names: + if len(size) > self.ndim: + names = [None] * (len(size) - self.ndim) + list(names) + elif self.ndim > len(size): + names = names[-len(size) :] + + def func(tensor, size=size): + feature_shape = tensor.shape[ndim:] + size = torch.Size((*size, *feature_shape)) + kwargs_copy = kwargs + if empty_lazy and is_tensor_collection(tensor): + kwargs_copy = dict(kwargs) + kwargs_copy["empty_lazy"] = empty_lazy + return getattr(tensor, funcname)( + size, + dtype=dtype, + requires_grad=requires_grad, + layout=layout, + **kwargs_copy, + ) + + result = self._fast_apply( + func, + call_on_nested=True, + device=device, + batch_size=size, + names=names, + ) + return result + + @classmethod + def _new_lazy_unsafe( + cls, + *tensordicts: T, + stack_dim: int = 0, + hook_out: callable | None = None, + hook_in: callable | None = None, + batch_size: Sequence[int] | None = None, + device: torch.device | None = None, + names: Sequence[str] | None = None, + stack_dim_name: str | None = None, + strict_shape: bool = False, + ) -> None: + self = cls.__new__(cls) + self._is_locked = None + + # sanity check + num_tds = len(tensordicts) + batch_size = torch.Size(batch_size) if batch_size is not None else None + if not num_tds: + # create an empty tensor + td0 = TensorDict(batch_size=batch_size, device=device, names=names) + self._device = torch.device(device) if device is not None else None + else: + td0 = tensordicts[0] + # device = td0.device + _batch_size = td0.batch_size + + for td in tensordicts[1:]: + _bs = td.batch_size + if _bs != _batch_size: + _batch_size = torch.Size( + [s if _bs[i] == s else -1 for i, s in enumerate(_batch_size)] + ) + self.tensordicts: list[TensorDictBase] = list(tensordicts) + self.stack_dim = stack_dim + self._batch_size = self._compute_batch_size(_batch_size, stack_dim, num_tds) + self.hook_out = hook_out + self.hook_in = hook_in + if stack_dim_name is not None: + self._td_dim_name = stack_dim_name + return self + + # These attributes should never be set + @property + def _is_shared(self): + return all(td._is_shared for td in self.tensordicts) + + @property + def _is_memmap(self): + return all(td._is_memmap for td in self.tensordicts) + + @property + @cache # noqa: B019 + def _has_exclusive_keys(self): + keys = None + for td in self.tensordicts: + _keys = set(td.keys(True, True)) + if keys is None: + keys = _keys + else: + if keys != _keys: + return True + else: + return False + + @_fails_exclusive_keys + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + ) -> dict[str, Any]: ... + + def _reduce_get_metadata(self): + metadata = {} + metadata["stack_dim"] = self.stack_dim + metadata["stack_dim_name"] = self._td_dim_name + metadata["is_locked"] = self.is_locked + return metadata + + @classmethod + def from_dict( + cls, + input_dict: Dict[str, Dict[NestedKey, Any]] | List[Dict[NestedKey, Any]], + *other, + auto_batch_size: bool = False, + batch_size=None, + device=None, + batch_dims=None, + stack_dim_name=None, + stack_dim=0, + names: list[str] | None = None, + ): + if names is not None: + if stack_dim_name is not None: + raise ValueError("names and stack_dim_name are exclusive.") + stack_dim_name = list(names).pop(stack_dim) + + return cls._new_lazy_unsafe( + *( + TensorDict.from_dict( + ( + input_dict[str(i)] + if isinstance(input_dict, Mapping) + else input_dict[i] + ), + *other, + auto_batch_size=auto_batch_size, + device=device, + batch_dims=batch_dims, + batch_size=batch_size, + names=names, + ) + for i in range(len(input_dict)) + ), + stack_dim=stack_dim, + stack_dim_name=stack_dim_name, + ) + + @classmethod + def from_list( + cls, + input: list[Mapping], + *, + auto_batch_size: bool | None = None, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: list[str] | None = None, + lazy: bool | None = None, + ) -> Self: + return cls.from_dict( + input, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + batch_size=batch_size, + device=device, + names=names, + ) + + @_fails_exclusive_keys + def state_dict( + self, + destination=None, + prefix="", + keep_vars=False, + flatten=False, + ) -> OrderedDict[str, Any]: ... + + @_fails_exclusive_keys + def flatten_keys( + self, + separator: str = ".", + inplace: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + ) -> Self: ... + + @_fails_exclusive_keys + def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> Self: ... + + @property + def device(self) -> torch.device | None: + # devices might have changed, so we check that they're all the same + if self.tensordicts: + device = self.tensordicts[0].device + for td in self.tensordicts: + if device != td.device: + return None + return device + return self._device + + @device.setter + def device(self, value: DeviceType) -> None: + if not self.tensordicts: + self._device = torch.device(value) if value is not None else value + return + for t in self.tensordicts: + t.device = value + + def clear_device_(self) -> Self: + for td in self.tensordicts: + td.clear_device_() + return self + + @property + def batch_size(self) -> torch.Size: + return self._batch_size + + @batch_size.setter + def batch_size(self, new_size: torch.Size) -> None: + return self._batch_size_setter(new_size) + + @property + @cache # noqa + def names(self): + names = list(self.tensordicts[0].names) + for td in self.tensordicts[1:]: + if names != td.names: + raise ValueError( + f"Not all dim names match, got {names} and {td.names}." + ) + names.insert(self.stack_dim, self._td_dim_name) + return names + + @names.setter + @erase_cache # a nested lazy stacked tensordict is not apparent to the root + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + for td in self.tensordicts: + td.names = None + self._td_dim_name = None + else: + names_c = list(names) + name = names_c[self.stack_dim] + self._td_dim_name = name + del names_c[self.stack_dim] + for td in self.tensordicts: + if td._check_dim_name(name): + # TODO: should reset names here + raise ValueError(f"The dimension name {name} is already taken.") + td.rename_(*names_c) + + def _rename_subtds(self, names): + # remove the name of the stack dim + names = list(names) + del names[self.stack_dim] + for td in self.tensordicts: + td.names = names + + def _has_names(self): + return all(td._has_names() for td in self.tensordicts) + + def _erase_names(self): + self._td_dim_name = None + for td in self.tensordicts: + td._erase_names() + + def get_item_shape(self, key): + """Gets the shape of an item in the lazy stack. + + Heterogeneous dimensions are returned as -1. + + This implementation is inefficient as it will attempt to stack the items + to compute their shape, and should only be used for printing. + """ + try: + item = self.get(key) + return item.shape + except RuntimeError as err: + if re.match( + r"Failed to stack tensors within a tensordict", + str(err), + ): + shape = None + for td in self.tensordicts: + if shape is None: + shape = list(td.get_item_shape(key)) + else: + _shape = td.get_item_shape(key) + if len(shape) != len(_shape): + shape = [-1] + return torch.Size(shape) + shape = [ + s1 if s1 == s2 else -1 + for (s1, s2) in _zip_strict(shape, _shape) + ] + shape.insert(self.stack_dim, len(self.tensordicts)) + return torch.Size(shape) + else: + raise err + + def is_shared(self) -> bool: + are_shared = [td.is_shared() for td in self.tensordicts] + are_shared = [value for value in are_shared if value is not None] + if not len(are_shared): + return None + if any(are_shared) and not all(are_shared): + raise RuntimeError( + f"tensordicts shared status mismatch, got {sum(are_shared)} " + f"shared tensordicts and " + f"{len(are_shared) - sum(are_shared)} non shared tensordict " + ) + return all(are_shared) + + def is_memmap(self) -> bool: + are_memmap = [td.is_memmap() for td in self.tensordicts] + if any(are_memmap) and not all(are_memmap): + raise RuntimeError( + f"tensordicts memmap status mismatch, got {sum(are_memmap)} " + f"memmap tensordicts and " + f"{len(are_memmap) - sum(are_memmap)} non memmap tensordict " + ) + return are_memmap[0] + + def _reset_batch_size( + self, + td0: TensorDictBase, + tensordicts: list[TensorDictBase], + device: torch.device, + num_tds: int, + strict_shape: bool, + ): + _batch_size = td0.batch_size + stack_dim = self.stack_dim + + for td in tensordicts[1:]: + if not is_tensor_collection(td): + raise TypeError( + "Expected all inputs to be TensorDictBase instances but got " + f"{type(td)} instead." + ) + _bs = td.batch_size + _device = td.device + if device != _device: + raise RuntimeError(f"devices differ, got {device} and {_device}") + if _bs != _batch_size: + if strict_shape or len(_bs) != len(_batch_size): + raise RuntimeError( + f"batch sizes in tensordicts differs, LazyStackedTensorDict " + f"cannot be created. Got td[0].batch_size={_batch_size} " + f"and td[i].batch_size={_bs}. If the length match and you wish " + f"to stack these tensordicts, set strict_shape to False." + ) + else: + _batch_size = torch.Size( + [s if _bs[i] == s else -1 for i, s in enumerate(_batch_size)] + ) + self._batch_size = self._compute_batch_size(_batch_size, stack_dim, num_tds) + + @staticmethod + def _compute_batch_size( + batch_size: torch.Size, stack_dim: int, num_tds: int + ) -> torch.Size: + s = list(batch_size) + s.insert(stack_dim, num_tds) + return torch.Size(s) + + def _set_str( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ) -> Self: + try: + inplace = self._convert_inplace(inplace, key) + except KeyError as e: + raise KeyError( + "setting a value in-place on a stack of TensorDict is only " + "permitted if all members of the stack have this key in " + "their register." + ) from e + if not validated: + value = self._validate_value( + value, + non_blocking=non_blocking, + check_shape=not (isinstance(value, list) and list_to_stack()), + ) + validated = True + if self._is_vmapped: + value = self.hook_in(value) + if isinstance(value, list): + if self.stack_dim == 0: + values = list(value) + else: + values = _recursive_unbind_list(value, self.stack_dim) + else: + values = value.unbind(self.stack_dim) + for tensordict, item in _zip_strict(self.tensordicts, values): + tensordict._set_str( + key, + item, + inplace=inplace, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + return self + + def _set_tuple( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + non_blocking: bool = False, + ) -> Self: + if len(key) == 1: + return self._set_str( + key[0], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + # if inplace is not False: # inplace could be None + # # we don't want to end up in the situation where one tensordict has + # # inplace=True and another one inplace=False because inplace was loose. + # # Worse could be writing with inplace=True up until some level then to + # # realize the key is missing in one td, raising an exception and having + # # messed up the data. Hence we must start by checking if the key + # # is present. + # has_key = key in self.keys(True) + # if inplace is True and not has_key: # inplace could be None + # raise KeyError( + # TensorDictBase.KEY_ERROR.format( + # key, type(self).__name__, sorted(self.keys()) + # ) + # ) + # inplace = has_key + if not validated: + value = self._validate_value(value, non_blocking=non_blocking) + validated = True + if self._is_vmapped: + value = self.hook_in(value) + values = value.unbind(self.stack_dim) + for tensordict, item in _zip_strict(self.tensordicts, values): + tensordict._set_tuple( + key, + item, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + return self + + def _split_index(self, index): + """Given a tuple index, split it in as many indices as the number of tensordicts. + + Returns: + a dictionary with {index-of-td: index-within-td} + the number of single dim indices until stack dim + a boolean indicating if the index along the stack dim is an integer + """ + if not isinstance(index, tuple): + index = (index,) + index = convert_ellipsis_to_idx(index, self.batch_size) + index = _broadcast_tensors(index) + out = [] + num_single = 0 + num_none = 0 + isinteger = False + is_nd_tensor = False + cursor = 0 # the dimension cursor + selected_td_idx = torch.arange(len(self.tensordicts)) + has_bool = False + num_squash = 0 + encountered_tensor = False + for i, idx in enumerate(index): # noqa: B007 + cursor_incr = 1 + # if idx is None: + # idx = True + if idx is None or idx is True: + out.append(None) + num_none += cursor <= self.stack_dim + continue + if cursor == self.stack_dim: + # we need to check which tds need to be indexed + if isinstance(idx, ftdim.Dim): + raise ValueError( + "Cannot index a lazy stacked tensordict along the stack dimension with " + "a first-class dimension index. Consider consolidating the tensordict first " + "using `tensordict.contiguous()`." + ) + elif isinstance(idx, slice) or _is_number(idx): + selected_td_idx = range(len(self.tensordicts))[idx] + if not isinstance(selected_td_idx, range): + isinteger = True + selected_td_idx = [selected_td_idx] + elif isinstance(idx, torch.Tensor): + if idx.dtype == torch.bool: + # we mark that we need to dispatch the indices across stack idx + has_bool = True + # split mask along dim + individual_masks = idx = idx.unbind(0) + selected_td_idx = range(len(self.tensordicts)) + out.append(idx) + split_dim = self.stack_dim - num_single + mask_loc = i + else: + is_nd_tensor = True + if not encountered_tensor: + # num_single -= idx.ndim - 1 + encountered_tensor = True + else: + num_single += 1 + selected_td_idx = idx + # out.append(idx.unbind(0)) + else: + raise TypeError(f"Invalid index type: {type(idx)}.") + else: + if _is_number(idx) and cursor < self.stack_dim: + num_single += 1 + if _is_number(idx) or isinstance( + idx, + ( + ftdim.Dim, + slice, + ), + ): + out.append(idx) + elif isinstance(idx, torch.Tensor): + if idx.dtype == torch.bool: + cursor_incr = idx.ndim + if cursor < self.stack_dim: + num_squash += cursor_incr - 1 + if ( + cursor < self.stack_dim + and cursor + cursor_incr > self.stack_dim + ): + # we mark that we need to dispatch the indices across stack idx + has_bool = True + # split mask along dim + # relative_stack_dim = self.stack_dim - cursor - cursor_incr + individual_masks = idx = idx.unbind(0) + selected_td_idx = range(self.shape[i]) + split_dim = cursor - num_single + mask_loc = i + elif cursor < self.stack_dim: + # we know idx is not a single integer, so it must have + # a dimension. We play with num_single, reducing it + # by the number of dims of idx: if idx has 3 dims, our + # indexed tensor will have 2 more dimensions, going in + # the opposite direction of indexing with a single integer, + # smth[torch.tensor(1)].ndim = smth.ndim-1 + # smth[torch.tensor([1])].ndim = smth.ndim + # smth[torch.tensor([[1]])].ndim = smth.ndim+1 + if not encountered_tensor: + num_single -= idx.ndim - 1 + encountered_tensor = True + else: + num_single += 1 + out.append(idx) + else: + raise TypeError(f"Invalid index type: {type(idx)}.") + cursor += cursor_incr + if has_bool: + out = tuple( + tuple(idx if not isinstance(idx, tuple) else idx[i] for idx in out) + for i in selected_td_idx + ) + return { + "index_dict": {i: out[i] for i in selected_td_idx}, + "num_single": num_single, + "isinteger": isinteger, + "has_bool": has_bool, + "individual_masks": individual_masks, + "split_dim": split_dim, + "mask_loc": mask_loc, + "is_nd_tensor": is_nd_tensor, + "num_none": num_none, + "num_squash": num_squash, + } + elif is_nd_tensor: + + def isindexable(idx): + if isinstance(idx, torch.Tensor): + if idx.dtype == torch.bool: + return False + return True + if isinstance(idx, (tuple, list, range)): + return True + return False + + def outer_list(tensor_index, tuple_index): + """Converts a tensor and a tuple to a nested list where each leaf is a (int, index) tuple where the index only points to one element.""" + if isinstance(tensor_index, torch.Tensor): + list_index = tensor_index.tolist() + else: + list_index = tensor_index + list_result = [] + + def index_tuple_index(i, convert=False): + for idx in tuple_index: + if isindexable(idx): + if convert: + yield int(idx[i]) + else: + yield idx[i] + else: + yield idx + + for i, idx in enumerate(list_index): + if isinstance(idx, int): + list_result.append( + (idx, tuple(index_tuple_index(i, convert=True))) + ) + elif isinstance(idx, list): + list_result.append(outer_list(idx, tuple(index_tuple_index(i)))) + else: + raise NotImplementedError + return list_result + + return { + "index_dict": outer_list(selected_td_idx, out), + "num_single": num_single, + "isinteger": isinteger, + "has_bool": has_bool, + "is_nd_tensor": is_nd_tensor, + "num_none": num_none, + "num_squash": num_squash, + } + return { + "index_dict": {i: tuple(out) for i in selected_td_idx}, + "num_single": num_single, + "isinteger": isinteger, + "has_bool": has_bool, + "is_nd_tensor": is_nd_tensor, + "num_none": num_none, + "num_squash": num_squash, + } + + def _set_at_str(self, key, value, index, *, validated, non_blocking: bool): + if not validated: + value = self._validate_value( + value, check_shape=False, non_blocking=non_blocking + ) + validated = True + if self._is_vmapped: + value = self.hook_in(value) + split_index = self._split_index(index) + converted_idx = split_index["index_dict"] + num_single = split_index["num_single"] + isinteger = split_index["isinteger"] + has_bool = split_index["has_bool"] + num_squash = split_index.get("num_squash", 0) + num_none = split_index.get("num_none", 0) + is_nd_tensor = split_index.get("is_nd_tensor", False) + if isinteger: + # this will break if the index along the stack dim is [0] or :1 or smth + for i, _idx in converted_idx.items(): + self.tensordicts[i]._set_at_str( + key, value, _idx, validated=validated, non_blocking=non_blocking + ) + return self + if is_nd_tensor: + unbind_dim = self.stack_dim - num_single + num_none - num_squash + value_unbind = value.unbind(unbind_dim) + + def set_at_str(converted_idx): + for i, item in enumerate(converted_idx): + if isinstance(item, list): + set_at_str(item) + else: + _value = value_unbind[i] + stack_idx, idx = item + self.tensordicts[stack_idx]._set_at_str( + key, + _value, + idx, + validated=validated, + non_blocking=non_blocking, + ) + + set_at_str(converted_idx) + return self + elif not has_bool: + unbind_dim = self.stack_dim - num_single + num_none - num_squash + value_unbind = value.unbind(unbind_dim) + for (i, _idx), _value in _zip_strict( + converted_idx.items(), + value_unbind, + ): + self.tensordicts[i]._set_at_str( + key, _value, _idx, validated=validated, non_blocking=non_blocking + ) + else: + # we must split, not unbind + mask_unbind = split_index["individual_masks"] + split_dim = split_index["split_dim"] + splits = [_mask_unbind.sum().item() for _mask_unbind in mask_unbind] + value_unbind = value.split(splits, split_dim) + if mask_unbind[0].ndim == 0: + # we can return a stack + for (i, _idx), mask, _value in _zip_strict( + converted_idx.items(), + mask_unbind, + value_unbind, + ): + if mask.any(): + self.tensordicts[i]._set_at_str( + key, + _value, + _idx, + validated=validated, + non_blocking=non_blocking, + ) + else: + for (i, _idx), _value in _zip_strict( + converted_idx.items(), value_unbind + ): + self_idx = (slice(None),) * split_index["mask_loc"] + (i,) + self[self_idx]._set_at_str( + key, + _value, + _idx, + validated=validated, + non_blocking=non_blocking, + ) + + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool): + if len(key) == 1: + return self._set_at_str( + key[0], value, idx, validated=validated, non_blocking=non_blocking + ) + # get the "last" tds + tds = [] + for td in self.tensordicts: + tds.append(td.get(key[:-1])) + # build only a single lazy stack from it + # (if the stack is a stack of stacks this won't be awesomely efficient + # but then we'd need to splut the value (which we can do) and recompute + # the sub-index for each td, which is a different story! + td = LazyStackedTensorDict( + *tds, stack_dim=self.stack_dim, hook_out=self.hook_out, hook_in=self.hook_in + ) + if not validated: + value = self._validate_value( + value, check_shape=False, non_blocking=non_blocking + ) + validated = True + if self._is_vmapped: + value = self.hook_in(value) + item = td._get_str(key, NO_DEFAULT) + item[idx] = value + td._set_str(key, item, inplace=True, validated=True, non_blocking=non_blocking) + return self + + def _legacy_unsqueeze(self, dim: int) -> Self: + if dim < 0: + dim = self.batch_dims + dim + 1 + + if (dim > self.batch_dims) or (dim < 0): + raise RuntimeError( + f"unsqueezing is allowed for dims comprised between " + f"`-td.batch_dims` and `td.batch_dims` only. Got " + f"dim={dim} with a batch size of {self.batch_size}." + ) + if dim <= self.stack_dim: + stack_dim = self.stack_dim + 1 + else: + dim = dim - 1 + stack_dim = self.stack_dim + return type(self)( + *(tensordict.unsqueeze(dim) for tensordict in self.tensordicts), + stack_dim=stack_dim, + stack_dim_name=self._td_dim_name, + ) + + def _legacy_squeeze(self, dim: int | None = None) -> Self: + """Squeezes all tensors for a dimension comprised in between `-td.batch_dims+1` and `td.batch_dims-1` and returns them in a new tensordict. + + Args: + dim (Optional[int]): dimension along which to squeeze. If dim is None, all singleton dimensions will be squeezed. dim is None by default. + + """ + if dim is None: + size = self.size() + if len(self.size()) == 1 or size.count(1) == 0: + return self + first_singleton_dim = size.index(1) + return self.squeeze(first_singleton_dim).squeeze() + + if dim < 0: + dim = self.batch_dims + dim + + if self.batch_dims and (dim >= self.batch_dims or dim < 0): + raise RuntimeError( + f"squeezing is allowed for dims comprised between 0 and " + f"td.batch_dims only. Got dim={dim} and batch_size" + f"={self.batch_size}." + ) + + if dim >= self.batch_dims or self.batch_size[dim] != 1: + return self + if dim == self.stack_dim: + return self.tensordicts[0] + elif dim < self.stack_dim: + stack_dim = self.stack_dim - 1 + else: + dim = dim - 1 + stack_dim = self.stack_dim + return type(self)( + *(tensordict.squeeze(dim) for tensordict in self.tensordicts), + stack_dim=stack_dim, + stack_dim_name=self._td_dim_name, + ) + + def _unbind(self, dim: int) -> tuple[TensorCollection, ...]: + if dim == self.stack_dim: + return tuple(self.tensordicts) + else: + # return a stack of unbound tensordicts + out = [] + new_dim = dim if dim < self.stack_dim else dim - 1 + new_stack_dim = ( + self.stack_dim if dim > self.stack_dim else self.stack_dim - 1 + ) + for td in self.tensordicts: + out.append(td._unbind(new_dim)) + return tuple( + self.lazy_stack(vals, new_stack_dim) for vals in _zip_strict(*out) + ) + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + if dim == self.stack_dim: + for source, tensordict_dest in _zip_strict(list_item, self.tensordicts): + tensordict_dest.update_(source) + else: + for i, td in enumerate(list_item): + idx = (slice(None),) * dim + (i,) + self.update_at_(td, idx) + return self + + def _maybe_get_list(self, key): + vals = [] + for td in self.tensordicts: + if isinstance(td, LazyStackedTensorDict): + val = td._maybe_get_list(key) + else: + val = td._get_str(key, None) + if _is_tensor_collection(type(val)): + return self._get_str(key, NO_DEFAULT) + elif val is None: + return None + vals.append(val) + return vals + + def get( + self, + key: NestedKey, + *args, + as_list: bool = False, + as_padded_tensor: bool = False, + as_nested_tensor: bool = False, + padding_side: str = "right", + layout: torch.layout = None, + padding_value: float | int | bool = 0.0, + **kwargs, + ) -> CompatibleType: + """Gets the value stored with the input key. + + Args: + key (str, tuple of str): key to be queried. If tuple of str it is + equivalent to chained calls of getattr. + default: default value if the key is not found in the tensordict. Defaults to ``None``. + + .. warning:: + Previously, if a key was not present in the tensordict and no default + was passed, a `KeyError` was raised. From v0.7, this behaviour has been changed + and a `None` value is returned instead (in accordance with the what dict.get behavior). + To adopt the old behavior, set the environment variable `export TD_GET_DEFAULTS_TO_NONE='0'` or call + :func`~tensordict.set_get_defaults_to_none(False)`. + + Keyword Args: + as_list (bool, optional): if ``True``, ragged tensors will be returned as list. + Exclusive with `as_padded_tensor` and `as_nested_tensor`. + Defaults to ``False``. + as_padded_tensor (bool, optional): if ``True``, ragged tensors will be returned as padded tensors. + The padding value can be controlled via the `padding_value` keyword argument, and the padding + side via the `padding_side` argument. + Exclusive with `as_list` and `as_nested_tensor`. + Defaults to ``False``. + as_nested_tensor (bool, optional): if ``True``, ragged tensors will be returned as list. + Exclusive with `as_list` and `as_padded_tensor`. + The layout can be controlled via the `torch.layout` argument. + Defaults to ``False``. + layout (torch.layout, optional): the layout when `as_nested_tensor=True`. + padding_side (str): The side of padding. Must be `"left"` or `"right"`. Defaults to `"right"`. + padding_value (scalar or bool, optional): The padding value. Defaults to 0.0. + + Examples: + >>> from tensordict import TensorDict, lazy_stack + >>> import torch + >>> td = lazy_stack([ + ... TensorDict({"x": torch.ones(1,)}), + ... TensorDict({"x": torch.ones(2,) * 2}), + ... ]) + >>> td.get("x", as_nested_tensor=True) + NestedTensor(size=(2, j1), offsets=tensor([0, 1, 3]), contiguous=True) + >>> td.get("x", as_padded_tensor=True) + tensor([[1., 0.], + [2., 2.]]) + + """ + return super().get( + key, + *args, + as_list=as_list, + as_padded_tensor=as_padded_tensor, + as_nested_tensor=as_nested_tensor, + padding_side=padding_side, + layout=layout, + padding_value=padding_value, + **kwargs, + ) + + @cache # noqa: B019 + def _get_str( + self, + key: NestedKey, + default: Any = NO_DEFAULT, + *, + as_list: bool = False, + as_padded_tensor: bool = False, + as_nested_tensor: bool = False, + padding_side: str = "right", + layout: torch.layout = None, + padding_value: float | int | bool = 0.0, + ) -> CompatibleType: + # we can handle the case where the key is a tuple of length 1 + tensors = [] + for td in self.tensordicts: + # tensors.append(td._get_str(key, default=default)) + tensors.append( + td._get_str( + key, + default=default, + as_list=as_list, + as_padded_tensor=as_padded_tensor, + as_nested_tensor=as_nested_tensor, + padding_side=padding_side, + layout=layout, + padding_value=padding_value, + ) + ) + if ( + tensors[-1] is default + and not isinstance(default, torch.Tensor) + and not is_tensor_collection(default) + ): + # then we consider this default as non-stackable and return prematurly + return default + if as_list: + # Only return as list if all items are plain tensors or non-tensor values + # If they are TensorDicts/tensor collections, we need to stack them properly + if all( + isinstance(item, torch.Tensor) or not is_tensor_collection(item) + for item in tensors + ): + return tensors + # Otherwise, pass to lazy_stack to handle nested structures + if as_nested_tensor: + # If all collected items are plain tensors, convert directly to nested tensor + if all(isinstance(item, torch.Tensor) for item in tensors): + if layout is None: + layout = torch.jagged + return torch.nested.as_nested_tensor(tensors, layout=layout) + # Otherwise, pass to lazy_stack to handle nested structures + if as_padded_tensor: + # If all collected items are plain tensors, convert directly to padded tensor + if all(isinstance(item, torch.Tensor) for item in tensors): + return pad_sequence( + tensors, + padding_value=padding_value, + padding_side=padding_side, + batch_first=True, + ) + # Otherwise, pass to lazy_stack to handle nested structures + try: + out = self.lazy_stack( + tensors, + self.stack_dim, + stack_dim_name=self._td_dim_name, + as_list=as_list, + as_padded_tensor=as_padded_tensor, + as_nested_tensor=as_nested_tensor, + padding_side=padding_side, + layout=layout, + padding_value=padding_value, + ) + if _is_tensor_collection(type(out)): + if isinstance(out, LazyStackedTensorDict): + # then it's a LazyStackedTD + out.hook_out = self.hook_out + out.hook_in = self.hook_in + out._is_vmapped = self._is_vmapped + incr = 0 if not self._is_vmapped else 1 + out._batch_size = ( + self._batch_size + + out.batch_size[(len(self._batch_size) + incr) :] + ) + elif is_tensorclass(out): + # then it's a tensorclass + out._tensordict.hook_out = self.hook_out + out._tensordict.hook_in = self.hook_in + out._tensordict._is_vmapped = self._is_vmapped + incr = 0 if not self._is_vmapped else 1 + out._tensordict._batch_size = ( + self._batch_size + + out._tensordict.batch_size[(len(self._batch_size) + incr) :] + ) + else: + raise RuntimeError + elif self.hook_out is not None: + out = self.hook_out(out) + return out + except RuntimeError as err: + if "stack expects each tensor to be equal size" in str(err): + shapes = {_shape(tensor) for tensor in tensors} + raise RuntimeError( + f"Found more than one unique shape in the tensors to be " + f"stacked ({shapes}). This is likely due to a modification " + f"of one of the stacked TensorDicts, where a key has been " + f"updated/created with an uncompatible shape. If the entries " + f"are intended to have a different shape, use the get_nestedtensor " + f"method instead." + ) + else: + raise err + + def _get_tuple(self, key, default, **kwargs): + first = self._get_str(key[0], None, **kwargs) + if first is None: + return self._default_get(key[0], default) + if len(key) == 1: + return first + try: + return first._get_tuple(key[1:], default=default, **kwargs) + except AttributeError as err: + if "has no attribute" in str(err): + raise ValueError( + f"Expected a TensorDictBase instance but got {type(first)} instead" + f" for key '{key[1:]}' in tensordict:\n{self}." + ) + + @classmethod + def lazy_stack( + cls, + items: Sequence[TensorDictBase], + dim: int = 0, + *, + device: DeviceType | None = None, + out: T | None = None, + stack_dim_name: str | None = None, + strict_shape: bool = False, + as_list: bool = False, + as_padded_tensor: bool = False, + as_nested_tensor: bool = False, + padding_side: str = "right", + layout: torch.layout | None = None, + padding_value: float | int | bool = 0.0, + ) -> Self: # noqa: D417 + """Stacks tensordicts in a LazyStackedTensorDict. + + Args: + items (Sequence of TensorDictBase instances): A sequence of TensorDictBase + instances to stack. + dim (int, optional): the dim along which to perform the lazy stack. + Defaults to 0. + + Keyword Args: + device (torch.device, optional): a device to set in the `LazyStackedTensorDict` + in case it cannot be inferred from the tensordict list (e.g., the list is empty). + out (TensorDictBase, optional): a `LazyStackedTensorDict` where to write the data. + stack_dim_name (str, optional): a name for the stacked dimension. + strict_shape (bool, optional): if ``True``, every tensordict's shapes must match. + Defaults to ``False``. + as_list (bool, optional): if ``True``, ragged tensors will be returned as list. + Exclusive with `as_padded_tensor` and `as_nested_tensor`. + Defaults to ``False``. + as_padded_tensor (bool, optional): if ``True``, ragged tensors will be returned as padded tensors. + The padding value can be controlled via the `padding_value` keyword argument, and the padding + side via the `padding_side` argument. + Exclusive with `as_list` and `as_nested_tensor`. + Defaults to ``False``. + as_nested_tensor (bool, optional): if ``True``, ragged tensors will be returned as list. + Exclusive with `as_list` and `as_padded_tensor`. + The layout can be controlled via the `torch.layout` argument. + Defaults to ``False``. + layout (torch.layout, optional): the layout when `as_nested_tensor=True`. + padding_side (str): The side of padding. Must be `"left"` or `"right"`. Defaults to `"right"`. + padding_value (scalar or bool, optional): The padding value. Defaults to 0.0. + + """ + if not items: + raise RuntimeError("items cannot be empty") + + if all(isinstance(item, torch.Tensor) for item in items): + # This must be implemented here and not in _get_str because we want to leverage this check + special_return = sum((as_list, as_padded_tensor, as_nested_tensor)) + if special_return > 1: + raise TypeError( + "as_list, as_padded_tensor and as_nested_tensor are exclusive." + ) + elif special_return: + if as_padded_tensor: + return pad_sequence( + items, + padding_value=padding_value, + padding_side=padding_side, + batch_first=True, + ) + if as_nested_tensor: + if layout is None: + layout = torch.jagged + return torch.nested.as_nested_tensor(items, layout=layout) + if as_list: + return items + try: + return torch.stack(items, dim=dim, out=out) + except RuntimeError as err: + raise RuntimeError( + "Failed to stack tensors within a tensordict. You can use nested tensors, " + "padded tensors or return lists via specialized keyword arguments. " + "Check the TensorDict.lazy_stack documentation!" + ) from err + if all(is_non_tensor(tensordict) for tensordict in items): + return items[0]._stack_non_tensor(items, dim=dim) + if all( + is_tensorclass(item) and type(item) == type(items[0]) # noqa: E721 + for item in items + ): + lazy_stack = cls.lazy_stack( + [item._tensordict for item in items], + dim=dim, + out=out, + stack_dim_name=stack_dim_name, + ) + # we take the first non_tensordict by convention + return type(items[0])._from_tensordict( + tensordict=lazy_stack, non_tensordict=items[0]._non_tensordict + ) + + batch_size = items[0].batch_size + if dim < 0: + dim = len(batch_size) + dim + 1 + + if strict_shape: + for td in items[1:]: + if td.batch_size != items[0].batch_size: + raise RuntimeError( + "stacking tensordicts requires them to have congruent batch sizes, " + f"got td1.batch_size={td.batch_size} and td2.batch_size=" + f"{items[0].batch_size}" + ) + + if out is None: + # We need to handle tensordicts with exclusive keys and tensordicts with + # mismatching shapes. + # The first case is handled within _check_keys which fails if keys + # don't match exactly. + # The second requires a check over the tensor shapes. + return LazyStackedTensorDict( + *items, + stack_dim=dim, + stack_dim_name=stack_dim_name, + strict_shape=strict_shape, + device=device, + ) + else: + batch_size = list(batch_size) + batch_size.insert(dim, len(items)) + batch_size = torch.Size(batch_size) + + if out.batch_size != batch_size: + raise RuntimeError( + "out.batch_size and stacked batch size must match, " + f"got out.batch_size={out.batch_size} and batch_size" + f"={batch_size}" + ) + + try: + out._stack_onto_(items, dim) + except KeyError as err: + raise err + return out + + @classmethod + def maybe_dense_stack( + cls, + items: Sequence[TensorDictBase], + dim: int = 0, + out: T | None = None, + strict: bool = False, + ) -> Self: + """Stacks tensors or tensordicts densly if possible, or onto a LazyStackedTensorDict otherwise. + + Examples: + >>> td0 = TensorDict({"a": 0}, []) + >>> td1 = TensorDict({"b": 0}, []) + >>> LazyStackedTensorDict.maybe_dense_stack([td0, td0]) # returns a TensorDict with shape [2] + >>> LazyStackedTensorDict.maybe_dense_stack([td0, td1]) # returns a LazyStackedTensorDict with shape [2] + >>> LazyStackedTensorDict.maybe_dense_stack(list(torch.randn(2))) # returns a torch.Tensor with shape [2] + """ + from ._torch_func import _stack + + return _stack(items, dim=dim, out=out, strict=strict, maybe_dense_stack=True) + + @cache # noqa: B019 + def _add_batch_dim(self, *, in_dim, vmap_level): + if self.is_memmap(): + td = LazyStackedTensorDict.lazy_stack( + [td.cpu().as_tensor() for td in self.tensordicts], 0 + ) + else: + td = self + if in_dim < 0: + in_dim = self.ndim + in_dim + if in_dim == self.stack_dim: + result = self._cached_add_batch_dims( + td, in_dim=in_dim, vmap_level=vmap_level + ) + else: + if in_dim < td.stack_dim: + # then we'll stack along a dim before + stack_dim = td.stack_dim - 1 + else: + in_dim = in_dim - 1 + stack_dim = td.stack_dim + + def addbatchdim(_arg): + return _add_batch_dim(_arg, in_dim, vmap_level) + + tds = [ + td._fast_apply( + addbatchdim, + batch_size=[b for i, b in enumerate(td.batch_size) if i != in_dim], + names=( + [name for i, name in enumerate(td.names) if i != in_dim] + if self._has_names() + else None + ), + ) + for td in td.tensordicts + ] + result = LazyStackedTensorDict(*tds, stack_dim=stack_dim) + if self.is_locked: + result.lock_() + return result + + @classmethod + def _cached_add_batch_dims(cls, td, in_dim, vmap_level): + # we return a stack with hook_out, and hack the batch_size and names + # Per se it is still a LazyStack but the stacking dim is "hidden" from + # the outside + out = td.copy() + + def hook_out(tensor, in_dim=in_dim, vmap_level=vmap_level): + if _is_tensor_collection(type(tensor)): + return tensor._add_batch_dim(in_dim=in_dim, vmap_level=vmap_level) + return _add_batch_dim(tensor, in_dim, vmap_level) + + n = len(td.tensordicts) + + def hook_in( + tensor, + out_dim=in_dim, + batch_size=n, + vmap_level=vmap_level, + ): + if _is_tensor_collection(type(tensor)): + return tensor._remove_batch_dim(vmap_level, batch_size, out_dim) + return _remove_batch_dim(tensor, vmap_level, batch_size, out_dim) + + out.hook_out = hook_out + out.hook_in = hook_in + out._is_vmapped = True + out._batch_size = torch.Size( + [dim for i, dim in enumerate(out._batch_size) if i != out.stack_dim] + ) + return out + + @cache # noqa: B019 + def _remove_batch_dim(self, vmap_level, batch_size, out_dim): + if self.hook_out is not None: + # this is the hacked version. We just need to remove the hook_out and + # reset a proper batch size + result = LazyStackedTensorDict( + *self.tensordicts, + stack_dim=out_dim, + ) + # return self._cache_remove_batch_dim(vmap_level=vmap_level, batch_size=batch_size, out_dim=out_dim) + else: + # we must call _remove_batch_dim on all tensordicts + # batch_size: size of the batch when we unhide it. + # out_dim: dimension where the output will be found + new_batch_size = list(self.batch_size) + new_batch_size.insert(out_dim, batch_size) + new_names = list(self.names) + new_names.insert(out_dim, None) + # rebuild the lazy stack + # the stack dim is the same if the out_dim is past it, but it + # must be incremented by one otherwise. + # In the first case, the out_dim must be decremented by one + if out_dim > self.stack_dim: + stack_dim = self.stack_dim + out_dim = out_dim - 1 + else: + stack_dim = self.stack_dim + 1 + result = LazyStackedTensorDict( + *[ + td._remove_batch_dim( + vmap_level=vmap_level, batch_size=batch_size, out_dim=out_dim + ) + for td in self.tensordicts + ], + stack_dim=stack_dim, + ) + if self.is_locked: + result.lock_() + return result + + @cache # noqa: B019 + def _maybe_remove_batch_dim(self, funcname, vmap_level, batch_size, out_dim): + if self.hook_out is not None: + # this is the hacked version. We just need to remove the hook_out and + # reset a proper batch size + result = LazyStackedTensorDict( + *self.tensordicts, + stack_dim=out_dim, + ) + # return self._cache_remove_batch_dim(vmap_level=vmap_level, batch_size=batch_size, out_dim=out_dim) + else: + # we must call _remove_batch_dim on all tensordicts + # batch_size: size of the batch when we unhide it. + # out_dim: dimension where the output will be found + new_batch_size = list(self.batch_size) + new_batch_size.insert(out_dim, batch_size) + new_names = list(self.names) + new_names.insert(out_dim, None) + # rebuild the lazy stack + # the stack dim is the same if the out_dim is past it, but it + # must be incremented by one otherwise. + # In the first case, the out_dim must be decremented by one + if out_dim > self.stack_dim: + stack_dim = self.stack_dim + out_dim = out_dim - 1 + else: + stack_dim = self.stack_dim + 1 + result = LazyStackedTensorDict( + *[ + td._maybe_remove_batch_dim( + funcname, + vmap_level=vmap_level, + batch_size=batch_size, + out_dim=out_dim, + ) + for td in self.tensordicts + ], + stack_dim=stack_dim, + ) + if self.is_locked: + result.lock_() + return result + + def get_nestedtensor( + self, + key: NestedKey, + default: Any = NO_DEFAULT, + *, + layout: torch.layout | None = None, + ) -> CompatibleType: + """Returns a nested tensor when stacking cannot be achieved. + + Args: + key (NestedKey): the entry to nest. + default (Any, optiona): the default value to return in case the key + isn't in all sub-tensordicts. + + .. note:: + In case the default is a tensor, this method will attempt + the construction of a nestedtensor with it. Otherwise, the default + value will be returned. + + Keyword Args: + layout (torch.layout, optional): the layout for the nested tensor. + + Examples: + >>> td0 = TensorDict({"a": torch.zeros(4), "b": torch.zeros(4)}, []) + >>> td1 = TensorDict({"a": torch.ones(5)}, []) + >>> td = torch.stack([td0, td1], 0) + >>> a = td.get_nestedtensor("a") + >>> # using a tensor as default uses this default to build the nested tensor + >>> b = td.get_nestedtensor("b", default=torch.ones(4)) + >>> assert (a == b).all() + >>> # using anything else as default returns the default + >>> b2 = td.get_nestedtensor("b", None) + >>> assert b2 is None + + """ + # disallow getting nested tensor if the stacking dimension is not 0 + if self.stack_dim != 0: + raise RuntimeError( + "Because nested tensors can only be stacked along their first " + "dimension, LazyStackedTensorDict.get_nestedtensor can only be called " + "when the stack_dim is 0." + ) + + # we can handle the case where the key is a tuple of length 1 + key = _unravel_key_to_tuple(key) + subkey = key[0] + if len(key) > 1: + tensordict = self.get(subkey, default) + if tensordict is default: + return default + return tensordict.get_nestedtensor(key[1:], default=default, layout=layout) + tensors = [td.get(subkey, default=default) for td in self.tensordicts] + if not isinstance(default, torch.Tensor) and any( + tensor is default for tensor in tensors + ): + # we don't stack but return the default + return default + return torch.nested.nested_tensor(tensors, layout=layout) + + def is_contiguous(self) -> bool: + return False + + def contiguous(self) -> Self: + source = {key: value.contiguous() for key, value in self.items()} + batch_size = self.batch_size + device = self.device + out = TensorDict._new_unsafe( + source=source, + batch_size=batch_size, + device=device, + names=self.names if self._has_names() else None, + lock=self.is_locked, + ) + return out + + def densify(self, *, layout: torch.layout = torch.strided): + """Attempts to represent the lazy stack with contiguous tensors (plain tensors or nested). + + Keyword Args: + layout (torch.layout): the layout of the nested tensors, if any. Defaults to + :class:`~torch.strided`. + + """ + result = TensorDict._new_unsafe( + batch_size=self.batch_size, device=self.device, names=self.names + ) + for key in self._exclusive_keys(): + list_of_entries = [ + td._get_str(key, default=None) for td in self.tensordicts + ] + is_tensor = all( + isinstance(item, torch.Tensor) or item is None + for item in list_of_entries + ) + if is_tensor: + shapes = { + tensor.shape if tensor is not None else None + for tensor in list_of_entries + } + if None in shapes: + # There must be at least one non-None value + a_shape = None + while a_shape is None: + a_shape = shapes.pop() + if not a_shape: + raise RuntimeError( + f"Cannot densify a tensordict with values with empty shape and exclusive keys: got shape {a_shape}." + ) + none_shape = a_shape[:-1] + (0,) + for tensor in list_of_entries: + if tensor is not None: + a_tensor = tensor.new_zeros(none_shape) + break + list_of_entries = [ + tensor if tensor is not None else a_tensor + for tensor in list_of_entries + ] + shapes.update({a_shape, none_shape}) + if len(shapes) == 1: + tensor = torch.stack(list_of_entries, self.stack_dim) + else: + if self.stack_dim == 0: + tensor = torch.nested.nested_tensor( + list_of_entries, layout=layout + ) + else: + raise NotImplementedError( + f"stack_dim is {self.stack_dim} but not 0. Densify canot be done." + ) + else: + tensor = self._get_str(key, None) + if tensor is not None: + tensor = tensor.densify(layout=layout) + else: + from tensordict import NonTensorData + + tensor = NonTensorData(None) + result._set_str(key, tensor, validated=True, inplace=False) + return result + + def empty( + self, recurse=False, *, batch_size=None, device=NO_DEFAULT, names=None + ) -> Self: + name = None + if batch_size is not None and ( + self.stack_dim + and batch_size[: self.stack_dim] != self.batch_size[: self.stack_dim] + ): + return TensorDict.empty( + self, + recurse=recurse, + batch_size=batch_size, + device=device if device is not NO_DEFAULT else self.device, + names=names if names is not None else None, + ) + if names is not None: + if len(names) > self.stack_dim: + name = names[self.stack_dim] + names = [name for i, name in enumerate(names) if i != self.stack_dim] + if batch_size is not None: + batch_size = torch.Size( + [b for i, b in enumerate(batch_size) if i != self.stack_dim] + ) + return type(self)( + *[ + td.empty( + recurse=recurse, batch_size=batch_size, device=device, names=names + ) + for td in self.tensordicts + ], + stack_dim=self.stack_dim, + stack_dim_name=name, + ) + + def _clone(self, recurse: bool = True) -> Self: + if recurse: + # This could be optimized using copy but we must be careful with + # metadata (_is_shared etc) + result = type(self)( + *[td._clone() for td in self.tensordicts], + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + ) + else: + result = type(self)( + *[td._clone(recurse=False) for td in self.tensordicts], + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + ) + return result + + def to(self, *args, **kwargs) -> Self: + if kwargs.get("batch_size") is not None: + raise TypeError("Cannot pass batch-size to a LazyStackedTensorDict.") + return super().to(*args, **kwargs) + + def _check_new_batch_size(self, new_size: torch.Size) -> None: + if len(new_size) <= self.stack_dim: + raise RuntimeError( + "Changing the batch_size of a LazyStackedTensorDicts can only " + "be done with sizes that are at least as long as the " + "stacking dimension." + ) + super()._check_new_batch_size(new_size) + + def _change_batch_size(self, new_size: torch.Size) -> None: + self._batch_size = new_size + + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _LazyStackedTensorDictKeysView: + keys = _LazyStackedTensorDictKeysView( + self, + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + return keys + + def values( + self, + include_nested=False, + leaves_only=False, + is_leaf=None, + *, + sort: bool = False, + ): + if is_leaf not in ( + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + ): + yield from super().values( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + else: + for td in self.tensordicts: + yield from td.values( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + def items( + self, + include_nested=False, + leaves_only=False, + is_leaf=None, + *, + sort: bool = False, + ): + if is_leaf not in ( + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + ): + yield from super().items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + else: + for i, td in enumerate(self.tensordicts): + for key, val in td.items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ): + if isinstance(key, str): + key = (str(i), key) + else: + key = (str(i), *key) + yield key, val + + valid_keys = keys + + def non_tensor_items(self, include_nested: bool = False): + """Returns all non-tensor leaves, maybe recursively.""" + items = self.tensordicts[0].non_tensor_items(include_nested=include_nested) + return tuple( + ( + key, + torch.stack( + [val0, *[td.get(key) for td in self.tensordicts[1:]]], + self.stack_dim, + ), + ) + for (key, val0) in items + ) + + def _iterate_over_keys(self) -> None: + # this is about 20x faster than the version above + yield from self._key_list() + + @cache # noqa: B019 + def _key_list(self): + if not self.tensordicts: + return [] + keys = set(self.tensordicts[0].keys()) + for td in self.tensordicts[1:]: + keys = keys.intersection(td.keys()) + return sorted(keys, key=str) + + @lock_blocked + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + key, val = self.tensordicts[0].popitem() + vals = [val] + for i, td in enumerate(self.tensordicts[1:]): + val = td.pop(key, None) + if val is not None: + vals.append(val) + else: + for j in range(i + 1): + self.tensordicts[j].set(key, vals[j]) + raise RuntimeError(f"Could not find key {key} in all tensordicts.") + return key, torch.stack(vals, dim=self.stack_dim) + + def entry_class(self, key: NestedKey) -> type: + data_type = type(self.tensordicts[0].get(key)) + if _is_tensor_collection(data_type): + return LazyStackedTensorDict + return data_type + + def apply_(self, fn: Callable, *others, **kwargs): + others = (other.unbind(self.stack_dim) for other in others) + for td, *_others in _zip_strict(self.tensordicts, *others): + td._fast_apply(fn, *_others, inplace=True, propagate_lock=True, **kwargs) + return self + + def _multithread_apply_nest(self, *args, **kwargs): + if kwargs.get("batch_size") is not None: + raise RuntimeError( + f"batch_size cannot be specified for {type(self).__name__}._multithread_apply_nest." + ) + return super()._multithread_apply_nest(*args, **kwargs) + + def _multithread_apply_flat( + self, + fn: Callable, + *others: T, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + is_leaf: Callable[[Type], bool] | None = None, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + ) -> None: + others = (other.unbind(self.stack_dim) for other in others) + if ( + call_on_nested + and named + and is_leaf + in (_NESTED_TENSORS_AS_LISTS, _NESTED_TENSORS_AS_LISTS_NONTENSOR) + ): + # When calling on nested with name and the name includes the TD index, we + # want to call the function on each td. + # If we were not keeping track of the TD's index, names would be the same for all + # tds and there's a risk that values would collide. + # nested_keys is irrelevant when named + call_on_nested are both true. + for i, (td, *oth) in enumerate(_zip_strict(self.tensordicts, *others)): + key = prefix + (str(i),) + if len(key) == 1: + key = key[0] + futures.append(executor.submit(fn, key, td, *oth)) + local_futures.append(futures[-1]) + else: + for i, (td, *oth) in enumerate(_zip_strict(self.tensordicts, *others)): + local_futures.append([]) + td._multithread_apply_flat( + fn, + *oth, + call_on_nested=call_on_nested, + default=default, + named=named, + nested_keys=nested_keys, + prefix=( + prefix + (str(i),) + if is_leaf + in ( + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + ) + else prefix + ), + is_leaf=is_leaf, + executor=executor, + futures=futures, + local_futures=local_futures[-1], + ) + + def _multithread_rebuild( + self, + *, + # We know batch_size is None, this has been checked earlier + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + out: TensorDictBase | None = None, + filter_empty: bool = False, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + subs_results: Dict[Future, Any] | None = None, + multithread_set: bool = False, # Experimental + **constructor_kwargs, + ) -> None: + if inplace and any( + arg for arg in (batch_size, device, names, constructor_kwargs) + ): + raise ValueError( + "Cannot pass other arguments to LazyStackedTensorDict.apply when inplace=True." + ) + if out is not None: + if not isinstance(out, LazyStackedTensorDict): + raise ValueError( + "out must be a LazyStackedTensorDict instance in lazy_stack.apply(..., out=out)." + ) + out = out.tensordicts + results = [] + for i, (td, local_future) in enumerate( + _zip_strict(self.tensordicts, local_futures) + ): + local_out = out[i] if out is not None else None + # Each local_future points to a list of futures for a single tensordict + local_out = td._multithread_rebuild( + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=checked, + out=local_out, + filter_empty=filter_empty, + executor=executor, + futures=futures, + local_futures=local_future, + subs_results=subs_results, + multithread_set=multithread_set, + ) + results.append(local_out) + if filter_empty and all(r is None for r in results): + return + if not inplace: + out = type(self)( + *results, + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + ) + else: + out = self + if names is not NO_DEFAULT: + out.names = names + return out + + def _apply_nest( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + filter_empty: bool | None = None, + is_leaf: Callable | None = None, + out: TensorDictBase | None = None, + **constructor_kwargs, + ) -> Self | None: + if inplace and any( + arg for arg in (batch_size, device, names, constructor_kwargs) + ): + raise ValueError( + "Cannot pass other arguments to LazyStackedTensorDict.apply when inplace=True. Got args " + f"batch_size={batch_size}, device={device}, names={names}, constructor_kwargs={constructor_kwargs}" + ) + if out is not None: + if not isinstance(out, LazyStackedTensorDict): + raise ValueError( + "out must be a LazyStackedTensorDict instance in lazy_stack.apply(..., out=out)." + ) + out = out.tensordicts + elif batch_size is not None: + # any op that modifies the batch-size will result in a regular TensorDict + batch_size = torch.Size(batch_size) + out = TensorDict._new_unsafe( + {}, + batch_size=batch_size, + device=device if device is not NO_DEFAULT else self.device, + names=names if names else self._maybe_names(), + ) + return TensorDict._apply_nest( + self, + fn, + *others, + batch_size=batch_size, + device=device, + names=names, + checked=checked, + call_on_nested=call_on_nested, + default=default, + named=named, + nested_keys=nested_keys, + prefix=prefix, + inplace=inplace, + filter_empty=filter_empty, + is_leaf=is_leaf, + out=out, + **constructor_kwargs, + ) + + others = (other.unbind(self.stack_dim) for other in others) + if ( + call_on_nested + and named + and is_leaf + in (_NESTED_TENSORS_AS_LISTS, _NESTED_TENSORS_AS_LISTS_NONTENSOR) + ): + # When calling on nested with name and the name includes the TD index, we + # want to call the function on each td. + # If we were not keeping track of the TD's index, names would be the same for all + # tds and there's a risk that values would collide. + # nested_keys is irrelevant when named + call_on_nested are both true. + results = [] + for i, (td, *oth) in enumerate(_zip_strict(self.tensordicts, *others)): + key = prefix + (str(i),) + if len(key) == 1: + key = key[0] + results.append(fn(key, td, *oth)) + else: + results = [ + td._apply_nest( + fn, + *oth, + checked=checked, + device=device, + call_on_nested=call_on_nested, + default=default, + named=named, + nested_keys=nested_keys, + prefix=( + prefix + (str(i),) + if is_leaf + in ( + _NESTED_TENSORS_AS_LISTS, + _NESTED_TENSORS_AS_LISTS_NONTENSOR, + ) + else prefix + ), + inplace=inplace, + filter_empty=filter_empty, + is_leaf=is_leaf, + out=out[i] if out is not None else None, + ) + for i, (td, *oth) in enumerate(_zip_strict(self.tensordicts, *others)) + ] + if all(r is None for r in results) and filter_empty in (None, True): + return + if not inplace: + if not results or any(r is not None for r in results): + try: + out = type(self)( + *results, + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + ) + except Exception as e: + raise RuntimeError( + f"Failed to reconstruct the lazy stack of tensordicts with class: {type(self)}. " + f"One common issue is that the outputs of apply are a mix of None and non-None " + f"values. Check that the outputs of apply() are all None or all non-None. " + f"Otherwise, please report this bug on tensordict github." + ) from e + else: + out = None + else: + out = self + if names is not NO_DEFAULT: + out.names = names + return out + + def _select( + self, + *keys: NestedKey, + inplace: bool = False, + strict: bool = False, + set_shared: bool = True, + ) -> LazyStackedTensorDict: + # the following implementation keeps the hidden keys in the tensordicts + tensordicts = [ + td._select(*keys, inplace=inplace, strict=strict, set_shared=set_shared) + for td in self.tensordicts + ] + if inplace: + return self + result = self._new_lazy_unsafe( + *tensordicts, stack_dim=self.stack_dim, stack_dim_name=self._td_dim_name + ) + return result + + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> LazyStackedTensorDict: + tensordicts = [ + tensordict._exclude(*keys, inplace=inplace, set_shared=set_shared) + for tensordict in self.tensordicts + ] + if inplace: + self.tensordicts = tensordicts + return self + result = type(self)( + *tensordicts, stack_dim=self.stack_dim, stack_dim_name=self._td_dim_name + ) + return result + + def __setitem__(self, index: IndexType, value: T) -> Self: + if isinstance(index, (tuple, str)): + # try: + index_unravel = _unravel_key_to_tuple(index) + if index_unravel: + self._set_tuple( + index_unravel, + value, + inplace=( + BEST_ATTEMPT_INPLACE + if isinstance(self, _SubTensorDict) + else False + ), + validated=False, + non_blocking=False, + ) + return + + if any( + isinstance(sub_index, (list, range, np.ndarray)) for sub_index in index + ): + index = tuple( + ( + torch.as_tensor(sub_index, device=self.device) + if isinstance(sub_index, (list, range, np.ndarray)) + else sub_index + ) + for sub_index in index + ) + + if index is Ellipsis or (isinstance(index, tuple) and Ellipsis in index): + index = convert_ellipsis_to_idx(index, self.batch_size) + elif isinstance(index, (list, range)): + index = torch.as_tensor(index, device=self.device) + elif isinstance(index, (type(None), bool)) or ( + isinstance(index, torch.Tensor) + and index.shape == () + and index.dtype == torch.bool + and index.all() + ): + self.unsqueeze(0).update(value) + return self + + if is_tensor_collection(value) or isinstance(value, dict): + indexed_bs = _getitem_batch_size(self.batch_size, index) + if isinstance(value, dict): + value = TensorDict._new_unsafe( + value, batch_size=indexed_bs, device=self.device + ) + if value.batch_size != indexed_bs: + # try to expand + try: + value = value.expand(indexed_bs) + except RuntimeError as err: + raise RuntimeError( + f"indexed destination TensorDict batch size is {indexed_bs} " + f"(batch_size = {self.batch_size}, index={index}), " + f"which differs from the source batch size {value.batch_size}" + ) from err + split_index = self._split_index(index) + converted_idx = split_index["index_dict"] + num_single = split_index["num_single"] + isinteger = split_index["isinteger"] + has_bool = split_index["has_bool"] + num_squash = split_index.get("num_squash", 0) + num_none = split_index.get("num_none", 0) + is_nd_tensor = split_index.get("is_nd_tensor", False) + if isinteger: + # this will break if the index along the stack dim is [0] or :1 or smth + for i, _idx in converted_idx.items(): + if _idx == (): + self.tensordicts[i].update(value, inplace=True) + else: + self.tensordicts[i][_idx] = value + return self + if is_nd_tensor: + unbind_dim = self.stack_dim - num_single + num_none - num_squash + + # converted_idx is a nested list with (int, index) items + def assign(converted_idx, value=value): + value = value.unbind(unbind_dim) + for i, item in enumerate(converted_idx): + if isinstance(item, list): + assign(item) + else: + stack_item, idx = item + if idx == (): + self.tensordicts[stack_item] = value[i] + else: + self.tensordicts[stack_item][idx] = value[i] + + assign(converted_idx) + return self + if not has_bool: + unbind_dim = self.stack_dim - num_single + num_none - num_squash + value_unbind = value.unbind(unbind_dim) + for (i, _idx), _value in _zip_strict( + converted_idx.items(), + value_unbind, + ): + if _idx == (): + self.tensordicts[i].update(_value, inplace=True) + else: + self.tensordicts[i][_idx] = _value + else: + # we must split, not unbind + mask_unbind = split_index["individual_masks"] + split_dim = split_index["split_dim"] + splits = [_mask_unbind.sum().item() for _mask_unbind in mask_unbind] + value_unbind = value.split(splits, split_dim) + if mask_unbind[0].ndim == 0: + # we can return a stack + for (i, _idx), mask, _value in _zip_strict( + converted_idx.items(), + mask_unbind, + value_unbind, + ): + if mask.any(): + self.tensordicts[i][_idx] = _value + else: + for (i, _idx), _value in _zip_strict( + converted_idx.items(), value_unbind + ): + self_idx = (slice(None),) * split_index["mask_loc"] + (i,) + self[self_idx][_idx] = _value + else: + for key in self.keys(): + self.set_at_(key, value, index) + + def __contains__(self, item: IndexType) -> bool: + if isinstance(item, TensorDictBase): + return any(item is td for td in self.tensordicts) + return super().__contains__(item) + + def __getitem__(self, index: IndexType) -> Self | Tensor | TensorCollection | Any: + if isinstance(index, (tuple, str)): + index_key = _unravel_key_to_tuple(index) + if index_key: + leaf = self._get_tuple(index_key, NO_DEFAULT) + if is_non_tensor(leaf): + # Only lazy stacks of non tensors are actually tensordict instances + if isinstance(leaf, TensorDictBase): + return leaf.tolist(as_linked_list=True) + return leaf.data + return leaf + if isinstance(index, (type(None), bool)) or ( + isinstance(index, torch.Tensor) + and index.shape == () + and index.dtype == torch.bool + and index.all() + ): + return self.unsqueeze(0) + split_index = self._split_index(index) + converted_idx = split_index["index_dict"] + isinteger = split_index["isinteger"] + has_bool = split_index["has_bool"] + is_nd_tensor = split_index["is_nd_tensor"] + num_single = split_index.get("num_single", 0) + num_none = split_index.get("num_none", 0) + num_squash = split_index.get("num_squash", 0) + if has_bool: + mask_unbind = split_index["individual_masks"] + cat_dim = split_index["mask_loc"] - num_single + result = [] + if mask_unbind[0].ndim == 0: + # we can return a stack + for (i, _idx), mask in _zip_strict(converted_idx.items(), mask_unbind): + if mask.any(): + if mask.all() and self.tensordicts[i].ndim == 0: + result.append(self.tensordicts[i]) + else: + result.append(self.tensordicts[i][_idx]) + result[-1] = result[-1].squeeze(cat_dim) + if not result: + batch_size = _getitem_batch_size(self.batch_size, index) + else: + batch_size = None + return self._new_lazy_unsafe( + *result, + stack_dim=cat_dim, + device=self.device, + names=self.names, + batch_size=batch_size, + ) + else: + for i, _idx in converted_idx.items(): + self_idx = (slice(None),) * split_index["mask_loc"] + (i,) + result.append(self[self_idx][_idx]) + return torch.cat(result, cat_dim) + elif is_nd_tensor: + new_stack_dim = self.stack_dim - num_single + num_none + + def recompose(converted_idx, stack_dim=new_stack_dim): + stack = [] + for item in converted_idx: + if isinstance(item, list): + stack.append(recompose(item, stack_dim=stack_dim)) + else: + stack_elt, idx = item + if idx != (): + stack.append(self.tensordicts[stack_elt][idx]) + else: + stack.append(self.tensordicts[stack_elt]) + + # TODO: this produces multiple dims with the same name + result = LazyStackedTensorDict.lazy_stack( + stack, stack_dim, stack_dim_name=self._td_dim_name + ) + if self.is_locked: + result.lock_() + return result + + return recompose(converted_idx) + else: + if isinteger: + for ( + i, + _idx, + ) in ( + converted_idx.items() + ): # for convenience but there's only one element + result = self.tensordicts[i] + if _idx is not None and _idx != (): + result = result[_idx] + return result + else: + result = [] + new_stack_dim = self.stack_dim - num_single + num_none - num_squash + for i, _idx in converted_idx.items(): + if _idx == (): + result.append(self.tensordicts[i]) + else: + result.append(self.tensordicts[i][_idx]) + result = LazyStackedTensorDict.lazy_stack( + result, new_stack_dim, stack_dim_name=self._td_dim_name + ) + if self.is_locked: + result.lock_() + return result + + def __eq__(self, other): + return self._dispatch_comparison(other, "__eq__", "__eq__", default=False) + + def __ne__(self, other): + return self._dispatch_comparison(other, "__ne__", "__ne__", default=True) + + def __or__(self, other): + return self._dispatch_comparison(other, "__or__", "__or__", default=NO_DEFAULT) + + def __xor__(self, other): + return self._dispatch_comparison( + other, "__xor__", "__xor__", default=NO_DEFAULT + ) + + def __ge__(self, other): + return self._dispatch_comparison(other, "__ge__", "__le__", default=NO_DEFAULT) + + def __gt__(self, other): + return self._dispatch_comparison(other, "__gt__", "__lt__", default=NO_DEFAULT) + + def __le__(self, other): + return self._dispatch_comparison(other, "__le__", "__ge__", default=NO_DEFAULT) + + def __lt__(self, other): + return self._dispatch_comparison(other, "__lt__", "__gt__", default=NO_DEFAULT) + + def _dispatch_comparison(self, other, comparison_str, inverse_str, default): + if is_tensorclass(other): + return getattr(other, inverse_str)(self) + if isinstance(other, (dict,)): + # we may want to broadcast it instead + other = TensorDict.from_dict(other, batch_size=self.batch_size) + if _is_tensor_collection(type(other)): + if other.batch_size != self.batch_size: + if self.ndim < other.ndim: + self_expand = self.expand(other.batch_size) + elif self.ndim > other.ndim: + other = other.expand(self.batch_size) + self_expand = self + else: + raise RuntimeError( + f"Could not compare tensordicts with shapes {self.shape} and {other.shape}" + ) + else: + self_expand = self + out = [] + for td0, td1 in _zip_strict( + self_expand.tensordicts, other.unbind(self_expand.stack_dim) + ): + out.append(getattr(td0, comparison_str)(td1)) + return LazyStackedTensorDict.lazy_stack(out, self.stack_dim) + if isinstance(other, (numbers.Number, Tensor)): + return LazyStackedTensorDict.lazy_stack( + [getattr(td, comparison_str)(other) for td in self.tensordicts], + self.stack_dim, + ) + if default is NO_DEFAULT: + raise ValueError( + f"Incompatible value {type(other)} for op {comparison_str}." + ) + return default + + def _cast_reduction( + self, + *, + reduction_name, + dim=NO_DEFAULT, + keepdim=NO_DEFAULT, + tuple_ok=True, + further_reduce: bool, + **kwargs, + ): + if further_reduce: + if dim is NO_DEFAULT: + # It is not very memory-efficient to do this, but it's the easiest to cover all use cases + agglomerate = [ + val.contiguous().flatten() + for val in self._values_list( + True, True, is_leaf=_NESTED_TENSORS_AS_LISTS + ) + ] + agglomerate = torch.cat(agglomerate, dim=-1) + return getattr(torch, reduction_name)(agglomerate, **kwargs) + elif dim == "feature": + + def proc_val(val): + val = val.contiguous() + if val.ndim > self.ndim: + val = val.flatten(self.ndim, -1) + else: + val = val.unsqueeze(-1) + return val + + agglomerate = [ + proc_val(val) + for val in self.values( + True, + True, + ) + ] + dim = -1 + cat_dim = -1 + keepdim = False + else: + agglomerate = [ + val.contiguous().unsqueeze(self.stack_dim) + for val in self.values(True, True) + ] + cat_dim = self.stack_dim + agglomerate = torch.cat(agglomerate, dim=cat_dim) + return getattr(torch, reduction_name)( + agglomerate, dim=dim, keepdim=keepdim, **kwargs + ) + + try: + td: TensorDict = self.to_tensordict() + except Exception: + raise RuntimeError( + f"{reduction_name} requires this object to be cast to a regular TensorDict. " + f"If you need {type(self).__name__} to support {reduction_name}, help us by filing an issue" + f" on github!" + ) + return td._cast_reduction( + reduction_name=reduction_name, + dim=dim, + keepdim=keepdim, + tuple_ok=tuple_ok, + further_reduce=further_reduce, + **kwargs, + ) + + def all(self, dim: int | None = None) -> bool | TensorDictBase: + if dim is not None and (dim >= self.batch_dims or dim < -self.batch_dims): + raise RuntimeError( + "dim must be greater than or equal to -tensordict.batch_dims and " + "smaller than tensordict.batch_dims" + ) + if dim is not None: + # TODO: we need to adapt this to LazyStackedTensorDict too + if dim < 0: + dim = self.batch_dims + dim + return TensorDict( + source={key: value.all(dim=dim) for key, value in self.items()}, + batch_size=[b for i, b in enumerate(self.batch_size) if i != dim], + device=self.device, + ) + return all(value.all() for value in self.tensordicts) + + def any(self, dim: int | None = None) -> bool | TensorDictBase: + if dim is not None and (dim >= self.batch_dims or dim < -self.batch_dims): + raise RuntimeError( + "dim must be greater than or equal to -tensordict.batch_dims and " + "smaller than tensordict.batch_dims" + ) + if dim is not None: + # TODO: we need to adapt this to LazyStackedTensorDict too + if dim < 0: + dim = self.batch_dims + dim + return TensorDict( + source={key: value.any(dim=dim) for key, value in self.items()}, + batch_size=[b for i, b in enumerate(self.batch_size) if i != dim], + device=self.device, + ) + return any(value.any() for value in self.tensordicts) + + def _send( + self, + dst: int, + _tag: int = -1, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + ) -> int: + for td in self.tensordicts: + _tag = td._send(dst, _tag=_tag, pseudo_rand=pseudo_rand, group=group) + return _tag + + def _isend( + self, + dst: int, + _tag: int = -1, + _futures: list[torch.Future] | None = None, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + return_early: bool = False, + ) -> int | list[torch.Future]: + if _futures is None: + is_root = True + _futures = [] + else: + is_root = False + for td in self.tensordicts: + _tag = td._isend( + dst, _tag=_tag, pseudo_rand=pseudo_rand, _futures=_futures, group=group + ) + if is_root and not return_early: + for _future in _futures: + _future.wait() + elif is_root and return_early: + return _futures + return _tag + + def _recv( + self, + src: int, + _tag: int = -1, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + ) -> int: + for td in self.tensordicts: + _tag = td._recv(src, _tag=_tag, pseudo_rand=pseudo_rand, group=group) + return _tag + + def _irecv( + self, + src: int, + return_premature: bool = False, + _tag: int = -1, + _future_list: list[torch.Future] | None = None, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: + root = False + if _future_list is None: + _future_list = [] + root = True + for td in self.tensordicts: + _tag, _future_list = td._irecv( + src=src, + return_premature=return_premature, + _tag=_tag, + _future_list=_future_list, + pseudo_rand=pseudo_rand, + group=group, + ) + + if not root: + return _tag, _future_list + elif return_premature: + return _future_list + else: + for future in _future_list: + future.wait() + return + + @lock_blocked + def del_(self, key: NestedKey, **kwargs: Any) -> Self: + # Use check-before-delete pattern for torch.compile compatibility + key_tuple = _unravel_key_to_tuple(key) + is_nested = len(key_tuple) > 1 + ids = set() + cur_len = len(ids) + is_deleted = False + for td in self.tensordicts: + # checking that the td has not been processed yet. + # It could be that not all sub-tensordicts have the appropriate + # entry but one must have it (or an error is thrown). + tdid = id(td) + ids.add(tdid) + new_cur_len = len(ids) + if new_cur_len == cur_len: + continue + cur_len = new_cur_len + if key_tuple[0] in td.keys(): + if is_nested: + # For nested keys, check the full path exists + if key in td.keys(True): + td.del_(key, **kwargs) + is_deleted = True + else: + td.del_(key, **kwargs) + is_deleted = True + if not is_deleted: + raise KeyError(f"Key {key} not found in any of the stacked tensordicts.") + return self + + def pop(self, key: NestedKey, default: Any = NO_DEFAULT) -> CompatibleType: + # using try/except for get/del is suboptimal, but + # this is faster that checkink if key in self keys + key = _unravel_key_to_tuple(key) + if len(key) == 1: + key = key[0] + present = False + if isinstance(key, tuple): + if key in self.keys(True): + present = True + value = self._get_tuple(key, NO_DEFAULT) + elif key in self.keys(): + present = True + value = self._get_str(key, NO_DEFAULT) + if present: + self.del_(key) + elif default is not NO_DEFAULT: + value = default + else: + raise KeyError( + f"You are trying to pop key `{key}` which is not in dict " + f"without providing default value." + ) + return value + + def share_memory_(self) -> Self: + for td in self.tensordicts: + td.share_memory_() + self.lock_() + return self + + def detach_(self) -> Self: + for td in self.tensordicts: + td.detach_() + return self + + def _memmap_( + self, + *, + prefix: str | None = None, + copy_existing: bool = False, + executor=None, + futures=None, + inplace=True, + like=False, + share_non_tensor, + robust_key: bool = False, + existsok, + ) -> Self: + if prefix is not None: + prefix = Path(prefix) + + def save_metadata(prefix=prefix, self=self): + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + with open(prefix / "meta.json", "wb") as f: + from tensordict.utils import json_dumps + + json_str = json_dumps( + {"_type": str(type(self)), "stack_dim": self.stack_dim} + ) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + f.write(json_str.encode("utf-8")) + else: + f.write(json_str) + + if executor is None: + save_metadata() + else: + futures.append(executor.submit(save_metadata)) + + results = [] + for i, td in enumerate(self.tensordicts): + results.append( + td._memmap_( + prefix=(prefix / str(i)) if prefix is not None else None, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + ) + if not inplace: + results = LazyStackedTensorDict.lazy_stack(results, dim=self.stack_dim) + else: + results = self + results._device = torch.device("cpu") + return results + + @classmethod + def _load_memmap( + cls, + prefix: str, + metadata: dict, + device: torch.device | None = None, + *, + out=None, + robust_key: bool = False, + **kwargs, + ) -> LazyStackedTensorDict: + tensordicts = [] + i = 0 + stack_dim = metadata["stack_dim"] + if out is not None: + out = out.unbind(stack_dim) + while (prefix / str(i)).exists(): + tensordicts.append( + TensorDict.load_memmap( + prefix / str(i), + device=device, + **kwargs, + non_blocking=True, + out=out[i] if out is not None else None, + robust_key=robust_key, + ) + ) + i += 1 + return cls(*tensordicts, stack_dim=stack_dim, **kwargs) + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for LazyStack as " + "it can't return a contiguous view of the lazy stacked tensors. " + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for LazyStack as " + "it can't return a contiguous view of the lazy stacked tensors. " + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_tensor( + self, key: NestedKey, tensor: torch.Tensor, *, copy_data: bool = True + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for LazyStack as " + "it can't return a contiguous view of the lazy stacked tensors. " + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def expand(self, *args: int, inplace: bool = False) -> Self: + if len(args) == 1 and isinstance(args[0], Sequence): + shape = tuple(args[0]) + else: + shape = args + # We need to reprod the elements if shape is (1,) + if shape[self.stack_dim - self.ndim] != len(self.tensordicts): + if len(self.tensordicts) == 1: + tds = [ + self.tensordicts[0].copy() + for _ in range(shape[self.stack_dim - self.ndim]) + ] + else: + raise ValueError( + f"Cannot expand LazyStackedTensorDict with shape {shape} because the number of elements in the stack " + f"({len(self.tensordicts)}) does not match the number of elements in the new shape ({shape[self.stack_dim - self.ndim]})." + ) + else: + tds = self.tensordicts + stack_dim = len(shape) + self.stack_dim - self.ndimension() + new_shape_tensordicts = [v for i, v in enumerate(shape) if i != stack_dim] + tensordicts = [td.expand(new_shape_tensordicts) for td in tds] + if inplace: + self.tensordicts = tensordicts + self.stack_dim = stack_dim + return self + return LazyStackedTensorDict.maybe_dense_stack(tensordicts, dim=stack_dim) + + @lock_blocked + def update( + self, + input_dict_or_td: T, + clone: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + non_blocking: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + **kwargs: Any, + ) -> Self: + # This implementation of update is compatible with exclusive keys + # as well as vmapped lazy stacks. + # We iterate over the tensordicts rather than iterating over the keys, + # which requires stacking and unbinding but is also not robust to missing keys. + if input_dict_or_td is self: + # no op + return self + if is_leaf is None: + is_leaf = _is_leaf_nontensor + if isinstance(input_dict_or_td, dict): + input_dict_or_td = TensorDict.from_dict( + input_dict_or_td, batch_size=self.batch_size + ) + + if keys_to_update is not None: + keys_to_update = unravel_key_list(keys_to_update) + if len(keys_to_update) == 0: + return self + + if ( + isinstance(input_dict_or_td, LazyStackedTensorDict) + and input_dict_or_td.stack_dim == self.stack_dim + ): + tds = list(self.tensordicts) + if len(input_dict_or_td.tensordicts) > len(self.tensordicts): + tds.extend( + [td.copy() for td in input_dict_or_td.tensordicts[len(tds) :]] + ) + elif len(input_dict_or_td.tensordicts) != len(self.tensordicts): + if update_batch_size: + keys_source = set(input_dict_or_td.keys(True)) + keys_dest = set(self.keys(True)) + if not keys_dest.issubset(keys_source): + raise RuntimeError( + "Some keys of the dest tensordict are not present in the source " + "during update with mismatching batch-size. " + f"batch_size of source={input_dict_or_td.batch_size}, batch_size of dest={self.batch_size}, " + f"keys in dest but not in source: {{{keys_dest - keys_source}}}." + ) + self.__init__( + *input_dict_or_td.tensordicts, + stack_dim=self.stack_dim, + hook_out=self.hook_out, + hook_in=self.hook_in, + stack_dim_name=self._td_dim_name, + ) + return self + + else: + raise ValueError( + "cannot update stacked tensordicts with different shapes when update_batch_size=False." + ) + for td_dest, td_source in _zip_strict(tds, input_dict_or_td.tensordicts): + td_dest.update( + td_source, + clone=clone, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + **kwargs, + ) + return self + + if self.hook_in is not None: + self_upd = self.hook_in(self) + input_dict_or_td = self.hook_in(input_dict_or_td) + else: + self_upd = self + # Then we can decompose the tensordict along its stack dim + if input_dict_or_td.ndim <= self_upd.stack_dim or input_dict_or_td.batch_size[ + self_upd.stack_dim + ] != len(self_upd.tensordicts): + # We receive a tensordict with a different batch-size than self. + # If update_batch_size is True, we can just convert the input tensordict to a lazy stack and update. + # This will change self, most likely removing a bunch of sub-tensordicts but we're good because the + # user is interested in modifying the batch-size. + # If update_batch_size is False, we need to try to change the batch-size of the input tensordict. + # That can only be done in restricted cases, so we raise an error if the batch-size of self (which must + # remain unchanged) is incompatible with the content of the input tensordict. + if update_batch_size: + return self.update( + input_dict_or_td.to_lazystack(self.stack_dim), + clone=clone, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + **kwargs, + ) + # if the batch-size does not permit unbinding, let's first try to reset the batch-size. + input_dict_or_td = input_dict_or_td.copy() + batch_size = self_upd.batch_size + if self_upd.hook_out is not None: + batch_size = list(batch_size) + batch_size.insert(self_upd.stack_dim, len(self_upd.tensordicts)) + try: + input_dict_or_td.batch_size = batch_size + except RuntimeError as err: + raise ValueError( + "cannot update stacked tensordicts with different shapes." + ) from err + for td_dest, td_source in _zip_strict( + self_upd.tensordicts, input_dict_or_td.unbind(self_upd.stack_dim) + ): + td_dest.update( + td_source, + clone=clone, + keys_to_update=keys_to_update, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + **kwargs, + ) + if self.hook_out is not None: + self_upd = self.hook_out(self_upd) + else: + self_upd = self + return self_upd + + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorDictBase, + clone: bool = False, + *, + non_blocking: bool = False, + **kwargs: Any, + ) -> Self: + if input_dict_or_td is self: + # no op + return self + if not is_tensor_collection(input_dict_or_td): + input_dict_or_td = TensorDict.from_dict( + input_dict_or_td, batch_dims=self.batch_dims + ) + if input_dict_or_td.batch_dims <= self.stack_dim: + raise RuntimeError( + f"Built tensordict with ndim={input_dict_or_td.ndim} does not have enough dims." + ) + if input_dict_or_td.batch_size[self.stack_dim] != len(self.tensordicts): + raise ValueError("cannot update stacked tensordicts with different shapes.") + for td_dest, td_source in _zip_strict( + self.tensordicts, input_dict_or_td.unbind(self.stack_dim) + ): + td_dest.update_(td_source, clone=clone, non_blocking=non_blocking, **kwargs) + return self + + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorDictBase, + index: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + ) -> Self: + if not _is_tensor_collection(type(input_dict_or_td)): + input_dict_or_td = TensorDict.from_dict( + input_dict_or_td, batch_size=self.batch_size + ) + split_index = self._split_index(index) + converted_idx = split_index["index_dict"] + num_single = split_index["num_single"] + isinteger = split_index["isinteger"] + if isinteger: + # this will break if the index along the stack dim is [0] or :1 or smth + for i, _idx in converted_idx.items(): + self.tensordicts[i].update_at_( + input_dict_or_td, + _idx, + non_blocking=non_blocking, + ) + return self + unbind_dim = self.stack_dim - num_single + for (i, _idx), _value in _zip_strict( + converted_idx.items(), + input_dict_or_td.unbind(unbind_dim), + ): + self.tensordicts[i].update_at_( + _value, + _idx, + non_blocking=non_blocking, + ) + return self + + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> Self: + for td in self.tensordicts: + td.rename_key_(old_key, new_key, safe=safe) + return self + + rename_key = _renamed_inplace_method(rename_key_) + + def where( + self, + condition: Tensor, + other: Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ): + from tensordict import lazy_stack + + if condition.ndim < self.ndim: + condition = expand_right(condition, self.batch_size) + condition = condition.unbind(self.stack_dim) + if _is_tensor_collection(type(other)) or ( + isinstance(other, Tensor) + and other.shape[: self.stack_dim] == self.shape[: self.stack_dim] + ): + other = other.unbind(self.stack_dim) + + def where(td, cond, other, pad): + if cond.numel() > 1: + return td.where(cond, other, pad=pad) + return other if not cond else td + + result = lazy_stack( + [ + where(td, cond, _other, pad=pad) + for td, cond, _other in _zip_strict( + self.tensordicts, condition, other + ) + ], + self.stack_dim, + ) + else: + result = lazy_stack( + [ + td.where(cond, other, pad=pad) + for td, cond in _zip_strict(self.tensordicts, condition) + ], + self.stack_dim, + ) + # We should not pass out to stack because this will overwrite the tensors in-place, but + # we don't want that + if out is not None: + out.update(result, update_batch_size=update_batch_size) + return out + return result + + def masked_fill_(self, mask: Tensor, value: float | bool) -> Self: + mask_unbind = mask.unbind(dim=self.stack_dim) + for _mask, td in _zip_strict(mask_unbind, self.tensordicts): + td.masked_fill_(_mask, value) + return self + + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: + td_copy = self.clone() + return td_copy.masked_fill_(mask, value) + + @lock_blocked + def insert(self, index: int, tensordict: T) -> None: + """Insert a TensorDict into the stack at the specified index. + + Analogous to list.insert. The inserted TensorDict must have compatible + batch_size and device. Insertion is in-place, nothing is returned. + + Args: + index (int): The index at which the new TensorDict should be inserted. + tensordict (TensorDictBase): The TensorDict to be inserted into the stack. + + """ + if not isinstance(tensordict, TensorDictBase): + raise TypeError( + "Expected new value to be TensorDictBase instance but got " + f"{type(tensordict)} instead." + ) + if self.tensordicts: + batch_size = self.tensordicts[0].batch_size + device = self.tensordicts[0].device + + _batch_size = tensordict.batch_size + _device = tensordict.device + + if device != _device: + raise ValueError( + f"Devices differ: stack has device={device}, new value has " + f"device={_device}." + ) + if _batch_size != batch_size: + raise ValueError( + f"Batch sizes in tensordicts differs: stack has " + f"batch_size={batch_size}, new_value has batch_size={_batch_size}." + ) + else: + batch_size = tensordict.batch_size + + self.tensordicts.insert(index, tensordict) + + N = len(self.tensordicts) + self._batch_size = self._compute_batch_size(batch_size, self.stack_dim, N) + + @lock_blocked + def append(self, tensordict: T) -> None: + """Append a TensorDict onto the stack. + + Analogous to list.append. The appended TensorDict must have compatible + batch_size and device. The append operation is in-place, nothing is returned. + + Args: + tensordict (TensorDictBase): The TensorDict to be appended onto the stack. + + """ + self.insert(len(self.tensordicts), tensordict) + + @lock_blocked + def extend(self, tensordict: list[T] | T) -> None: + """Extends the lazy stack with new tensordicts.""" + if _is_tensor_collection(type(tensordict)): + tensordict = list(tensordict.unbind(self.stack_dim)) + if any(not isinstance(tensordict, TensorDictBase) for tensordict in tensordict): + raise TypeError( + "Expected new value to be TensorDictBase instance but got " + f"{[type(tensordict) for tensordict in tensordict]} instead." + ) + if self.tensordicts: + batch_size = self.tensordicts[0].batch_size + device = self.tensordicts[0].device + + for _td in tensordict: + _batch_size = _td.batch_size + _device = _td.device + + if device != _device: + raise ValueError( + f"Devices differ: stack has device={device}, new value has " + f"device={_device}." + ) + if _batch_size != batch_size: + raise ValueError( + f"Batch sizes in tensordicts differs: stack has " + f"batch_size={batch_size}, new_value has batch_size={_batch_size}." + ) + else: + batch_size = tensordict.batch_size + + self.tensordicts.extend(tensordict) + + N = len(self.tensordicts) + self._batch_size = self._compute_batch_size(batch_size, self.stack_dim, N) + + @property + def is_locked(self) -> bool: + if self._is_locked is not None: + # if tensordicts have been locked through this Lazy stack, then we can + # trust this lazy stack to contain the info. + # In all other cases we must check + return self._is_locked + # If any of the tensordicts is not locked, we assume that the lazy stack + # is not locked either. Caching is then disabled and + for td in self.tensordicts: + if not td.is_locked: + return False + else: + if not self.tensordicts: + return False + # In this case, all tensordicts were locked before the lazy stack + # was created and they were not locked through the lazy stack. + # This means we cannot cache the value because this lazy stack + # if not part of the graph. We don't want it to be part of the graph + # because this object being locked is only a side-effect. + # Calling self.lock_() here could however speed things up. + return True + + @is_locked.setter + def is_locked(self, value: bool) -> None: + if value: + self.lock_() + else: + self.unlock_() + + @property + def _lock_parents_weakrefs(self): + """Weakrefs of all tensordicts that need to be unlocked for this to be unlocked.""" + _lock_parents_weakrefs = [] + for tensordict in self.tensordicts: + _lock_parents_weakrefs = ( + _lock_parents_weakrefs + tensordict._lock_parents_weakrefs + ) + _lock_parents_weakrefs = [ + item for item in _lock_parents_weakrefs if item is not weakref.ref(self) + ] + return _lock_parents_weakrefs + + def _propagate_lock(self, lock_parents_weakrefs=None, *, is_compiling): + """Registers the parent tensordict that handles the lock.""" + self._is_locked = True + if not is_compiling: + is_root = lock_parents_weakrefs is None + if is_root: + lock_parents_weakrefs = [] + + lock_parents_weakrefs = copy(lock_parents_weakrefs) + [weakref.ref(self)] + for dest in self.tensordicts: + dest._propagate_lock(lock_parents_weakrefs, is_compiling=is_compiling) + + @erase_cache + def _propagate_unlock(self): + # we can't set _is_locked to False because after it's unlocked, anything + # can happen to a child tensordict. + self._is_locked = None + sub_tds = defaultdict() + for child in self.tensordicts: + # we want to make sure that if the same child is present twice in the + # stack we won't iterate multiple times over it + sub_tds[id(child)] = child._propagate_unlock() + [child] + sub_tds = [item for value in sub_tds.values() for item in value] + return sub_tds + + def __repr__(self): + fields = _td_fields(self) + field_str = indent(f"fields={{{fields}}}", 4 * " ") + exclusive_fields_str = indent( + f"exclusive_fields={{{self._repr_exclusive_fields()}}}", 4 * " " + ) + batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ") + device_str = indent(f"device={self.device}", 4 * " ") + is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ") + stack_dim = indent(f"stack_dim={self.stack_dim}", 4 * " ") + string = ",\n".join( + [ + field_str, + exclusive_fields_str, + batch_size_str, + device_str, + is_shared_str, + stack_dim, + ] + ) + return f"{type(self).__name__}(\n{string})" + + def _exclusive_keys(self): + return {key for td in self.tensordicts for key in td.keys()} + + def _repr_exclusive_fields(self): + keys = set(self.keys()) + exclusive_keys = [ + _td_fields(td, [k for k in td.keys() if k not in keys]) + for td in self.tensordicts + ] + exclusive_key_str = ",\n".join( + [ + indent(f"{i} ->{line}", 4 * " ") + for i, line in enumerate(exclusive_keys) + if line != "\n" + ] + ) + + return "\n" + exclusive_key_str + + def _view(self, *args, raise_if_not_view: bool = True, **kwargs) -> Self: + shape = _get_shape_from_args(*args, **kwargs) + if any(dim < 0 for dim in shape): + shape = _infer_size_impl(shape, self.numel()) + + # Then we just need to reorganize the lazy stack + shape = torch.Size(shape) + is_flatten, (i, j) = _check_is_flatten( + shape, self.batch_size, return_flatten_dim=True + ) + if is_flatten: + # we need to get a flat representation of all the elements from dim i to j, starting from j + tds = [self] + for _ in range(i, j + 1): + # for k in range(j, i-1, -1): + # unbind along k + tds = [_td for local_td in tds for _td in local_td.unbind(i)] + # the dim along which to stack is the first, ie, i + tds = self._new_lazy_unsafe(*tds, stack_dim=i) + if self.is_locked: + return tds.lock_() + return tds + + is_unflatten, (i, j) = _check_is_unflatten( + shape, self.batch_size, return_flatten_dim=True + ) + if is_unflatten: + # we are going to organize our list of (A*B*C) elements in a nested list of (A * (B * (C))) elements + tds = self + for k in range(i, j): + tds = self._new_lazy_unsafe( + *list(tds.chunk(shape[k], dim=k)), stack_dim=k + ) + if self.is_locked: + return tds.lock_() + return tds + if raise_if_not_view: + raise RuntimeError( + "Cannot call `view` on a lazy stacked tensordict. Call `reshape` instead." + ) + return TensorDict.reshape(self, shape) + + def reshape( + self, + *args, + **kwargs, + ) -> Self: + return self._view(*args, raise_if_not_view=False, **kwargs) + + def flatten(self, start_dim: int = 0, end_dim=-1): + end_dim = _maybe_correct_neg_dim(end_dim, shape=self.batch_size) + start_dim = _maybe_correct_neg_dim(start_dim, shape=self.batch_size) + new_shape = [ + s for i, s in enumerate(self.batch_size) if i < start_dim or i > end_dim + ] + new_shape.insert(start_dim, -1) + return self.view(new_shape) + + def unflatten(self, dim, unflattened_size): + dim = _maybe_correct_neg_dim(dim, shape=self.batch_size) + new_shape = self.batch_size + if dim == 0: + new_shape = torch.Size(unflattened_size) + new_shape[1:] + else: + new_shape = ( + new_shape[:dim] + torch.Size(unflattened_size) + new_shape[dim + 1 :] + ) + return self.view(new_shape) + + def _transpose(self, dim0, dim1): + if self._is_vmapped: + raise RuntimeError("cannot call transpose within vmap.") + if dim0 == self.stack_dim: + # we know dim0 and dim1 are sorted so dim1 comes after dim0 + # example: shape = [5, 4, 3, 2, 1], stack_dim=1, dim0=1, dim1=4 + # resulting shape: [5, 1, 3, 2, 4] + if dim1 == dim0 + 1: + result = type(self)( + *self.tensordicts, stack_dim=dim1, stack_dim_name=self._td_dim_name + ) + else: + result = type(self)( + *(td.transpose(dim0, dim1 - 1) for td in self.tensordicts), + stack_dim=dim1, + stack_dim_name=self._td_dim_name, + ) + elif dim1 == self.stack_dim: + # example: shape = [5, 4, 3, 2, 1], stack_dim=3, dim0=1, dim1=3 + # resulting shape: [5, 2, 3, 4, 1] + if dim0 + 1 == dim1: + result = type(self)( + *self.tensordicts, stack_dim=dim0, stack_dim_name=self._td_dim_name + ) + else: + result = type(self)( + *(td.transpose(dim0 + 1, dim1) for td in self.tensordicts), + stack_dim=dim0, + stack_dim_name=self._td_dim_name, + ) + else: + dim0 = dim0 if dim0 < self.stack_dim else dim0 - 1 + dim1 = dim1 if dim1 < self.stack_dim else dim1 - 1 + result = type(self)( + *(td.transpose(dim0, dim1) for td in self.tensordicts), + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + ) + return result + + def _repeat(self, *repeats: int) -> TensorDictBase: + repeats = list(repeats) + r_dim = repeats.pop(self.stack_dim) + tds = [td.repeat(*repeats) for td in self.tensordicts] + tds = [td for _ in range(r_dim) for td in tds] + return type(self)( + *tds, + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + hook_in=self.hook_in, + hook_out=self.hook_out, + ) + + def repeat_interleave( + self, + repeats: torch.Tensor | int, + dim: int | None = None, + *, + output_size: int | None = None, + ) -> TensorDictBase: + if self.ndim == 0: + return self.unsqueeze(0).repeat_interleave( + repeats=repeats, dim=dim, output_size=output_size + ) + if dim is None: + if self.ndim > 1: + return self.reshape(-1).repeat_interleave(repeats, dim=0) + return self.repeat_interleave(repeats, dim=0) + dim_corrected = dim if dim >= 0 else self.ndim + dim + if not (dim_corrected >= 0): + raise ValueError( + f"dim {dim} is out of range for tensordict with shape {self.shape}." + ) + if dim_corrected == self.stack_dim: + if isinstance(repeats, int): + repeats: list[int] = [repeats] * len(self.tensordicts) + else: + repeats = repeats.tolist() + new_list_of_tds = [ + t for t, r in zip(self.tensordicts, repeats) for _ in range(r) + ] + result = type(self)( + *new_list_of_tds, + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + hook_out=self.hook_out, + hook_in=self.hook_in, + ) + else: + dim_corrected = ( + dim_corrected if dim_corrected < self.stack_dim else dim_corrected - 1 + ) + result = type(self)( + *( + td.repeat_interleave( + repeats=repeats, dim=dim_corrected, output_size=output_size + ) + for td in self.tensordicts + ), + stack_dim=self.stack_dim, + stack_dim_name=self._td_dim_name, + hook_in=self.hook_in, + hook_out=self.hook_out, + ) + return result + + def _permute( + self, + *args, + **kwargs, + ): + dims_list = _get_shape_from_args(*args, kwarg_name="dims", **kwargs) + dims_list = [dim if dim >= 0 else self.ndim + dim for dim in dims_list] + dims_list_sort = np.argsort(dims_list) + # find the new stack dim + stack_dim = dims_list_sort[self.stack_dim] + # remove that dim from the dims_list + dims_list = [ + d if d < self.stack_dim else d - 1 for d in dims_list if d != self.stack_dim + ] + result = LazyStackedTensorDict.lazy_stack( + [td.permute(dims_list) for td in self.tensordicts], + stack_dim, + stack_dim_name=self._td_dim_name, + ) + return result + + def _squeeze(self, dim=None): + if dim is not None: + new_dim = dim + if new_dim < 0: + new_dim = self.batch_dims + new_dim + if new_dim > self.batch_dims - 1 or new_dim < 0: + raise RuntimeError( + f"The dim provided to squeeze is incompatible with the tensordict shape: dim={dim} and batch_size={self.batch_size}." + ) + dim = new_dim + if self.batch_size[dim] != 1: + return self + if dim == self.stack_dim: + return self.tensordicts[0] + if dim > self.stack_dim: + dim = dim - 1 + stack_dim = self.stack_dim + else: + stack_dim = self.stack_dim - 1 + result = LazyStackedTensorDict.lazy_stack( + [td.squeeze(dim) for td in self.tensordicts], + stack_dim, + stack_dim_name=self._td_dim_name, + ) + else: + result = self + for dim in range(self.batch_dims - 1, -1, -1): + if self.batch_size[dim] == 1: + result = result.squeeze(dim) + return result + + def _unsqueeze(self, dim: int): + new_dim = dim + if new_dim < 0: + new_dim = self.batch_dims + new_dim + 1 + if new_dim > self.batch_dims or new_dim < 0: + raise RuntimeError( + f"The dim provided to unsqueeze is incompatible with the tensordict shape: dim={dim} and batch_size={self.batch_size}." + ) + dim = new_dim + if dim > self.stack_dim: + dim = dim - 1 + stack_dim = self.stack_dim + else: + stack_dim = self.stack_dim + 1 + result = LazyStackedTensorDict.lazy_stack( + [td.unsqueeze(dim) for td in self.tensordicts], + stack_dim, + stack_dim_name=self._td_dim_name, + ) + return result + + def split(self, split_size: int | list[int], dim: int = 0) -> list[TensorDictBase]: + dim = _maybe_correct_neg_dim(dim, shape=self.shape) + if dim == self.stack_dim: + if isinstance(split_size, int): + split_size = [split_size] * -(len(self.tensordicts) // -split_size) + split_size[-1] = len(self.tensordicts) - sum(split_size[:-1]) + + def iter_across_tds(): + start = 0 + for s in split_size: + if s == 0: + batch_size = list(self._batch_size) + batch_size.pop(self.stack_dim) + yield LazyStackedTensorDict( + batch_size=batch_size, + device=self.device, + stack_dim=self.stack_dim, + ) + continue + stop = start + s + yield self._new_lazy_unsafe( + *self.tensordicts[slice(start, stop)], stack_dim=self.stack_dim + ) + start = stop + + return tuple(iter_across_tds()) + tds = [] + split_dim = dim if dim < self.stack_dim else dim - 1 + for td in self.tensordicts: + tds.append(td.split(split_size, split_dim)) + return tuple( + self._new_lazy_unsafe(*tds, stack_dim=self.stack_dim) + for tds in _zip_strict(*tds) + ) + + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorCollection, ...]: + splits = -(self.batch_size[dim] // -chunks) + return self.split(splits, dim) + + lock_ = TensorDictBase.lock_ + lock = _renamed_inplace_method(lock_) + + unlock_ = TensorDictBase.unlock_ + unlock = _renamed_inplace_method(unlock_) + + _check_device = TensorDict._check_device + _check_is_shared = TensorDict._check_is_shared + _convert_to_tensordict = TensorDict._convert_to_tensordict + _index_tensordict = TensorDict._index_tensordict + masked_select = TensorDict.masked_select + _to_module = TensorDict._to_module + from_dict_instance = TensorDict.from_dict_instance + + +class _CustomOpTensorDict(TensorDictBase): + """Encodes lazy operations on tensors contained in a TensorDict.""" + + _safe = False + _lazy = True + + def __init__( + self, + source: T, + custom_op: str, + inv_op: str | None = None, + custom_op_kwargs: dict | None = None, + inv_op_kwargs: dict | None = None, + batch_size: Sequence[int] | None = None, + ) -> None: + + if not isinstance(source, TensorDictBase): + raise TypeError( + f"Expected source to be a TensorDictBase isntance, " + f"but got {type(source)} instead." + ) + self._source = source + self.custom_op = custom_op + self.inv_op = inv_op + self.custom_op_kwargs = custom_op_kwargs if custom_op_kwargs is not None else {} + self.inv_op_kwargs = inv_op_kwargs if inv_op_kwargs is not None else {} + self._batch_size = None + if batch_size is not None and batch_size != self.batch_size: + raise RuntimeError("batch_size does not match self.batch_size.") + + # These attributes should never be set + @property + @cache # noqa + def _is_shared(self): + return self._source._is_shared + + @property + @cache # noqa + def _is_memmap(self): + return self._source._is_memmap + + def is_empty(self) -> bool: + return self._source.is_empty() + + def is_memmap(self) -> bool: + return self._source.is_memmap() + + def is_shared(self) -> bool: + return self._source.is_shared() + + def _update_custom_op_kwargs(self, source_tensor: Tensor) -> dict[str, Any]: + """Allows for a transformation to be customized for a certain shape, device or dtype. + + By default, this is a no-op on self.custom_op_kwargs + + Args: + source_tensor: corresponding Tensor + + Returns: + a dictionary with the kwargs of the operation to execute + for the tensor + + """ + return self.custom_op_kwargs + + def _update_inv_op_kwargs(self, source_tensor: Tensor) -> dict[str, Any]: + """Allows for an inverse transformation to be customized for a certain shape, device or dtype. + + By default, this is a no-op on self.inv_op_kwargs + + Args: + source_tensor: corresponding tensor + + Returns: + a dictionary with the kwargs of the operation to execute for + the tensor + + """ + return self.inv_op_kwargs + + def entry_class(self, key: NestedKey) -> type: + return type(self._source.get(key)) + + @classmethod + def from_dict( + cls, input_dict, batch_size=None, device=None, batch_dims=None, names=None + ): + raise NotImplementedError(f"from_dict not implemented for {cls.__name__}.") + + @property + def device(self) -> torch.device | None: + return self._source.device + + @device.setter + def device(self, value: DeviceType) -> None: + self._source.device = value + + @property + def batch_size(self) -> torch.Size: + if self._batch_size is None: + self._batch_size = getattr( + torch.zeros(self._source.batch_size, device="meta"), self.custom_op + )(**self.custom_op_kwargs).shape + return self._batch_size + + @batch_size.setter + def batch_size(self, new_size: torch.Size) -> None: + self._batch_size_setter(new_size) + + def _has_names(self): + return self._source._has_names() + + def _erase_names(self): + raise RuntimeError( + f"Cannot erase names of a {type(self).__name__}. " + f"Erase source TensorDict's names instead." + ) + + def _rename_subtds(self, names): + for key in self.keys(): + if _is_tensor_collection(self.entry_class(key)): + raise RuntimeError( + "Cannot rename dimensions of a lazy TensorDict with " + "nested collections. Convert the instance to a regular " + "tensordict by using the `to_tensordict()` method first." + ) + + def _change_batch_size(self, new_size: torch.Size) -> None: + self._batch_size = new_size + + def _get_str(self, key, default, **kwargs): + tensor = self._source._get_str(key, default, **kwargs) + if tensor is default: + return tensor + return self._transform_value(tensor) + + def _get_tuple(self, key, default, **kwargs): + tensor = self._source._get_tuple(key, default, **kwargs) + if tensor is default: + return tensor + return self._transform_value(tensor) + + def _transform_value(self, item): + return getattr(item, self.custom_op)(**self._update_custom_op_kwargs(item)) + + def _set_str( + self, + key, + value, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ): + if not validated: + value = self._validate_value( + value, check_shape=True, non_blocking=non_blocking + ) + validated = True + value = getattr(value, self.inv_op)(**self._update_inv_op_kwargs(value)) + self._source._set_str( + key, + value, + inplace=inplace, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + return self + + def _set_tuple( + self, key, value, *, inplace: bool, validated: bool, non_blocking: bool + ): + if len(key) == 1: + return self._set_str( + key[0], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + source = self._source._get_str(key[0], None) + if source is None: + source = self._source._create_nested_str(key[0]) + nested = type(self)( + source, + custom_op=self.custom_op, + inv_op=self.inv_op, + custom_op_kwargs=self._update_custom_op_kwargs(source), + inv_op_kwargs=self._update_inv_op_kwargs(source), + ) + nested._set_tuple( + key[1:], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + return self + + def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool): + transformed_tensor, original_tensor = self._get_str( + key, NO_DEFAULT + ), self._source._get_str(key, NO_DEFAULT) + if transformed_tensor.data_ptr() != original_tensor.data_ptr(): + raise RuntimeError( + f"{self} original tensor and transformed_in do not point to the " + f"same storage. Setting values in place is not currently " + f"supported in this setting, consider calling " + f"`td.clone()` before `td.set_at_(...)`" + ) + transformed_tensor[idx] = value + return self + + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool): + transformed_tensor, original_tensor = self._get_tuple( + key, NO_DEFAULT + ), self._source._get_tuple(key, NO_DEFAULT) + if transformed_tensor.data_ptr() != original_tensor.data_ptr(): + raise RuntimeError( + f"{self} original tensor and transformed_in do not point to the " + f"same storage. Setting values in place is not currently " + f"supported in this setting, consider calling " + f"`td.clone()` before `td.set_at_(...)`" + ) + if not validated: + value = self._validate_value( + value, check_shape=False, non_blocking=non_blocking + ) + + transformed_tensor[idx] = value + return self + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + raise RuntimeError( + f"stacking tensordicts is not allowed for type {type(self).__name__}" + f"consider calling 'to_tensordict()` first" + ) + + def __repr__(self) -> str: + custom_op_kwargs_str = ", ".join( + [f"{key}={value}" for key, value in self.custom_op_kwargs.items()] + ) + indented_source = textwrap.indent(f"source={self._source}", "\t") + return ( + f"{type(self).__name__}(\n{indented_source}, " + f"\n\top={self.custom_op}({custom_op_kwargs_str}))" + ) + + # @cache # noqa: B019 + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _TensorDictKeysView: + return self._source.keys( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + def _select( + self, + *keys: NestedKey, + inplace: bool = False, + strict: bool = True, + set_shared: bool = True, + ) -> _CustomOpTensorDict: + if inplace: + raise RuntimeError("Cannot call select inplace on a lazy tensordict.") + return self.to_tensordict(retain_none=True)._select( + *keys, inplace=False, strict=strict, set_shared=set_shared + ) + + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> _CustomOpTensorDict: + if inplace: + raise RuntimeError("Cannot call exclude inplace on a lazy tensordict.") + return self.to_tensordict()._exclude( + *keys, inplace=False, set_shared=set_shared + ) + + def _clone(self, recurse: bool = True) -> Self: + """Clones the Lazy TensorDict. + + Args: + recurse (bool, optional): if ``True`` (default), a regular + :class:`~.tensordict.TensorDict` instance will be returned. + Otherwise, another :class:`~.tensordict.SubTensorDict` with identical content + will be returned. + """ + if not recurse: + return type(self)( + source=self._source.clone(False), + custom_op=self.custom_op, + inv_op=self.inv_op, + custom_op_kwargs=self.custom_op_kwargs, + inv_op_kwargs=self.inv_op_kwargs, + batch_size=self.batch_size, + ) + return self.to_tensordict() + + def is_contiguous(self) -> bool: + return all([value.is_contiguous() for _, value in self.items()]) + + def contiguous(self) -> Self: + def contiguous(x): + return x.contiguous() + + return self._fast_apply(contiguous, propagate_lock=True) + + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> _CustomOpTensorDict: + self._source.rename_key_(old_key, new_key, safe=safe) + return self + + rename_key = _renamed_inplace_method(rename_key_) + + @lock_blocked + def del_(self, key: NestedKey) -> _CustomOpTensorDict: + self._source = self._source.del_(key) + return self + + def to(self, *args, **kwargs) -> Self: + non_blocking = kwargs.pop("non_blocking", None) + ( + device, + dtype, + _, + convert_to_format, + batch_size, + pin_memory, + num_threads, + inplace, + ) = _parse_to(*args, **kwargs) + if inplace: + raise TypeError(f"Cannot use inplace=True with {type(self).__name__}.to().") + + if batch_size is not None: + raise TypeError(f"Cannot pass batch-size to {type(self).__name__}.to().") + result = self + + if device is not None and dtype is None and device == self.device: + return result + + td = self._source.to(*args, non_blocking=non_blocking, **kwargs) + self_copy = copy(self) + self_copy._source = td + return self_copy + + def pin_memory( + self, *, num_threads: int | str = 0, inplace: bool = False + ) -> _CustomOpTensorDict: + _source = self._source.pin_memory(num_threads=num_threads, inplace=inplace) + if not inplace: + return type(self)( + source=_source, + custom_op=self.custom_op, + inv_op=self.inv_op, + custom_op_kwargs=self.custom_op_kwargs, + inv_op_kwargs=self.inv_op_kwargs, + batch_size=self.batch_size, + ) + return self + + @lock_blocked + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + key, val = self._source.popitem() + return key, self._transform_value(val) + + def detach_(self) -> _CustomOpTensorDict: + self._source.detach_() + return self + + def where( + self, + condition: Tensor, + other: Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ): + return self.to_tensordict().where( + condition=condition, + other=other, + out=out, + pad=pad, + update_batch_size=update_batch_size, + ) + + def masked_fill_(self, mask: Tensor, value: float | bool) -> _CustomOpTensorDict: + for key, item in self.items(): + val = self._source.get(key) + mask_exp = expand_right( + mask, list(mask.shape) + list(val.shape[self._source.batch_dims :]) + ) + mask_proc_inv = getattr(mask_exp, self.inv_op)( + **self._update_inv_op_kwargs(item) + ) + val[mask_proc_inv] = value + self._source.set(key, val) + return self + + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: + td_copy = self.clone() + return td_copy.masked_fill_(mask, value) + + def _memmap_( + self, + *, + prefix: str | None, + copy_existing: bool, + executor, + futures, + inplace, + like, + share_non_tensor, + existsok, + robust_key, + ) -> Self: + def save_metadata(data: TensorDictBase, filepath, metadata=None): + if metadata is None: + metadata = {} + metadata.update( + { + "shape": list(data.shape), + "device": str(data.device), + "_type": str(type(data)), + "custom_op": data.custom_op, + "inv_op": data.inv_op, + "custom_op_kwargs": data.custom_op_kwargs, + "inv_op_kwargs": data.inv_op_kwargs, + } + ) + with open(filepath, "wb") as json_metadata: + from tensordict.utils import json_dumps + + json_str = json_dumps(metadata) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + json_metadata.write(json_str.encode("utf-8")) + else: + json_metadata.write(json_str) + + if prefix is not None: + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + metadata = {} + + dest_source = self._source._memmap_( + prefix=None if prefix is None else prefix / "_source", + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if not inplace: + dest = type(self)( + dest_source, + custom_op=self.custom_op, + inv_op=self.inv_op, + custom_op_kwargs=self.custom_op_kwargs, + inv_op_kwargs=self.inv_op_kwargs, + batch_size=self.batch_size, + ) + else: + dest = self + + if prefix is not None: + if executor is None: + save_metadata( + dest, + prefix / "meta.json", + metadata=metadata, + ) + else: + futures.append( + executor.submit(save_metadata, dest, prefix / "meta.json", metadata) + ) + return dest + + @classmethod + def _load_memmap(cls, prefix: str, metadata: dict, **kwargs) -> _CustomOpTensorDict: + custom_op = metadata.pop("custom_op") + inv_op = metadata.pop("inv_op") + custom_op_kwargs = metadata.pop("custom_op_kwargs") + inv_op_kwargs = metadata.pop("inv_op_kwargs") + + source = TensorDict.load_memmap(prefix / "_source", **kwargs, non_blocking=True) + + return cls( + source, + custom_op=custom_op, + inv_op=inv_op, + custom_op_kwargs=custom_op_kwargs, + inv_op_kwargs=inv_op_kwargs, + ) + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for lazy tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for lazy tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_tensor( + self, key: NestedKey, tensor: torch.Tensor, *, copy_data: bool = True + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for lazy tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def share_memory_(self) -> _CustomOpTensorDict: + self._source.share_memory_() + self.lock_() + return self + + @property + def _td_dim_names(self): + # we also want for _td_dim_names to be accurate + if self._source._td_dim_names is None: + return None + return self.names + + @property + def is_locked(self) -> bool: + return self._source.is_locked + + @is_locked.setter + def is_locked(self, value) -> bool: + if value: + self.lock_() + else: + self.unlock_() + + @_as_context_manager("is_locked") + def lock_(self) -> Self: + self._source.lock_() + return self + + @erase_cache + @_as_context_manager("is_locked") + def unlock_(self) -> Self: + self._source.unlock_() + return self + + def _remove_lock(self, lock_id): + return self._source._remove_lock(lock_id) + + @erase_cache + def _propagate_lock(self, lock_ids, *, is_compiling): + return self._source._propagate_lock(lock_ids, is_compiling=is_compiling) + + @erase_cache + def _propagate_unlock(self): + return self._source._propagate_unlock() + + lock = _renamed_inplace_method(lock_) + unlock = _renamed_inplace_method(unlock_) + + def __del__(self): + pass + + @property + def sorted_keys(self): + return self._source.sorted_keys + + def _view(self, *args, **kwargs): + raise RuntimeError( + "Cannot call `view` on a lazy tensordict. Call `reshape` instead." + ) + + def _transpose(self, dim0, dim1): + raise RuntimeError( + "Cannot call `transpose` on a lazy tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _permute( + self, + *args, + **kwargs, + ): + raise RuntimeError( + "Cannot call `permute` on a lazy tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _squeeze(self, dim=None): + raise RuntimeError( + "Cannot call `squeeze` on a lazy tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _unsqueeze(self, dim: int): + raise RuntimeError( + "Cannot call `unsqueeze` on a lazy tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _cast_reduction( + self, + *, + reduction_name, + dim=NO_DEFAULT, + keepdim=NO_DEFAULT, + tuple_ok=True, + **kwargs, + ): + try: + td = self.to_tensordict() + except Exception: + raise RuntimeError( + f"{reduction_name} requires this object to be cast to a regular TensorDict. " + f"If you need {type(self).__name__} to support {reduction_name}, help us by filing an issue" + f" on github!" + ) + return td._cast_reduction( + reduction_name=reduction_name, + dim=dim, + keepdim=keepdim, + tuple_ok=tuple_ok, + **kwargs, + ) + + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]: + splits = -(self.batch_size[dim] // -chunks) + return self.split(splits, dim) + + __xor__ = TensorDict.__xor__ + __or__ = TensorDict.__or__ + __eq__ = TensorDict.__eq__ + __ne__ = TensorDict.__ne__ + __ge__ = TensorDict.__ge__ + __gt__ = TensorDict.__gt__ + __le__ = TensorDict.__le__ + __lt__ = TensorDict.__lt__ + __setitem__ = TensorDict.__setitem__ + _add_batch_dim = TensorDict._add_batch_dim + _check_device = TensorDict._check_device + _check_is_shared = TensorDict._check_is_shared + _convert_to_tensordict = TensorDict._convert_to_tensordict + _index_tensordict = TensorDict._index_tensordict + + _apply_nest = TensorDict._apply_nest + _get_names_idx = TensorDict._get_names_idx + _maybe_remove_batch_dim = TensorDict._maybe_remove_batch_dim + _multithread_apply_flat = TensorDict._multithread_apply_flat + _multithread_rebuild = TensorDict._multithread_rebuild + _remove_batch_dim = TensorDict._remove_batch_dim + _to_module = TensorDict._to_module + _unbind = TensorDict._unbind + all = TensorDict.all + any = TensorDict.any + expand = TensorDict.expand + from_dict_instance = TensorDict.from_dict_instance + masked_select = TensorDict.masked_select + _repeat = TensorDict._repeat + repeat_interleave = TensorDict.repeat_interleave + reshape = TensorDict.reshape + split = TensorDict.split + + +class _UnsqueezedTensorDict(_CustomOpTensorDict): + """A lazy view on an unsqueezed TensorDict. + + When calling `tensordict.unsqueeze(dim)`, a lazy view of this operation is + returned such that the following code snippet works without raising an + exception: + + >>> assert tensordict.unsqueeze(dim).squeeze(dim) is tensordict + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict({'a': torch.randn(3, 4)}, batch_size=[3]) + >>> td_unsqueeze = td.unsqueeze(-1) + >>> print(td_unsqueeze.shape) + torch.Size([3, 1]) + >>> print(td_unsqueeze.squeeze(-1) is td) + True + """ + + def _legacy_squeeze(self, dim: int | None) -> Self: + if dim is not None and dim < 0: + dim = self.batch_dims + dim + if dim == self.custom_op_kwargs.get("dim"): + return self._source + return super()._legacy_squeeze(dim) + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + unsqueezed_dim = self.custom_op_kwargs["dim"] + diff_to_apply = 1 if dim < unsqueezed_dim else 0 + list_item_unsqueeze = [ + item.squeeze(unsqueezed_dim - diff_to_apply) for item in list_item + ] + return self._source._stack_onto_(list_item_unsqueeze, dim) + + @property + def names(self): + names = list(self._source.names) + dim = self.custom_op_kwargs.get("dim") + names.insert(dim, None) + return names + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a lazy tensordict cannot be modified. Call to_tensordict() first." + ) + + +class _SqueezedTensorDict(_CustomOpTensorDict): + """A lazy view on a squeezed TensorDict. + + See the `UnsqueezedTensorDict` class documentation for more information. + + """ + + def _legacy_unsqueeze(self, dim: int) -> Self: + if dim < 0: + dim = self.batch_dims + dim + 1 + inv_op_dim = self.inv_op_kwargs.get("dim") + if inv_op_dim < 0: + inv_op_dim = self.batch_dims + inv_op_dim + 1 + if dim == inv_op_dim: + return self._source + return super()._legacy_unsqueeze(dim) + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + squeezed_dim = self.custom_op_kwargs["dim"] + # dim=0, squeezed_dim=2, [3, 4, 5] [3, 4, 1, 5] [[4, 5], [4, 5], [4, 5]] => unsq 1 + # dim=1, squeezed_dim=2, [3, 4, 5] [3, 4, 1, 5] [[3, 5], [3, 5], [3, 5], [3, 4]] => unsq 1 + # dim=2, squeezed_dim=2, [3, 4, 5] [3, 4, 1, 5] [[3, 4], [3, 4], ...] => unsq 2 + diff_to_apply = 1 if dim < squeezed_dim else 0 + list_item_unsqueeze = [ + item.unsqueeze(squeezed_dim - diff_to_apply) for item in list_item + ] + return self._source._stack_onto_(list_item_unsqueeze, dim) + + @property + def names(self): + names = list(self._source.names) + dim = self.custom_op_kwargs["dim"] + if self._source.batch_size[dim] == 1: + del names[dim] + return names + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a lazy tensordict cannot be modified. Call to_tensordict() first." + ) + + +class _ViewedTensorDict(_CustomOpTensorDict): + def _update_custom_op_kwargs(self, source_tensor: Tensor) -> dict[str, Any]: + new_dim_list = list(self.custom_op_kwargs.get("size")) + new_dim_list += list(source_tensor.shape[self._source.batch_dims :]) + new_dim = torch.Size(new_dim_list) + new_dict = deepcopy(self.custom_op_kwargs) + new_dict.update({"size": new_dim}) + return new_dict + + def _update_inv_op_kwargs(self, tensor: Tensor) -> dict: + size = list(self.inv_op_kwargs.get("size")) + size += list(_shape(tensor)[self.batch_dims :]) + new_dim = torch.Size(size) + new_dict = deepcopy(self.inv_op_kwargs) + new_dict.update({"size": new_dim}) + return new_dict + + def _legacy_view( + self, *shape: int, size: list | tuple | torch.Size | None = None + ) -> Self: + if len(shape) == 0 and size is not None: + return self._legacy_view(*size) + elif len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)): + return self._legacy_view(*shape[0]) + elif not isinstance(shape, torch.Size): + shape = infer_size_impl(shape, self.numel()) + shape = torch.Size(shape) + if shape == self._source.batch_size: + return self._source + return super()._legacy_view(*shape) + + @property + def names(self): + return [None] * self.ndim + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a lazy tensordict cannot be modified. Call to_tensordict() first." + ) + + +class _TransposedTensorDict(_CustomOpTensorDict): + """A lazy view on a TensorDict with two batch dimensions transposed. + + When calling `tensordict.permute(dims_list, dim)`, a lazy view of this operation is + returned such that the following code snippet works without raising an + exception: + + >>> assert tensordict.transpose(dims_list, dim).transpose(dims_list, dim) is tensordict + + """ + + def _legacy_transpose(self, dim0, dim1) -> Self: + if dim0 < 0: + dim0 = self.ndim + dim0 + if dim1 < 0: + dim1 = self.ndim + dim1 + if any((dim0 < 0, dim1 < 0)): + raise ValueError( + "The provided dimensions are incompatible with the tensordict batch-size." + ) + if dim0 == dim1: + return self + dims = (self.inv_op_kwargs.get("dim0"), self.inv_op_kwargs.get("dim1")) + if dim0 in dims and dim1 in dims: + return self._source + return super()._legacy_transpose(dim0, dim1) + + def add_missing_dims( + self, num_dims: int, batch_dims: tuple[int, ...] + ) -> tuple[int, ...]: + dim_diff = num_dims - len(batch_dims) + all_dims = list(range(num_dims)) + for i, x in enumerate(batch_dims): + if x < 0: + x = x - dim_diff + all_dims[i] = x + return tuple(all_dims) + + def _update_custom_op_kwargs(self, source_tensor: Tensor) -> dict[str, Any]: + return self.custom_op_kwargs + + def _update_inv_op_kwargs(self, tensor: Tensor) -> dict[str, Any]: + return self.custom_op_kwargs + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + trsp = self.custom_op_kwargs["dim0"], self.custom_op_kwargs["dim1"] + if dim == trsp[0]: + dim = trsp[1] + elif dim == trsp[1]: + dim = trsp[0] + + list_permuted_items = [] + for item in list_item: + list_permuted_items.append(item.transpose(*trsp)) + self._source._stack_onto_(list_permuted_items, dim) + return self + + @property + def names(self): + names = list(self._source.names) + dim0 = self.custom_op_kwargs["dim0"] + dim1 = self.custom_op_kwargs["dim1"] + names = [ + names[dim0] if i == dim1 else names[dim1] if i == dim0 else name + for i, name in enumerate(names) + ] + return names + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a lazy tensordict cannot be modified. Call to_tensordict() first." + ) + + +class _PermutedTensorDict(_CustomOpTensorDict): + """A lazy view on a TensorDict with the batch dimensions permuted. + + When calling `tensordict.permute(dims_list, dim)`, a lazy view of this operation is + returned such that the following code snippet works without raising an + exception: + + >>> assert tensordict.permute(dims_list, dim).permute(dims_list, dim) is tensordict + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict({'a': torch.randn(4, 5, 6, 9)}, batch_size=[3]) + >>> td_permute = td.permute(dims=(2, 1, 0)) + >>> print(td_permute.shape) + torch.Size([6, 5, 4]) + >>> print(td_permute.permute(dims=(2, 1, 0)) is td) + True + + """ + + def _legacy_permute( + self, + *dims_list: int, + dims: Sequence[int] | None = None, + ) -> Self: + if len(dims_list) == 0: + dims_list = dims + elif len(dims_list) == 1 and not isinstance(dims_list[0], int): + dims_list = dims_list[0] + if len(dims_list) != len(self.shape): + raise RuntimeError( + f"number of dims don't match in permute (got {len(dims_list)}, expected {len(self.shape)}" + ) + if not len(dims_list) and not self.batch_dims: + return self + if np.array_equal(dims_list, range(self.batch_dims)): + return self + if np.array_equal(np.argsort(dims_list), self.inv_op_kwargs.get("dims")): + return self._source + return super()._legacy_permute(*dims_list) + + def add_missing_dims( + self, num_dims: int, batch_dims: tuple[int, ...] + ) -> tuple[int, ...]: + # Adds the feature dimensions to the permute dims + dim_diff = num_dims - len(batch_dims) + all_dims = list(range(num_dims)) + for i, x in enumerate(batch_dims): + if x < 0: + x = x - dim_diff + all_dims[i] = x + return tuple(all_dims) + + def _update_custom_op_kwargs(self, source_tensor: Tensor) -> dict[str, Any]: + new_dims = self.add_missing_dims( + len(source_tensor.shape), self.custom_op_kwargs["dims"] + ) + kwargs = deepcopy(self.custom_op_kwargs) + kwargs.update({"dims": new_dims}) + return kwargs + + def _update_inv_op_kwargs(self, tensor: Tensor) -> dict[str, Any]: + new_dims = self.add_missing_dims( + self._source.batch_dims + len(_shape(tensor)[self.batch_dims :]), + self.custom_op_kwargs["dims"], + ) + kwargs = deepcopy(self.custom_op_kwargs) + kwargs.update({"dims": tuple(np.argsort(new_dims))}) + return kwargs + + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + permute_dims = self.custom_op_kwargs["dims"] + inv_permute_dims = np.argsort(permute_dims) + new_dim = [i for i, v in enumerate(inv_permute_dims) if v == dim][0] + inv_permute_dims = [p for p in inv_permute_dims if p != dim] + inv_permute_dims = np.argsort(np.argsort(inv_permute_dims)) + + list_permuted_items = [] + for item in list_item: + perm = list(inv_permute_dims) + list( + range(self.batch_dims - 1, item.ndimension()) + ) + list_permuted_items.append(item.permute(*perm)) + self._source._stack_onto_(list_permuted_items, new_dim) + return self + + @property + def names(self): + names = list(self._source.names) + return [names[i] for i in self.custom_op_kwargs["dims"]] + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a lazy tensordict cannot be modified. Call to_tensordict() first." + ) + + +def _iter_items_lazystack( + tensordict: LazyStackedTensorDict, return_none_for_het_values: bool = False +) -> Iterator[tuple[str, CompatibleType]]: + for key in tensordict.tensordicts[0].keys(): + values = tensordict._maybe_get_list(key) + if values is not None: + yield key, values + + +_register_tensor_class(LazyStackedTensorDict) +_register_tensor_class(_CustomOpTensorDict) +_register_tensor_class(_PermutedTensorDict) +_register_tensor_class(_SqueezedTensorDict) +_register_tensor_class(_UnsqueezedTensorDict) +_register_tensor_class(_TransposedTensorDict) +_register_tensor_class(_ViewedTensorDict) diff --git a/lib/python3.12/site-packages/tensordict/_nestedkey.py b/lib/python3.12/site-packages/tensordict/_nestedkey.py new file mode 100644 index 0000000000000000000000000000000000000000..730c115e2662c52badb775b2fd3c8a232ed3d681 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_nestedkey.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import abc + + +class _NestedKeyMeta(abc.ABCMeta): + def __instancecheck__(self, instance): + return isinstance(instance, str) or ( + isinstance(instance, tuple) + and len(instance) + and all(isinstance(subkey, NestedKey) for subkey in instance) + ) + + +class NestedKey(metaclass=_NestedKeyMeta): + """An abstract class for nested keys. + + Nested keys are the generic key type accepted by TensorDict. + + A nested key is either a string or a non-empty tuple of NestedKeys instances. + + The NestedKey class supports instance checks. + + """ + + pass diff --git a/lib/python3.12/site-packages/tensordict/_nestedkey.pyi b/lib/python3.12/site-packages/tensordict/_nestedkey.pyi new file mode 100644 index 0000000000000000000000000000000000000000..bb1b87b347058797943cc1a2acb4dfe5382a0214 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_nestedkey.pyi @@ -0,0 +1,10 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +from typing import Tuple, TypeAlias, Union + +NestedKeyType = Union[str, Tuple["NestedKeyType", ...]] +NestedKey: TypeAlias = NestedKeyType diff --git a/lib/python3.12/site-packages/tensordict/_pytree.py b/lib/python3.12/site-packages/tensordict/_pytree.py new file mode 100644 index 0000000000000000000000000000000000000000..af8d6d0c27bdeae2c3cc14cea38b2d309d04c8bc --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_pytree.py @@ -0,0 +1,272 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from collections import defaultdict +from typing import Any, Dict, List, Tuple + +import torch +from tensordict._lazy import LazyStackedTensorDict +from tensordict._td import _SubTensorDict, TensorDict, TensorDictBase +from tensordict.base import _NESTED_TENSORS_AS_LISTS +from tensordict.persistent import PersistentTensorDict +from tensordict.utils import _shape, implement_for + +try: + from torch.utils._pytree import Context, MappingKey, register_pytree_node +except ImportError: + from torch.utils._pytree import ( + _register_pytree_node as register_pytree_node, + Context, + ) + +PYTREE_REGISTERED_TDS = ( + _SubTensorDict, + TensorDict, + PersistentTensorDict, +) +PYTREE_REGISTERED_LAZY_TDS = (LazyStackedTensorDict,) + + +def _str_to_dict(str_spec: str) -> Tuple[List[str], str]: + if str_spec[1] != "(" or str_spec[-1] != ")": + raise ValueError( + f"string must have '(' as a second character and ')' in last position. Got {str_spec}." + ) + context_and_child_strings = str_spec[2:-1] + + child_strings = [] + context_strings = [] + nested_parentheses = 0 + start_index = 0 + for i, char in enumerate(context_and_child_strings): + if char == ":": + if nested_parentheses == 0: + context_strings.append(context_and_child_strings[start_index:i]) + start_index = i + 1 + elif char == "(": + nested_parentheses += 1 + elif char == ")": + nested_parentheses -= 1 + + if nested_parentheses == 0 and char == ",": + child_strings.append(context_and_child_strings[start_index:i]) + start_index = i + 1 + + child_strings.append(context_and_child_strings[start_index:]) + return context_strings, ",".join(child_strings) + + +def _str_to_tensordictdict(str_spec: str) -> Tuple[List[str], str]: + context_and_child_strings = str_spec[2:-1] + + child_strings = [] + context_strings = [] + nested_parentheses = 0 + start_index = 0 + for i, char in enumerate(context_and_child_strings): + if char == ":": + if nested_parentheses == 0: + context_strings.append(context_and_child_strings[start_index:i]) + start_index = i + 1 + elif char == "(": + nested_parentheses += 1 + elif char == ")": + nested_parentheses -= 1 + + if nested_parentheses == 0 and char == ",": + child_strings.append(context_and_child_strings[start_index:i]) + start_index = i + 1 + + child_strings.append(context_and_child_strings[start_index:]) + return context_strings, ",".join(child_strings) + + +def _tensordict_flatten(d: TensorDict) -> Tuple[List[Any], Context]: + items = tuple(d.items()) + if items: + keys, values = zip(*items) + keys = list(keys) + values = list(values) + else: + keys = [] + values = [] + return values, { + "keys": keys, + "batch_size": d.batch_size, + "names": d.names if d._has_names() else None, + "device": d.device, + "constructor": _constructor(type(d)), + "non_tensor_data": d.non_tensor_items(), + "cls": type(d), + } + + +def _lazy_tensordict_flatten(d: LazyStackedTensorDict) -> Tuple[List[Any], Context]: + return list(d.tensordicts), { + "stack_dim_name": d._td_dim_name, + "stack_dim": d.stack_dim, + "constructor": _lazy_tensordict_constructor, + "cls": type(d), + } + + +def _tensordict_unflatten(values: List[Any], context: Context) -> Dict[Any, Any]: + device = context["device"] + if device is not None: + device = ( + device + if all(val.device == device for val in values if hasattr(val, "device")) + else None + ) + batch_size = context["batch_size"] + names = context["names"] + keys = context["keys"] + constructor = context["constructor"] + non_tensor_items = context["non_tensor_data"] + cls = context["cls"] + batch_dims = len(batch_size) + if any(tensor is None for tensor in values): + return + if any(_shape(tensor)[:batch_dims] != batch_size for tensor in values): + batch_size = torch.Size([]) + names = None + return constructor( + cls=cls, + keys=keys, + values=values, + batch_size=batch_size, + names=names, + device=device, + non_tensor_items=non_tensor_items, + ) + + +def _lazy_tensordict_unflatten(values: List[Any], context: Context) -> Dict[Any, Any]: + stack_dim = context["stack_dim"] + return cls(*values, stack_dim=stack_dim, stack_dim_name=context["stack_dim_name"]) + + +def _td_flatten_with_keys( + d: TensorDictBase, +): + items = tuple(d.items(is_leaf=_NESTED_TENSORS_AS_LISTS)) + if items: + keys, values = zip(*items) + keys = list(keys) + values = list(values) + else: + keys = [] + values = [] + return [(MappingKey(k), v) for k, v in zip(keys, values)], { + "keys": keys, + "batch_size": d.batch_size, + "names": d._maybe_names(), + "device": d.device, + "constructor": _constructor(type(d)), + "non_tensor_data": d.non_tensor_items(), + "cls": type(d), + } + + +def _lazy_td_flatten_with_keys( + d: LazyStackedTensorDict, +): + raise NotImplementedError + + +@implement_for("torch", None, "2.3") +def _register_td_node(cls): + register_pytree_node( + cls, + _tensordict_flatten, + _tensordict_unflatten, + ) + + +@implement_for("torch", "2.3") +def _register_td_node(cls): # noqa: F811 + register_pytree_node( + cls, + _tensordict_flatten, + _tensordict_unflatten, + flatten_with_keys_fn=_td_flatten_with_keys, + ) + + +@implement_for("torch", None, "2.3") +def _register_lazy_td_node(cls): + register_pytree_node( + cls, + _lazy_tensordict_flatten, + _lazy_tensordict_unflatten, + ) + + +@implement_for("torch", "2.3") +def _register_lazy_td_node(cls): # noqa: F811 + register_pytree_node( + cls, + _lazy_tensordict_flatten, + _lazy_tensordict_unflatten, + flatten_with_keys_fn=_lazy_td_flatten_with_keys, + ) + + +def _constructor(cls): + return _CONSTRUCTORS[cls] + + +def _tensorclass_constructor( + *, cls, keys, values, batch_size, names, device, non_tensor_items +): + result = _tensordict_constructor( + cls=TensorDict, + keys=keys, + values=values, + batch_size=batch_size, + names=names, + device=device, + non_tensor_items=(), + ) + result = cls._from_tensordict(result, dict(non_tensor_items)) + return result + + +def _tensordict_constructor( + *, cls, keys, values, batch_size, names, device, non_tensor_items +): + result = cls._new_unsafe( + dict(zip(keys, values)), + batch_size=batch_size, + names=names, + device=device, + ) + for key, item in non_tensor_items: + result.set_non_tensor(key, item) + return result + + +def _lazy_tensordict_constructor( + *, cls, keys, values, batch_size, names, device, non_tensor_items +): + + result = cls._new_unsafe( + dict(zip(keys, values)), + batch_size=batch_size, + names=names, + device=device, + ) + for key, item in non_tensor_items: + result.set_non_tensor(key, item) + return result + + +_CONSTRUCTORS = defaultdict(lambda: _tensordict_constructor) +_CONSTRUCTORS[LazyStackedTensorDict] = _lazy_tensordict_constructor + + +for cls in PYTREE_REGISTERED_TDS: + _register_td_node(cls) +for cls in PYTREE_REGISTERED_LAZY_TDS: + _register_lazy_td_node(cls) diff --git a/lib/python3.12/site-packages/tensordict/_reductions.py b/lib/python3.12/site-packages/tensordict/_reductions.py new file mode 100644 index 0000000000000000000000000000000000000000..3c9bc2547d70d6f45a8d26b8f299368488dc4eaf --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_reductions.py @@ -0,0 +1,182 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import copyreg +from multiprocessing import reduction + +import torch +from tensordict._lazy import LazyStackedTensorDict +from tensordict._td import TensorDict + +from tensordict.tensorclass import NonTensorData, NonTensorStack +from tensordict.utils import _is_tensorclass, _STR_DTYPE_TO_DTYPE + +CLS_MAP = { + "TensorDict": TensorDict, + "LazyStackedTensorDict": LazyStackedTensorDict, + "NonTensorData": NonTensorData, + "NonTensorStack": NonTensorStack, +} + + +def _rebuild_tensordict_files(flat_key_values, metadata_dict, is_shared: bool = False): + def from_metadata(metadata=metadata_dict, prefix=None): + non_tensor = metadata.pop("non_tensors") + leaves = metadata.pop("leaves") + cls = metadata.pop("cls") + cls_metadata = metadata.pop("cls_metadata") + is_locked = cls_metadata.pop("is_locked", False) + + d = { + key: NonTensorData(data, batch_size=batch_size) + for (key, (data, batch_size)) in non_tensor.items() + } + for key in leaves.keys(): + total_key = (key,) if prefix is None else prefix + (key,) + if total_key[-1].startswith(""): + nested_values = flat_key_values[total_key] + nested_lengths = None + continue + if total_key[-1].startswith(""): + nested_lengths = flat_key_values[total_key] + continue + elif total_key[-1].startswith("", "") + value = torch.nested.nested_tensor_from_jagged( + nested_values, offsets=offsets, lengths=nested_lengths + ) + del nested_values + del nested_lengths + else: + value = flat_key_values[total_key] + d[key] = value + for k, v in metadata.items(): + # Each remaining key is a tuple pointing to a sub-tensordict + d[k] = from_metadata( + v, prefix=prefix + (k,) if prefix is not None else (k,) + ) + if isinstance(cls, str): + cls = CLS_MAP[cls] + result = cls._from_dict_validated(d, **cls_metadata) + if is_locked: + result.lock_() + # if is_shared: + # result._is_shared = is_shared + return result + + return from_metadata() + + +def _rebuild_tensordict_files_shared(flat_key_values, metadata_dict): + return _rebuild_tensordict_files(flat_key_values, metadata_dict, is_shared=True) + + +def _rebuild_tensordict_files_consolidated( + metadata, + storage, +): + def from_metadata(metadata=metadata, prefix=None): + consolidated = {"storage": storage, "metadata": metadata} + metadata = dict(metadata) + non_tensor = metadata.pop("non_tensors") + leaves = metadata.pop("leaves") + cls = metadata.pop("cls") + cls_metadata = dict(metadata.pop("cls_metadata")) + is_locked = cls_metadata.pop("is_locked", False) + # size can be there to tell what the size of the file is + _ = metadata.pop("size", None) + + d = { + key: NonTensorData( + data, + batch_size=batch_size, + device=torch.device(device) if device is not None else None, + ) + for (key, (data, batch_size, device)) in non_tensor.items() + } + for key, (dtype, local_shape, start, stop, pad) in leaves.items(): + dtype = _STR_DTYPE_TO_DTYPE[dtype] + # device = torch.device(device) + local_shape = torch.Size(local_shape) + value = storage[start:stop].view(dtype) + if pad: + value = value[: local_shape.numel()] + value = value.view(local_shape) + if key.startswith(""): + raise RuntimeError + elif key.startswith(""): + nested_values = value + nested_lengths = None + continue + elif key.startswith(""): + nested_lengths = value + continue + elif key.startswith(""): + from torch.nested._internal.nested_tensor import NestedTensor + + offsets = value + value = NestedTensor( + nested_values, offsets=offsets, lengths=nested_lengths + ) + key = key.replace("", "") + d[key] = value + for k, v in metadata.items(): + # Each remaining key is a tuple pointing to a sub-tensordict + d[k] = from_metadata( + v, prefix=prefix + (k,) if prefix is not None else (k,) + ) + if isinstance(cls, str): + cls = CLS_MAP[cls] + result = cls._from_dict_validated(d, **cls_metadata) + if is_locked: + result = result.lock_() + if _is_tensorclass(cls): + result._tensordict._consolidated = consolidated + else: + result._consolidated = consolidated + return result + + return from_metadata() + + +def _make_td(cls, state): + td = cls.__new__(cls) + td.__setstate__(state) + return td + + +def _reduce_td(data: TensorDict): + consolidated = getattr(data, "_consolidated", None) + if consolidated and consolidated["metadata"] is not None: + storage = consolidated["storage"] + storge_metadata = consolidated["metadata"] + return ( + _rebuild_tensordict_files_consolidated, + (storge_metadata, storage), + ) + + # This is faster than the solution below. + return ( + _make_td, + ( + type(data), + data.__getstate__(), + ), + ) + # metadata_dict, flat_key_values, _, _ = data._reduce_vals_and_metadata( + # requires_metadata=True + # ) + # return (_rebuild_tensordict_files, (flat_key_values, metadata_dict)) + + +reduction.register(TensorDict, _reduce_td) + +copyreg.pickle(TensorDict, _reduce_td) + +reduction.register(LazyStackedTensorDict, _reduce_td) + +copyreg.pickle(LazyStackedTensorDict, _reduce_td) diff --git a/lib/python3.12/site-packages/tensordict/_td.py b/lib/python3.12/site-packages/tensordict/_td.py new file mode 100644 index 0000000000000000000000000000000000000000..0bec65b856befa750a3db8d4ceee192db658cd52 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_td.py @@ -0,0 +1,5469 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import numbers +import os +import weakref +from collections import defaultdict +from concurrent.futures import Future, ThreadPoolExecutor, wait +from copy import copy +from numbers import Number +from pathlib import Path +from textwrap import indent +from typing import ( + Any, + Callable, + Dict, + Iterable, + Iterator, + List, + Sequence, + Tuple, + Type, + TYPE_CHECKING, +) +from warnings import warn + +import numpy as np + +import torch +from tensordict._nestedkey import NestedKey +from tensordict._tensorcollection import TensorCollection + +from tensordict.base import ( + _ACCEPTED_CLASSES, + _default_is_leaf, + _device_recorder, + _expand_to_match_shape, + _is_leaf_nontensor, + _is_tensor_collection, + _load_metadata, + _maybe_broadcast_other, + _NESTED_TENSORS_AS_LISTS, + _register_tensor_class, + _UNSET, + BEST_ATTEMPT_INPLACE, + CompatibleType, + is_tensor_collection, + NO_DEFAULT, + T, + TensorDictBase, +) +from tensordict.memmap import MemoryMappedTensor +from tensordict.utils import ( + _add_batch_dim_pre_hook, + _as_context_manager, + _BatchedUninitializedBuffer, + _BatchedUninitializedParameter, + _check_inbuild, + _clone_value, + _create_segments_from_int, + _create_segments_from_list, + _get_item, + _get_leaf_tensordict, + _get_shape_from_args, + _getitem_batch_size, + _index_preserve_data_ptr, + _infer_size_impl, + _is_shared, + _KEY_ERROR, + _LOCK_ERROR, + _maybe_correct_neg_dim, + _mismatch_keys, + _NON_STR_KEY_ERR, + _NON_STR_KEY_TUPLE_ERR, + _parse_to, + _pass_through, + _prune_selected_keys, + _set_item, + _set_max_batch_size, + _shape, + _STR_DTYPE_TO_DTYPE, + _StringKeys, + _StringOnlyDict, + _sub_index, + _unravel_key_to_tuple, + _zip_strict, + cache, + convert_ellipsis_to_idx, + DeviceType, + expand_as_right, + IndexType, + is_non_tensor, + is_tensorclass, + lock_blocked, + unravel_key, + unravel_key_list, +) +from torch import nn, Tensor +from torch._functorch.vmap import _maybe_remove_batch_dim +from torch.nn.parameter import UninitializedTensorMixin +from torch.nn.utils._named_member_accessor import swap_tensor +from torch.utils._pytree import tree_map + +try: + from functorch import dim as ftdim + + _has_funcdim = True +except ImportError: + from tensordict.utils import _ftdim_mock as ftdim + + _has_funcdim = False +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + +try: + from torch.nn.parameter import Buffer +except ImportError: + from tensordict.utils import Buffer + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + +_register_tensor_class(ftdim.Tensor) + +__base__setattr__ = torch.nn.Module.__setattr__ + +try: + from tensordict.utils import _import_and_wrap_functorch + + _add_batch_dim, _remove_batch_dim = _import_and_wrap_functorch( + "_add_batch_dim", "_remove_batch_dim" + ) +except ImportError: + + def _add_batch_dim(*args, **kwargs) -> Tensor: + raise NotImplementedError + + def _remove_batch_dim(*args, **kwargs) -> Tensor: + raise NotImplementedError + + +class TensorDict(TensorDictBase): + """A batched dictionary of tensors. + + TensorDict is a tensor container where all tensors are stored in a + key-value pair fashion and where each element shares the same first ``N`` + leading dimensions shape, where is an arbitrary number with ``N >= 0``. + + Additionally, if the tensordict has a specified device, then each element + must share that device. + + TensorDict instances support many regular tensor operations with the notable + exception of algebraic operations: + + - operations on shape: when a shape operation is called (indexing, + reshape, view, expand, transpose, permute, + unsqueeze, squeeze, masking etc), the operations is done as if it + was executed on a tensor of the same shape as the batch size then + expended to the right, e.g.: + + >>> td = TensorDict({'a': torch.zeros(3, 4, 5)}, batch_size=[3, 4]) + >>> # returns a TensorDict of batch size [3, 4, 1]: + >>> td_unsqueeze = td.unsqueeze(-1) + >>> # returns a TensorDict of batch size [12] + >>> td_view = td.view(-1) + >>> # returns a tensor of batch size [12, 4] + >>> a_view = td.view(-1).get("a") + + - casting operations: a TensorDict can be cast on a different device using + + >>> td_cpu = td.to("cpu") + >>> dictionary = td.to_dict() + + A call of the `.to()` method with a dtype will return an error. + + - Cloning (:meth:`~TensorDictBase.clone`), contiguous (:meth:`~TensorDictBase.contiguous`); + + - Reading: `td.get(key)`, `td.get_at(key, index)` + + - Content modification: :obj:`td.set(key, value)`, :obj:`td.set_(key, value)`, + :obj:`td.update(td_or_dict)`, :obj:`td.update_(td_or_dict)`, :obj:`td.fill_(key, + value)`, :obj:`td.rename_key_(old_name, new_name)`, etc. + + - Operations on multiple tensordicts: `torch.cat(tensordict_list, dim)`, + `torch.stack(tensordict_list, dim)`, `td1 == td2`, `td.apply(lambda x+y, other_td)` etc. + + Args: + source (TensorDict or Dict[NestedKey, Union[Tensor, TensorDictBase]]): a + data source. If empty, the tensordict can be populated subsequently. + A ``TensorDict`` can also be built via a sequence of keyword arguments, + as it is the case for ``dict(...)``. + batch_size (iterable of int, optional): a batch size for the + tensordict. The batch size can be modified subsequently as long + as it is compatible with its content. + If not batch-size is provided, an empty batch-size is assumed (it + is not inferred automatically from the data). To automatically set + the batch-size, refer to :meth:`~.auto_batch_size_`. + device (torch.device or compatible type, optional): a device for the + TensorDict. If provided, all tensors will be stored on that device. + If not, tensors on different devices are allowed. + names (lsit of str, optional): the names of the dimensions of the + tensordict. If provided, its length must match the one of the + ``batch_size``. Defaults to ``None`` (no dimension name, or ``None`` + for every dimension). + non_blocking (bool, optional): if ``True`` and a device is passed, the tensordict + is delivered without synchronization. This is the fastest option but is only + safe when casting from cpu to cuda (otherwise a synchronization call must be + implemented by the user). + If ``False`` is passed, every tensor movement will be done synchronously. + If ``None`` (default), the device casting will be done asynchronously but + a synchronization will be executed after creation if required. This option + should generally be faster than ``False`` and potentially slower than ``True``. + lock (bool, optional): if ``True``, the resulting tensordict will be + locked. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> source = {'random': torch.randn(3, 4), + ... 'zeros': torch.zeros(3, 4, 5)} + >>> batch_size = [3] + >>> td = TensorDict(source, batch_size=batch_size) + >>> print(td.shape) # equivalent to td.batch_size + torch.Size([3]) + >>> td_unqueeze = td.unsqueeze(-1) + >>> print(td_unqueeze.get("zeros").shape) + torch.Size([3, 1, 4, 5]) + >>> print(td_unqueeze[0].shape) + torch.Size([1]) + >>> print(td_unqueeze.view(-1).shape) + torch.Size([3]) + >>> print((td.clone()==td).all()) + True + + """ + + _td_dim_names = None + _is_shared = False + _is_memmap = False + _has_exclusive_keys = False + + def __init__( + self, + source: T | dict[NestedKey, CompatibleType] | None = None, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: DeviceType | None = None, + names: Sequence[str] | None = None, + non_blocking: bool | None = None, + lock: bool = False, + **kwargs: Any, + ) -> None: + if (source is not None) and kwargs: + raise ValueError( + "Either a dictionary or a sequence of kwargs must be provided, not both." + ) + source = source if not kwargs else kwargs + + self._tensordict = _StringOnlyDict() + + # if names and is_compiling(): + # graph_break() + has_device = device is not None + sub_non_blocking = False + call_sync = False + if has_device: + if non_blocking is None: + sub_non_blocking = True + else: + sub_non_blocking = non_blocking + device = torch.device(device) + # Auto-index the device + if device.index is None: + if device.type == "cuda": + device = torch.device( + device.type, index=torch.cuda.current_device() + ) + elif device.type not in ("cpu", "meta"): + device = torch.device(device.type, index=0) + if device.type == "cuda": + # CUDA does its sync by itself + call_sync = False + else: + call_sync = non_blocking is None + if call_sync: + _device_recorder.mark() + try: + self._device = device + + if source is None: + source = {} + if not isinstance(source, (TensorDictBase, dict)): + raise ValueError( + "A TensorDict source is expected to be a TensorDictBase " + f"sub-type or a dictionary, found type(source)={type(source)}." + ) + self._batch_size = self._parse_batch_size(source, batch_size) + # TODO: this breaks when stacking tensorclasses with dynamo + if not is_compiling(): + self._set_names(names) + + for key, value in source.items(): + self.set(key, value, non_blocking=sub_non_blocking) + if call_sync: + if _device_recorder.has_transfer(): + self._sync_all() + _device_recorder.unmark() + call_sync = False + + if lock: + self.lock_() + finally: + if call_sync: + _device_recorder.unmark() + + @classmethod + def _new_unsafe( + cls, + source: T | dict[NestedKey, CompatibleType] | None = None, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: DeviceType | None = None, + names: Sequence[str] | None = None, + non_blocking: bool | None = None, + lock: bool = False, + nested: bool = True, + is_shared: bool = False, + is_memmap: bool = False, + **kwargs: Any, + ) -> TensorDict: + if is_compiling() and cls is TensorDict: + # If the cls is not TensorDict, we must escape this to keep the same class. + # That's unfortunate because as of now it graph breaks but that's the best we can do. + td = TensorDict( + source, + batch_size=batch_size, + device=device, + names=names, + non_blocking=non_blocking, + lock=lock, + **kwargs, + ) + if is_shared: + td.share_memory_() + if is_memmap: + td.memmap_() + return td + if kwargs and not source: + source = kwargs + self = cls.__new__(cls) + sub_non_blocking = False + if device is not None: + if non_blocking is None: + sub_non_blocking = True + non_blocking = False + else: + sub_non_blocking = non_blocking + device = torch.device(device) if device is not None else None + if self._has_mps: + # With MPS, an explicit sync is required + sub_non_blocking = True + self._device = device + self._tensordict = _tensordict = _StringOnlyDict() + self._batch_size = batch_size + if source: # faster than calling items + for key, value in source.items(): + if nested and isinstance(value, dict): + value = cls._new_unsafe( + source=value, + batch_size=self._batch_size, + device=self._device, + non_blocking=sub_non_blocking, + ) + _tensordict[key] = value + # assert names is None or len(names) == self.batch_dims, (names, batch_size) + # assert (names is None) or (not all(name is None for name in names)) + self._td_dim_names = names + if lock: + self.lock_() + if is_shared: + self._is_shared = True + if is_memmap: + self._is_memmap = True + return self + + @classmethod + def from_module( + cls, + module: torch.nn.Module, + as_module: bool = False, + lock: bool = False, + use_state_dict: bool = False, + filter_empty: bool = True, + ): + result = cls._from_module( + module=module, + as_module=as_module, + use_state_dict=use_state_dict, + filter_empty=filter_empty, + ) + if result is None: + result = cls._new_unsafe({}, batch_size=torch.Size(())) + if lock: + result.lock_() + return result + + @classmethod + def _from_module( + cls, + module: torch.nn.Module, + as_module: bool = False, + use_state_dict: bool = False, + prefix="", + filter_empty: bool = True, + ): + from tensordict.nn import TensorDictParams + + if isinstance(module, TensorDictParams): + return module + destination = {} + if use_state_dict: + keep_vars = False + # do we need this feature atm? + local_metadata = {} + # if hasattr(destination, "_metadata"): + # destination._metadata[prefix[:-1]] = local_metadata + for hook in module._state_dict_pre_hooks.values(): + hook(module, prefix, keep_vars) + module._save_to_state_dict(destination, "", keep_vars) + else: + for name, param in module._parameters.items(): + if param is None: + continue + destination[name] = param + for name, buffer in module._buffers.items(): + if buffer is None: + continue + destination[name] = buffer + + if use_state_dict: + for hook in module._state_dict_hooks.values(): + hook_result = hook(module, destination, prefix, local_metadata) + if hook_result is not None: + destination = hook_result + if not filter_empty or destination: + destination_set = True + destination = cls._new_unsafe(destination, batch_size=torch.Size(())) + else: + destination_set = False + for name, submodule in module._modules.items(): + if submodule is not None: + subtd = cls._from_module( + module=submodule, + as_module=False, + use_state_dict=use_state_dict, + prefix=prefix + name + ".", + filter_empty=filter_empty, + ) + if subtd is not None: + if not destination_set: + destination = cls._new_unsafe(batch_size=torch.Size(())) + destination_set = True + destination._set_str( + name, subtd, validated=True, inplace=False, non_blocking=False + ) + if not destination_set: + return + + if as_module: + from tensordict.nn.params import TensorDictParams + + return TensorDictParams(destination, no_convert=True) + return destination + + def is_empty(self): + + for item in self._tensordict.values(): + # we need to check if item is empty + if _is_tensor_collection(type(item)): + if not item.is_empty(): + return False + + if is_non_tensor(item): + return False + else: + return False + return True + + def _to_module( + self, + module: nn.Module, + *, + inplace: bool | None = None, + return_swap: bool = True, + swap_dest=None, + memo=None, + use_state_dict: bool = False, + non_blocking: bool = False, + is_dynamo: bool | None = None, + ): + if is_dynamo is None: + is_dynamo = is_compiling() + if is_dynamo: + _check_inbuild() + + if not use_state_dict and isinstance(module, TensorDictBase): + if return_swap: + swap = module.copy() + module._param_td = getattr(self, "_param_td", self) + return swap + else: + module.update(self) + return + + hooks = memo["hooks"] + if return_swap: + _swap = {} + if not is_dynamo: + memo[weakref.ref(module)] = _swap + + if use_state_dict: + if inplace is not None: + raise RuntimeError( + "inplace argument cannot be passed when use_state_dict=True." + ) + # execute module's pre-hooks + state_dict = self.flatten_keys(".") + prefix = "" + strict = True + local_metadata = {} + missing_keys = [] + unexpected_keys = [] + error_msgs = [] + for hook in module._load_state_dict_pre_hooks.values(): + hook( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def convert_type(x, y): + if isinstance(y, nn.Parameter): + return nn.Parameter(x) + if isinstance(y, Buffer): + return Buffer(x) + return x + + input = state_dict.unflatten_keys(".")._fast_apply( + convert_type, self, propagate_lock=True + ) + else: + input = self + inplace = bool(inplace) + + # we use __dict__ directly to avoid the getattr/setattr overhead whenever we can + if not is_dynamo and type(module).__setattr__ is __base__setattr__: + # if type(module).__setattr__ is __base__setattr__: + __dict__ = module.__dict__ + _parameters = __dict__["_parameters"] + _buffers = __dict__["_buffers"] + else: + __dict__ = None + + for key, value in input.items(): + if isinstance(value, (Tensor, ftdim.Tensor)): + # For Dynamo, we use regular set/delattr as we're not + # much afraid by overhead (and dynamo doesn't like those + # hacks we're doing). + if __dict__ is not None: + # if setattr is the native nn.Module.setattr, we can rely on _set_tensor_dict + local_out = _set_tensor_dict( + __dict__, + _parameters, + _buffers, + hooks, + module, + key, + value, + inplace, + ) + else: + if not inplace: + local_out = swap_tensor(module, key, value) + else: + new_val = local_out + if return_swap: + local_out = local_out.clone() + new_val.data.copy_(value.data, non_blocking=non_blocking) + else: + if __dict__ is not None: + child = __dict__["_modules"][key] + else: + child = module._modules.get(key) + + if not is_dynamo: + local_out = memo.get(weakref.ref(child), NO_DEFAULT) + + if is_dynamo or local_out is NO_DEFAULT: + local_out = value._to_module( + child, + inplace=inplace, + return_swap=return_swap, + swap_dest={}, # we'll be calling update later + memo=memo, + use_state_dict=use_state_dict, + non_blocking=non_blocking, + is_dynamo=is_dynamo, + ) + + if return_swap: + _swap[key] = local_out + + if return_swap: + if isinstance(swap_dest, dict): + return _swap + elif swap_dest is not None: + + def _quick_set(swap_dict, swap_td): + for key, val in swap_dict.items(): + if isinstance(val, dict): + _quick_set(val, swap_td._get_str(key, default=NO_DEFAULT)) + elif swap_td._get_str(key, None) is not val: + swap_td._set_str( + key, + val, + inplace=False, + validated=True, + non_blocking=non_blocking, + ) + + _quick_set(_swap, swap_dest) + return swap_dest + else: + return self._new_unsafe(_swap, batch_size=torch.Size(())) + + @_maybe_broadcast_other("__ne__") + def __ne__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other != self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + raise KeyError( + f"keys in {self} and {other} mismatch, got {keys1} and {keys2}" + ) + d = {} + for key, item1 in self.items(): + d[key] = item1 != other.get(key) + return TensorDict(batch_size=self.batch_size, source=d, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value != other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return True + + @_maybe_broadcast_other("__xor__") + def __xor__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other ^ self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + raise KeyError( + f"keys in {self} and {other} mismatch, got {keys1} and {keys2}" + ) + d = {} + for key, item1 in self.items(): + d[key] = item1 ^ other.get(key) + return TensorDict(batch_size=self.batch_size, source=d, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value ^ other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return True + + @_maybe_broadcast_other("__or__") + def __or__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other | self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + raise KeyError( + f"keys in {self} and {other} mismatch, got {keys1} and {keys2}" + ) + d = {} + for key, item1 in self.items(): + d[key] = item1 | other.get(key) + return TensorDict(batch_size=self.batch_size, source=d, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value | other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + @_maybe_broadcast_other("__eq__") + def __eq__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other == self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + _mismatch_keys(keys1, keys2) + d = {} + for key, item1 in self.items(): + d[key] = item1 == other.get(key) + return TensorDict(source=d, batch_size=self.batch_size, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value == other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + @_maybe_broadcast_other("__ge__") + def __ge__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other <= self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + _mismatch_keys(keys1, keys2) + d = {} + for key, item1 in self.items(): + d[key] = item1 >= other.get(key) + return TensorDict(source=d, batch_size=self.batch_size, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value >= other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + @_maybe_broadcast_other("__gt__") + def __gt__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other < self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + _mismatch_keys(keys1, keys2) + d = {} + for key, item1 in self.items(): + d[key] = item1 > other.get(key) + return TensorDict(source=d, batch_size=self.batch_size, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value > other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + @_maybe_broadcast_other("__le__") + def __le__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other >= self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + _mismatch_keys(keys1, keys2) + d = {} + for key, item1 in self.items(): + d[key] = item1 <= other.get(key) + return TensorDict(source=d, batch_size=self.batch_size, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value <= other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + @_maybe_broadcast_other("__lt__") + def __lt__(self, other: Any) -> Self | bool: + if is_tensorclass(other): + return other > self + if isinstance(other, (dict,)): + other = self.from_dict_instance(other, auto_batch_size=False) + if _is_tensor_collection(type(other)): + keys1 = set(self.keys()) + keys2 = set(other.keys()) + if len(keys1.difference(keys2)) or len(keys1) != len(keys2): + _mismatch_keys(keys1, keys2) + d = {} + for key, item1 in self.items(): + d[key] = item1 < other.get(key) + return TensorDict(source=d, batch_size=self.batch_size, device=self.device) + if isinstance(other, (numbers.Number, Tensor)): + return TensorDict( + {key: value < other for key, value in self.items()}, + self.batch_size, + device=self.device, + ) + return False + + def __setitem__( + self, + index: IndexType, + value: Any, + ) -> None: + istuple = isinstance(index, tuple) + if istuple or isinstance(index, str): + # try: + index_unravel = _unravel_key_to_tuple(index) + if index_unravel: + self._set_tuple( + index_unravel, + value, + inplace=( + BEST_ATTEMPT_INPLACE + if isinstance(self, _SubTensorDict) + else False + ), + validated=False, + non_blocking=False, + ) + return + + # we must use any and because using Ellipsis in index can break with some indices + if index is Ellipsis or ( + isinstance(index, tuple) and any(idx is Ellipsis for idx in index) + ): + index = convert_ellipsis_to_idx(index, self.batch_size) + # Convert index like (True,) or True to (0,) over unsqueezed self + if isinstance(index, tuple) and len(index) == 1: + index = index[0] + if isinstance(index, (bool, type(None))) or ( + isinstance(index, torch.Tensor) + and index.shape == () + and index.dtype == torch.bool + and index.all() + ): + with self.unsqueeze(0) as td_unsqueezed: + td_unsqueezed[:] = value + return + + if isinstance(value, (TensorDictBase, dict)): + indexed_bs = _getitem_batch_size(self.batch_size, index) + if isinstance(value, dict): + value = self.from_dict_instance( + value, batch_size=indexed_bs, device=self.device + ) + elif value.device != self.device: + value = value.to(self.device) + # value = self.empty(recurse=True)[index].update(value) + if value.batch_size != indexed_bs: + if value.shape == indexed_bs[-len(value.shape) :]: + # try to expand on the left (broadcasting) + value = value.expand(indexed_bs) + else: + try: + # copy and change batch_size if can't be expanded + value = value.copy() + value.batch_size = indexed_bs + except RuntimeError as err: + raise RuntimeError( + f"indexed destination TensorDict batch size is {indexed_bs} " + f"(batch_size = {self.batch_size}, index={index}), " + f"which differs from the source batch size {value.batch_size}" + ) from err + + keys = set(self.keys()) + subtd = None + for value_key, item in value.items(): + if value_key in keys: + self._set_at_str( + value_key, item, index, validated=True, non_blocking=False + ) + else: + if subtd is None: + subtd = self._get_sub_tensordict(index) + subtd.set(value_key, item, inplace=True, non_blocking=False) + else: + for key in self.keys(): + self.set_at_(key, value, index) + + def all(self, dim: int | None = None) -> bool | TensorCollection: + if dim is not None and (dim >= self.batch_dims or dim < -self.batch_dims): + raise RuntimeError( + "dim must be greater than or equal to -tensordict.batch_dims and " + "smaller than tensordict.batch_dims" + ) + if dim is not None: + dim = _maybe_correct_neg_dim(dim, self.batch_size) + + names = None + if self._has_names(): + names = [name for i, name in enumerate(self.names) if i != dim] + + return TensorDict( + source={key: value.all(dim=dim) for key, value in self.items()}, + batch_size=[b for i, b in enumerate(self.batch_size) if i != dim], + device=self.device, + names=names, + ) + return all(value.all() for value in self.values()) + + def any(self, dim: int | None = None) -> bool | TensorCollection: + if dim is not None and (dim >= self.batch_dims or dim < -self.batch_dims): + raise RuntimeError( + "dim must be greater than or equal to -tensordict.batch_dims and " + "smaller than tensordict.batch_dims" + ) + if dim is not None: + dim = _maybe_correct_neg_dim(dim, self.batch_size) + + names = None + if self._has_names(): + names = [name for i, name in enumerate(self.names) if i != dim] + + return TensorDict( + source={key: value.any(dim=dim) for key, value in self.items()}, + batch_size=[b for i, b in enumerate(self.batch_size) if i != dim], + device=self.device, + names=names, + ) + return any([value.any() for value in self.values()]) + + def _cast_reduction( + self, + *, + reduction_name, + dim=NO_DEFAULT, + keepdim=NO_DEFAULT, + tuple_ok=True, + further_reduce: bool, + values_only: bool = True, + call_on_nested: bool = True, + batch_size=None, + **kwargs, + ): + if further_reduce: + # It is not very memory-efficient to do this, but it's the easiest to cover all use cases + if dim is NO_DEFAULT: + agglomerate = [ + val.contiguous().flatten() + for val in self._values_list( + True, True, is_leaf=_NESTED_TENSORS_AS_LISTS + ) + ] + agglomerate = torch.cat(agglomerate, dim=0) + if reduction_name == "quantile": + q = kwargs.pop("q") + return getattr(torch, reduction_name)(agglomerate, q, **kwargs) + return getattr(torch, reduction_name)(agglomerate, **kwargs) + else: + agglomerate = list( + self._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + ) + if dim == "feature": + agglomerate = [ + ( + val.flatten(self.ndim, -1) + if val.ndim > self.ndim + else val.unsqueeze(-1) + ) + for val in agglomerate + ] + cat_dim = -1 + dim = -1 + keepdim = False + elif isinstance(dim, tuple): + cat_dim = dim[0] + else: + cat_dim = dim + agglomerate = torch.cat(agglomerate, dim=cat_dim) + kwargs_copy = {} + if keepdim is not NO_DEFAULT: + kwargs_copy["keepdim"] = keepdim + if reduction_name == "quantile": + q = kwargs.pop("q") + kwargs_copy.update(kwargs) + return getattr(torch, reduction_name)( + agglomerate, q, dim=dim, **kwargs_copy + ) + kwargs_copy.update(kwargs) + return getattr(torch, reduction_name)( + agglomerate, dim=dim, **kwargs_copy + ) + + # IMPORTANT: do not directly access batch_dims (or any other property) + # via self.batch_dims otherwise a reference cycle is introduced + def proc_dim(dim, batch_dims, tuple_ok=True): + if dim is None: + return dim + if isinstance(dim, tuple): + if tuple_ok: + return tuple( + _d + for d in dim + for _d in proc_dim(d, batch_dims, tuple_ok=False) + ) + return dim + return (_maybe_correct_neg_dim(dim, None, batch_dims),) + + dim_needs_proc = (dim is not NO_DEFAULT) and (dim not in ("feature",)) + if dim_needs_proc: + dim = proc_dim(dim, self.batch_dims, tuple_ok=tuple_ok) + if not tuple_ok: + dim = dim[0] + if dim in ("feature",): + if keepdim: + raise TypeError("dim='feature' is incompatible with keepdim=True.") + + ndim = self.ndim + + def reduction(val): + if _is_tensor_collection(type(val)): + local_dim = dim + else: + if val.ndim > ndim: + val = val.flatten(ndim, -1) + else: + val = val.unsqueeze(-1) + local_dim = -1 + if reduction_name == "quantile": + # Make a copy of kwargs to avoid consuming q multiple times + kwargs_copy = kwargs.copy() + q = kwargs_copy.pop("q") + result = getattr(val, reduction_name)(q, local_dim, **kwargs_copy) + else: + result = getattr(val, reduction_name)(dim=local_dim, **kwargs) + if isinstance(result, tuple): + if values_only: + result = result.values + else: + return TensorDict.from_namedtuple(result) + return result + + if self._has_names(): + names = list(self.names) + else: + names = None + if not call_on_nested: + raise RuntimeError( + f"reduction {reduction_name} must be called with call_on_nested=True when dim='feature'." + ) + return self._fast_apply( + reduction, + call_on_nested=call_on_nested, + device=self.device, + names=names, + ) + + elif dim is not NO_DEFAULT or keepdim: + names = None + if self._has_names(): + if not keepdim and isinstance(dim, tuple): + names = [name for i, name in enumerate(self.names) if i not in dim] + else: + names = [name for i, name in enumerate(self.names) if i != dim] + if dim is not NO_DEFAULT: + kwargs["dim"] = dim + if keepdim is not NO_DEFAULT: + kwargs["keepdim"] = keepdim + + def reduction(val): + if reduction_name == "quantile": + # Make a copy of kwargs to avoid consuming q multiple times + kwargs_copy = kwargs.copy() + q = kwargs_copy.pop("q") + # Handle dim parameter properly for quantile + if "dim" in kwargs_copy: + dim_val = kwargs_copy.pop("dim") + # torch.quantile doesn't support tuple dimensions, so we need to handle this + if isinstance(dim_val, tuple): + # For tuple dimensions, we'll use the first dimension + # This is a limitation of torch.quantile compared to other reductions + dim_val = dim_val[0] + result = getattr(val, reduction_name)(q, dim_val, **kwargs_copy) + else: + result = getattr(val, reduction_name)(q, **kwargs_copy) + else: + result = getattr(val, reduction_name)(**kwargs) + if isinstance(result, tuple): + if values_only: + result = result.values + else: + return TensorDict.from_namedtuple(result) + return result + + if batch_size is not None: + pass + elif dim is not None and dim is not NO_DEFAULT: + if not keepdim: + if isinstance(dim, tuple): + batch_size = [ + b for i, b in enumerate(self.batch_size) if i not in dim + ] + else: + batch_size = [ + b for i, b in enumerate(self.batch_size) if i != dim + ] + else: + if isinstance(dim, tuple): + batch_size = [ + b if i not in dim else 1 + for i, b in enumerate(self.batch_size) + ] + else: + batch_size = [ + b if i != dim else 1 for i, b in enumerate(self.batch_size) + ] + + else: + batch_size = [1 for b in self.batch_size] + + return self._fast_apply( + reduction, + call_on_nested=call_on_nested, + batch_size=torch.Size(batch_size), + device=self.device, + names=names, + ) + + def reduction(val): + if reduction_name == "quantile": + # Make a copy of kwargs to avoid consuming q multiple times + kwargs_copy = kwargs.copy() + q = kwargs_copy.pop("q") + return getattr(val, reduction_name)(q, **kwargs_copy) + return getattr(val, reduction_name)(**kwargs) + + return self._fast_apply( + reduction, + call_on_nested=True, + batch_size=torch.Size([]), + device=self.device, + names=None, + ) + + def _multithread_apply_flat( + self, + fn: Callable, + *others: T, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + is_leaf: Callable[[Type], bool] | None = None, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + ) -> None: + if is_leaf is None: + is_leaf = _default_is_leaf + for key, item in self.items(): + if ( + not call_on_nested + and not is_leaf(type(item)) + # and not is_non_tensor(item) + ): + if default is not NO_DEFAULT: + _others = [_other._get_str(key, default=None) for _other in others] + _others = [ + self.empty(recurse=True) if _other is None else _other + for _other in _others + ] + else: + _others = [ + _other._get_str(key, default=NO_DEFAULT) for _other in others + ] + local_futures.append([]) + item._multithread_apply_flat( + fn, + *_others, + named=named, + nested_keys=nested_keys, + prefix=prefix + (key,), + is_leaf=is_leaf, + executor=executor, + futures=futures, + local_futures=local_futures[-1], + ) + else: + _others = [_other._get_str(key, default=default) for _other in others] + if named: + if nested_keys: + future = executor.submit( + fn, prefix + (key,) if prefix != () else key, item, *_others + ) + else: + future = executor.submit(fn, key, item, *_others) + else: + future = executor.submit(fn, item, *_others) + futures.append(future) + local_futures.append(future) + + def _multithread_rebuild( + self, + *, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + out: TensorCollection | None = None, + filter_empty: bool = False, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + subs_results: Dict[Future, Any] | None = None, + multithread_set: bool = False, # Experimental + **constructor_kwargs, + ) -> None: + if constructor_kwargs: + raise RuntimeError( + f"constructor_kwargs not supported for class {type(self).__name__}." + ) + # Rebuilds a tensordict from the futures of its leaves + if inplace: + result = self + is_locked = result.is_locked + elif out is not None: + result = out + if out.is_locked: + raise RuntimeError(_LOCK_ERROR) + is_locked = False + if batch_size is not None and batch_size != out.batch_size: + raise RuntimeError( + "batch_size and out.batch_size must be equal when both are provided." + ) + if device is not NO_DEFAULT and device != out.device: + raise RuntimeError( + "device and out.device must be equal when both are provided." + ) + else: + + def make_result(names=names, batch_size=batch_size): + if names is NO_DEFAULT: + if batch_size is not None: + # erase names + names = None + elif batch_size is None: + names = self.names if self._has_names() else None + return self.empty(batch_size=batch_size, device=device, names=names) + + result = make_result() + is_locked = False + + any_set = set() + + if isinstance(result, _SubTensorDict): + + def setter( + item_trsf, + key, + inplace=inplace, + result=result, + ): + set_item = item_trsf is not None + any_set.add(set_item) + if not set_item: + return + result.set(key, item_trsf, inplace=inplace) + + elif checked and isinstance(result, TensorDict) and (inplace is not True): + + def setter( + item_trsf, + key, + result=result, + ): + set_item = item_trsf is not None + any_set.add(set_item) + if not set_item: + return + result._tensordict[key] = item_trsf + + else: + + local_inplace = BEST_ATTEMPT_INPLACE if inplace else False + + def setter( + item_trsf, + key, + result=result, + checked=checked, + ): + set_item = item_trsf is not None + any_set.add(set_item) + if not set_item: + return + + result._set_str( + key, + item_trsf, + inplace=local_inplace, + validated=checked, + non_blocking=False, + ) + + for i, (key, local_future) in enumerate( + _zip_strict(self.keys(), local_futures) + ): + + if isinstance(local_future, list): + # We can't make this a future as it could cause deadlocks: + # If we put a future over the root and this triggers another + # call on the leaves, the root will occupy a spot in the execution queue + # and wait for completion, potentially preventing the leaf of + # getting in the execution queue at all. + td = self._get_str(key, default=None) + item_trsf = td._multithread_rebuild( + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=checked, + out=out, + filter_empty=filter_empty, + executor=executor, + futures=futures, + local_futures=local_future, + subs_results=subs_results, + multithread_set=multithread_set, + **constructor_kwargs, + ) + if multithread_set: + local_future = executor.submit(setter, item_trsf=item_trsf, key=key) + local_futures[i] = local_future + futures.append(local_future) + else: + setter(item_trsf=item_trsf, key=key) + else: + if multithread_set: + if subs_results is not None: + local_result = subs_results[local_future] + else: + # TODO: check if add_done_callback can safely be used here + # The issue is that it does not raises an exception encountered during the + # execution, resulting in UBs. + local_result = local_future.result() + local_future = executor.submit( + setter, item_trsf=local_result, key=key + ) + futures.append(local_future) + local_futures[i] = local_future + else: + local_result = local_future.result() + setter(item_trsf=local_result, key=key) + + if multithread_set: + wait(local_futures) + any_set = True in any_set or is_non_tensor(self) + + if filter_empty and not any_set: + return + elif not filter_empty and not inplace and is_locked: + result.lock_() + return result + + def _apply_nest( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + filter_empty: bool | None = None, + is_leaf: Callable[[Type], bool] | None = None, + out: TensorDictBase | None = None, + **constructor_kwargs, + ) -> Self | None: + if inplace: + result = self + is_locked = result.is_locked + elif out is not None: + result = out + if out.is_locked: + raise RuntimeError(_LOCK_ERROR) + is_locked = False + if batch_size is not None and batch_size != out.batch_size: + raise RuntimeError( + "batch_size and out.batch_size must be equal when both are provided." + ) + if device is not NO_DEFAULT and device != out.device: + if not checked: + raise RuntimeError( + f"device and out.device must be equal when both are provided. Got device={device} and out.device={out.device}." + ) + else: + device = torch.device(device) + out._device = device + for node in out.values(True, True, is_leaf=_is_tensor_collection): + if is_tensorclass(node): + node._tensordict._device = device + else: + node._device = device + else: + + def make_result(names=names, batch_size=batch_size): + if names is NO_DEFAULT: + if batch_size is not None: + # erase names + names = None + else: + names = self.names if self._has_names() else None + return self.empty(batch_size=batch_size, device=device, names=names) + + result = None + is_locked = False + + any_set = False + if is_leaf is None: + is_leaf = _default_is_leaf + + for key, item in self.items(): + if ( + not call_on_nested + and not is_leaf(type(item)) + # and not is_non_tensor(item) + ): + if default is not NO_DEFAULT: + _others = [_other._get_str(key, default=None) for _other in others] + _others = [ + self.empty(recurse=True) if _other is None else _other + for _other in _others + ] + else: + _others = [ + _other._get_str(key, default=NO_DEFAULT) for _other in others + ] + + item_trsf = item._apply_nest( + fn, + *_others, + inplace=inplace, + batch_size=batch_size, + device=device, + checked=checked, + named=named, + nested_keys=nested_keys, + default=default, + prefix=prefix + (key,), + filter_empty=filter_empty, + is_leaf=is_leaf, + out=out._get_str(key, default=None) if out is not None else None, + **constructor_kwargs, + ) + else: + _others = [_other._get_str(key, default=default) for _other in others] + if named: + if nested_keys: + item_trsf = fn( + prefix + (key,) if prefix != () else key, item, *_others + ) + else: + item_trsf = fn(key, item, *_others) + else: + item_trsf = fn(item, *_others) + if item_trsf is not None: + if not any_set: + if result is None: + result = make_result() + any_set = True + if isinstance(self, _SubTensorDict): + result.set(key, item_trsf, inplace=inplace) + else: + result._set_str( + key, + item_trsf, + inplace=BEST_ATTEMPT_INPLACE if inplace else False, + validated=checked, + non_blocking=False, + ) + + if filter_empty and not any_set: + return + elif filter_empty is None and not any_set and not self.is_empty(): + # we raise the deprecation warning only if the tensordict wasn't already empty. + # After we introduce the new behaviour, we will have to consider what happens + # to empty tensordicts by default: will they disappear or stay? + return + if result is None: + result = make_result() + + if not inplace and is_locked: + result.lock_() + return result + + # Functorch compatibility + @cache # noqa: B019 + def _add_batch_dim(self, *, in_dim: int, vmap_level: int) -> Self: + td = self + + def _add_batch_dim_wrapper(key: str, value: Any) -> Any: + if is_tensor_collection(value): + return value._add_batch_dim(in_dim=in_dim, vmap_level=vmap_level) + + if isinstance( + value, (_BatchedUninitializedParameter, _BatchedUninitializedBuffer) + ): + value.in_dim = in_dim + value.vmap_level = vmap_level + return value + return _add_batch_dim(value, in_dim, vmap_level) + + out = self._new_unsafe( + {key: _add_batch_dim_wrapper(key, value) for key, value in td.items()}, + batch_size=torch.Size( + [b for i, b in enumerate(td.batch_size) if i != in_dim] + ), + names=( + [name for i, name in enumerate(td.names) if i != in_dim] + if self._has_names() + else None + ), + lock=self.is_locked, + ) + return out + + @cache # noqa: B019 + def _remove_batch_dim(self, vmap_level: int, batch_size: int, out_dim: int) -> Self: + new_batch_size = list(self.batch_size) + new_batch_size.insert(out_dim, batch_size) + names = self._maybe_names() + if names: + new_names = list(names) + new_names.insert(out_dim, None) + else: + new_names = None + out = TensorDict( + { + key: ( + value._remove_batch_dim( + vmap_level=vmap_level, batch_size=batch_size, out_dim=out_dim + ) + if is_tensor_collection(value) + else _remove_batch_dim(value, vmap_level, batch_size, out_dim) + ) + for key, value in self.items() + }, + batch_size=new_batch_size, + names=new_names, + lock=self.is_locked, + ) + return out + + @cache # noqa: B019 + def _maybe_remove_batch_dim( + self, funcname: str, vmap_level: int, batch_size: int, out_dim: int + ) -> Self: + new_batch_size = list(self.batch_size) + new_batch_size.insert(out_dim, batch_size) + names = self._maybe_names() + if names: + new_names = list(names) + new_names.insert(out_dim, None) + else: + new_names = None + out = TensorDict( + { + key: ( + value._maybe_remove_batch_dim( + funcname=funcname, + vmap_level=vmap_level, + batch_size=batch_size, + out_dim=out_dim, + ) + if is_tensor_collection(value) + else _maybe_remove_batch_dim( + funcname, value, vmap_level, batch_size, out_dim + ) + ) + for key, value in self.items() + }, + batch_size=new_batch_size, + names=new_names, + lock=self.is_locked, + ) + return out + + def _convert_to_tensordict( + self, dict_value: dict[str, Any], non_blocking: bool | None = None + ) -> Self: + return TensorDict( + dict_value, + batch_size=self.batch_size, + device=self.device, + names=self._maybe_names(), + lock=self.is_locked, + non_blocking=non_blocking, + ) + + def _index_tensordict( + self, + index: IndexType, + new_batch_size: torch.Size | None = None, + names: List[str] | None = None, + ) -> Self: + batch_size = self.batch_size + batch_dims = len(batch_size) + + def _check_for_invalid_index(index): + if batch_size: + return + if index is None: + return + if ( + isinstance(index, torch.Tensor) + and index.dtype == torch.bool + and not index.ndim + ): + return + if isinstance(index, tuple): + if len(index) == 1: + return _check_for_invalid_index(index[0]) + elif all(idx is None for idx in index): + return + raise RuntimeError( + f"indexing a tensordict with td.batch_dims==0 is not permitted. Got index {index}." + ) + + _check_for_invalid_index(index) + + if new_batch_size is not None: + batch_size = new_batch_size + else: + batch_size = _getitem_batch_size(batch_size, index) + + if names is None: + names = self._get_names_idx(index) + + source = {} + for key, item in self.items(): + if isinstance(item, TensorDict): + # this is the simplest case, we can pre-compute the batch size easily + new_batch_size = batch_size + item.batch_size[batch_dims:] + source[key] = item._index_tensordict( + index, new_batch_size=new_batch_size + ) + else: + source[key] = _get_item(item, index) + result = self._new_unsafe( + source=source, + batch_size=batch_size, + device=self.device, + names=names, + # lock=self.is_locked, + ) + if self._is_memmap and _index_preserve_data_ptr(index): + result._is_memmap = True + result.lock_() + elif self._is_shared and _index_preserve_data_ptr(index): + result._is_shared = True + result.lock_() + return result + + def expand(self, *args, **kwargs) -> Self: + tensordict_dims = self.batch_dims + shape = _get_shape_from_args(*args, **kwargs) + + # new shape dim check + if len(shape) < len(self.shape): + raise RuntimeError( + f"the number of sizes provided ({len(shape)}) must be greater or equal to the number of " + f"dimensions in the TensorDict ({tensordict_dims})" + ) + + # new shape compatibility check + for old_dim, new_dim in zip(self.batch_size, shape[-tensordict_dims:]): + if old_dim != 1 and new_dim != old_dim: + raise RuntimeError( + "Incompatible expanded shape: The expanded shape length at non-singleton dimension should be same " + f"as the original length. target_shape = {shape}, existing_shape = {self.batch_size}" + ) + + if self._has_names(): + names = [None] * (len(shape) - tensordict_dims) + self.names + else: + names = None + + def _expand(tensor): + tensor_shape = tensor.shape + tensor_dims = len(tensor_shape) + last_n_dims = tensor_dims - tensordict_dims + if last_n_dims > 0: + new_shape = (*shape, *tensor_shape[-last_n_dims:]) + else: + new_shape = shape + return tensor.expand(new_shape) + + return self._fast_apply( + _expand, + batch_size=shape, + call_on_nested=True, + names=names, + propagate_lock=True, + ) + + def _unbind(self, dim: int): + batch_size = torch.Size([s for i, s in enumerate(self.batch_size) if i != dim]) + names = None + if self._has_names(): + names = [name for i, name in enumerate(self.names) if i != dim] + # We could use any() but dynamo doesn't like generators + for name in names: + if name is not None: + break + else: + names = None + device = self.device + + is_shared = self._is_shared + is_memmap = self._is_memmap + + def empty( + batch_size=batch_size, + names=names, + device=device, + is_shared=is_shared, + is_memmap=is_memmap, + ): + result = self._new_unsafe( + {}, batch_size=batch_size, names=names, device=device + ) + result._is_shared = is_shared + result._is_memmap = is_memmap + return result + + tds = tuple(empty() for _ in range(self.batch_size[dim])) + + def unbind(key, val, tds=tds): + unbound = ( + val.unbind(dim) + if not isinstance(val, TensorDictBase) + # tensorclass is also unbound using plain unbind + else val._unbind(dim) + ) + for td, _val in _zip_strict(tds, unbound): + td._set_str( + key, _val, validated=True, inplace=False, non_blocking=False + ) + + for key, val in self.items(): + unbind(key, val) + return tds + + def split( + self, split_size: int | list[int], dim: int = 0 + ) -> tuple[TensorDictBase, ...]: + # we must use slices to keep the storage of the tensors + WRONG_TYPE = "split(): argument 'split_size' must be int or list of ints" + batch_size = self.batch_size + dim = _maybe_correct_neg_dim(dim, batch_size) + max_size = batch_size[dim] + if isinstance(split_size, int): + if split_size <= 0: + raise ValueError( + f"TensorDict.split: split_size must be positive, got {split_size}." + ) + split_size = min(split_size, max_size) + segments = _create_segments_from_int(split_size, max_size) + splits = [end - start for start, end in segments] + splits = {k: v.split(splits, dim) for k, v in self.items()} + elif isinstance(split_size, (list, tuple)): + if len(split_size) == 0: + raise RuntimeError("Insufficient number of elements in split_size.") + if not all(isinstance(x, int) for x in split_size): + raise TypeError(WRONG_TYPE) + splits = {k: v.split(split_size, dim) for k, v in self.items()} + segments = _create_segments_from_list(split_size, max_size) + else: + raise TypeError(WRONG_TYPE) + names = self._maybe_names() + batch_sizes = [ + torch.Size( + tuple(d if i != dim else end - start for i, d in enumerate(batch_size)) + ) + for start, end in segments + ] + splits = [ + {k: v[ss] for k, v in splits.items()} for ss in range(len(batch_sizes)) + ] + device = self.device + is_shared = self._is_shared + is_memmap = self._is_memmap + is_locked = self.is_locked + result = tuple( + self._new_unsafe( + source=split, + batch_size=bsz, + names=names, + device=device, + lock=is_locked, + is_shared=is_shared, + is_memmap=is_memmap, + ) + for split, bsz in _zip_strict(splits, batch_sizes) + ) + return result + + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorCollection, ...]: + if chunks < 1: + raise ValueError( + f"chunks must be a strictly positive integer, got {chunks}." + ) + # fall back on split, using upper rounding + batch_size = self.batch_size + dim = _maybe_correct_neg_dim(dim, batch_size) + max_size = batch_size[dim] + split_size = -(max_size // -chunks) + segments = _create_segments_from_int(split_size, max_size) + splits = {k: v.chunk(chunks, dim) for k, v in self.items()} + names = self._maybe_names() + batch_sizes = [ + torch.Size( + tuple(d if i != dim else end - start for i, d in enumerate(batch_size)) + ) + for start, end in segments + ] + splits = [ + {k: v[ss] for k, v in splits.items()} for ss in range(len(batch_sizes)) + ] + device = self.device + is_shared = self._is_shared + is_memmap = self._is_memmap + is_locked = self.is_locked + result = tuple( + self._new_unsafe( + source=split, + batch_size=bsz, + names=names, + device=device, + lock=is_locked, + is_shared=is_shared, + is_memmap=is_memmap, + ) + for split, bsz in _zip_strict(splits, batch_sizes) + ) + return result + + def masked_select(self, mask: Tensor) -> Self: + d = {} + mask_expand = mask + while mask_expand.ndimension() > self.batch_dims: + mndim = mask_expand.ndimension() + mask_expand = mask_expand.squeeze(-1) + if mndim == mask_expand.ndimension(): # no more squeeze + break + for key, value in self.items(): + d[key] = value[mask_expand] + dim = int(mask.sum().item()) + other_dim = self.shape[mask.ndim :] + return TensorDict( + device=self.device, source=d, batch_size=torch.Size([dim, *other_dim]) + ) + + def _view( + self, + *args, + **kwargs, + ) -> Self: + shape = _get_shape_from_args(*args, **kwargs) + if any(dim < 0 for dim in shape): + shape = _infer_size_impl(shape, self.numel()) + if torch.Size(shape) == self.shape: + return self + batch_dims = self.batch_dims + + def _view(tensor): + return tensor.view((*shape, *tensor.shape[batch_dims:])) + + result = self._fast_apply( + _view, batch_size=shape, call_on_nested=True, propagate_lock=True + ) + self._maybe_set_shared_attributes(result) + return result + + def reshape( + self, + *args, + **kwargs, + ) -> Self: + shape = _get_shape_from_args(*args, **kwargs) + if any(dim < 0 for dim in shape): + shape = _infer_size_impl(shape, self.numel()) + shape = torch.Size(shape) + if torch.Size(shape) == self.shape: + return self + batch_dims = self.batch_dims + + def _reshape(tensor): + return tensor.reshape((*shape, *tensor.shape[batch_dims:])) + + return self._fast_apply( + _reshape, + batch_size=shape, + call_on_nested=True, + propagate_lock=True, + ) + + def repeat_interleave( + self, + repeats: torch.Tensor | int, + dim: int | None = None, + *, + output_size: int | None = None, + ) -> Self: + if self.ndim == 0: + return self.unsqueeze(0).repeat_interleave( + repeats=repeats, dim=dim, output_size=output_size + ) + if dim is None: + if self.ndim > 1: + return self.reshape(-1).repeat_interleave(repeats, dim=0) + return self.repeat_interleave(repeats, dim=0) + dim_corrected = dim if dim >= 0 else self.ndim + dim + if not (dim_corrected >= 0): + raise ValueError( + f"dim {dim} is out of range for tensordict with shape {self.shape}." + ) + new_batch_size = [] + for i, s in enumerate(self.batch_size): + if i == dim_corrected: + if isinstance(repeats, int): + new_batch_size.append(s * repeats) + else: + new_batch_size.append(repeats.sum().item()) + else: + new_batch_size.append(s) + new_batch_size = torch.Size(new_batch_size) + + def rep(leaf): + return leaf.repeat_interleave( + repeats=repeats, dim=dim_corrected, output_size=output_size + ) + + return self._fast_apply( + rep, + batch_size=new_batch_size, + call_on_nested=True, + propagate_lock=True, + names=self._maybe_names(), + ) + + def _repeat(self, *repeats: int) -> TensorCollection: + new_batch_size = torch.Size([i * r for i, r in zip(self.batch_size, repeats)]) + + def rep(leaf): + return leaf.repeat(*repeats, *((1,) * (leaf.ndim - self.ndim))) + + return self._fast_apply( + rep, + batch_size=new_batch_size, + call_on_nested=True, + propagate_lock=True, + names=self._maybe_names(), + ) + + def _transpose(self, dim0, dim1): + def _transpose(tensor): + return tensor.transpose(dim0, dim1) + + batch_size = list(self.batch_size) + v0 = batch_size[dim0] + v1 = batch_size[dim1] + batch_size[dim1] = v0 + batch_size[dim0] = v1 + if self._has_names(): + names = self.names + names = [ + names[dim0] if i == dim1 else names[dim1] if i == dim0 else names[i] + for i in range(self.ndim) + ] + else: + names = None + result = self._fast_apply( + _transpose, + batch_size=torch.Size(batch_size), + call_on_nested=True, + names=names, + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + return result + + def _permute(self, *args, **kwargs): + dims_list = _get_shape_from_args(*args, kwarg_name="dims", **kwargs) + dims_list = [dim if dim >= 0 else self.ndim + dim for dim in dims_list] + if any(dim < 0 or dim >= self.ndim for dim in dims_list): + raise ValueError( + "Received an permutation order incompatible with the tensordict shape." + ) + # note: to allow this to work recursively, we must allow permutation order with fewer elements than dims, + # as long as this list is complete. + if not np.array_equal(sorted(dims_list), range(len(dims_list))): + raise ValueError( + f"Cannot compute the permutation, got dims={dims_list} but expected a permutation of {list(range(len(dims_list)))}." + ) + if not len(dims_list) and not self.batch_dims: + return self + if np.array_equal(dims_list, range(len(dims_list))): + return self + + def _permute(tensor): + return tensor.permute(*dims_list, *range(len(dims_list), tensor.ndim)) + + batch_size = self.batch_size + batch_size = [batch_size[p] for p in dims_list] + list( + batch_size[len(dims_list) :] + ) + if self._has_names(): + names = self.names + names = [names[i] for i in dims_list] + else: + names = None + result = self._fast_apply( + _permute, + batch_size=batch_size, + call_on_nested=True, + names=names, + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + return result + + def _squeeze(self, dim=None): + batch_size = self.batch_size + if dim is None: + names = list(self.names) if self._has_names() else None + if names is not None: + batch_size, names = _zip_strict( + *[ + (size, name) + for size, name in _zip_strict(batch_size, names) + if size != 1 + ] + ) + else: + batch_size = [size for size in batch_size if size != 1] + batch_size = torch.Size(batch_size) + if batch_size == self.batch_size: + return self + + # we only want to squeeze dimensions lower than the batch dim, and view + # is the perfect op for this + def _squeeze(tensor): + from tensordict import LazyStackedTensorDict + + if ( + isinstance(tensor, LazyStackedTensorDict) + and tensor.shape[tensor.stack_dim] == 1 + ): + return tensor.tensordicts[0] + return tensor.view(*batch_size, *tensor.shape[self.batch_dims :]) + + return self._fast_apply( + _squeeze, + batch_size=batch_size, + names=names, + inplace=False, + call_on_nested=True, + propagate_lock=True, + ) + # make the dim positive + newdim = _maybe_correct_neg_dim(dim, batch_size) + if batch_size[dim] != 1: + return self + batch_size = list(batch_size) + batch_size.pop(dim) + batch_size = list(batch_size) + names = list(self.names) if self._has_names() else None + if names: + names.pop(dim) + + def squeeze(x): + return x.squeeze(newdim) + + result = self._fast_apply( + squeeze, + batch_size=batch_size, + names=names, + inplace=False, + call_on_nested=True, + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + return result + + def _unsqueeze(self, dim: int): + # make the dim positive + if dim < 0: + newdim = self.batch_dims + dim + 1 + else: + newdim = dim + + if (newdim > self.batch_dims) or (newdim < 0): + raise RuntimeError( + f"unsqueezing is allowed for dims comprised between " + f"`-td.batch_dims - 1` and `td.batch_dims` only. Got " + f"dim={dim} with a batch size of {self.batch_size}." + ) + batch_size = list(self.batch_size) + batch_size.insert(newdim, 1) + batch_size = torch.Size(batch_size) + + names = list(self.names) if self._has_names() else None + if names: + names.insert(newdim, None) + + def _unsqueeze(tensor): + return tensor.unsqueeze(newdim) + + result = self._fast_apply( + _unsqueeze, + batch_size=batch_size, + names=names, + inplace=False, + call_on_nested=True, + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + return result + + @classmethod + def from_dict( + cls, + input_dict: dict[NestedKey, CompatibleType] | TensorDictBase, + *, + auto_batch_size: bool | None = None, + batch_size: Sequence[int] | torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: Sequence[str] | None = None, + ) -> Self: + if _is_tensor_collection(type(input_dict)): + return input_dict + + if batch_dims is not None and batch_size is not None: + raise ValueError( + "Cannot pass both batch_size and batch_dims to `from_dict`." + ) + + batch_size_set = torch.Size(()) if batch_size is None else batch_size + input_dict = dict(input_dict) + for key, value in list(input_dict.items()): + # we don't know if another tensor of smaller size is coming + # so we can't be sure that the batch-size will still be valid later + input_dict[key] = TensorDict.from_any( + value, + auto_batch_size=False, + device=device, + batch_size=batch_size, + ) + # regular __init__ breaks because a tensor may have the same batch-size as the tensordict + out = cls( + input_dict, + batch_size=batch_size_set, + device=device, + names=names, + ) + if batch_size is None: + if auto_batch_size is None and batch_dims is None: + auto_batch_size = False + elif auto_batch_size is None: + auto_batch_size = True + if auto_batch_size: + _set_max_batch_size(out, batch_dims) + else: + out.batch_size = batch_size + return out + + @classmethod + def _from_dict_validated( + cls, input_dict, batch_size=None, device=None, batch_dims=None, names=None + ): + return cls._new_unsafe( + input_dict, + batch_size=torch.Size(batch_size), + device=torch.device(device) if device is not None else device, + names=names if any(name is not None for name in names) else None, + ) + + def from_dict_instance( + self, + input_dict, + *, + auto_batch_size: bool | None = None, + batch_size=None, + device=None, + batch_dims=None, + names=None, + ): + if batch_dims is not None and batch_size is not None: + raise ValueError( + "Cannot pass both batch_size and batch_dims to `from_dict`." + ) + from tensordict import TensorDict + + batch_size_set = torch.Size(()) if batch_size is None else batch_size + if is_compiling(): + input_dict = type(input_dict)(input_dict) + else: + input_dict = copy(input_dict) + for key, value in list(input_dict.items()): + if isinstance(value, (dict,)): + cur_value = self.get(key) + if cur_value is not None: + input_dict[key] = cur_value.from_dict_instance( + value, + device=device, + auto_batch_size=False, + ) + continue + else: + # we don't know if another tensor of smaller size is coming + # so we can't be sure that the batch-size will still be valid later + input_dict[key] = TensorDict.from_dict( + value, + device=device, + auto_batch_size=False, + ) + else: + input_dict[key] = TensorDict.from_any( + value, + auto_batch_size=False, + ) + + out = TensorDict.from_dict( + input_dict, + batch_size=batch_size_set, + device=device, + names=names, + ) + if batch_size is None: + if auto_batch_size is None and batch_dims is None: + auto_batch_size = False + elif auto_batch_size is None: + auto_batch_size = True + if auto_batch_size: + _set_max_batch_size(out, batch_dims) + else: + out.batch_size = batch_size + return out + + @staticmethod + def _parse_batch_size( + source: T | dict | None, + batch_size: Sequence[int] | torch.Size | int | None = None, + ) -> torch.Size: + ERR = "batch size was not specified when creating the TensorDict instance and it could not be retrieved from source." + + if is_compiling(): + if isinstance(batch_size, torch.Size): + return batch_size + elif isinstance(batch_size, tuple): + return torch.Size(batch_size) + elif isinstance(batch_size, list): + return torch.Size(tuple(batch_size)) + if batch_size is None: + return torch.Size([]) + elif isinstance(batch_size, Number): + return torch.Size([batch_size]) + elif isinstance(source, TensorDictBase): + return source.batch_size + raise ValueError() + + try: + return torch.Size(batch_size) + except Exception: + if batch_size is None: + return torch.Size([]) + elif isinstance(batch_size, Number): + return torch.Size([batch_size]) + elif isinstance(source, TensorDictBase): + return source.batch_size + raise ValueError(ERR) + + @property + def batch_dims(self) -> int: + return len(self.batch_size) + + @batch_dims.setter + def batch_dims(self, value: int) -> None: + raise RuntimeError( + f"Setting batch dims on {type(self).__name__} instances is not allowed." + ) + + def _has_names(self): + return self._td_dim_names is not None + + def _erase_names(self): + self._td_dim_names = None + + @property + def names(self): + names = self._td_dim_names + if names is None: + return [None for _ in range(self.batch_dims)] + # assert len(names) == self.batch_dims, (names, self.batch_dims) + # Return a copy but don't use copy to make dynamo happy + return list(names) + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + # we don't run checks on types for efficiency purposes + if names is None: + self._rename_subtds(names) + self._erase_names() + return + value = list(names) + # Faster but incompatible with dynamo + # num_none = sum(v is None for v in value) + num_none = 0 + for v in value: + num_none += v is None + if num_none == self.batch_dims: + self._set_names(None) + return + if num_none: + num_none -= 1 + if len(set(value)) != len(value) - num_none: + raise ValueError(f"Some dimension names are non-unique: {value}.") + if len(value) != self.batch_dims: + raise ValueError( + "the length of the dimension names must equate the tensordict batch_dims attribute. " + f"Got {value} for batch_dims {self.batch_dims}." + ) + self._rename_subtds(value) + self._td_dim_names = list(value) + + def _rename_subtds(self, names): + if names is None: + for item in self._tensordict.values(): + if _is_tensor_collection(type(item)): + item._erase_names() + return + for item in self._tensordict.values(): + if isinstance(item, TensorDict): + # For TensorDict items, we can directly set _td_dim_names + item_names = item._td_dim_names + if item_names is None: + # Extend names with None for the remaining dimensions + td_names = list(names) + [None] * (item.batch_dims - len(names)) + else: + td_names = list(names) + list(item_names)[len(names) :] + item._td_dim_names = td_names + # Recursively rename nested tensor collections + item._rename_subtds(td_names) + elif _is_tensor_collection(type(item)): + # For other tensor collections (tensorclasses, NonTensorData, etc.), + # use the public API which handles the renaming correctly + item_names = item.names + td_names = list(names) + item_names[len(names) :] + item.rename_(*td_names) + + @property + def device(self) -> torch.device | None: + """Device of the tensordict. + + Returns `None` if device hasn't been provided in the constructor or set via `tensordict.to(device)`. + The `to()` method can also be used as a context manager for temporary device changes. + + """ + return self._device + + @device.setter + def device(self, value: DeviceType) -> None: + raise RuntimeError( + "device cannot be set using tensordict.device = device, " + "because device cannot be updated in-place. To update device, use " + "tensordict.to(new_device), which will return a new tensordict " + "on the new device. You can also use tensordict.to(new_device) as a " + "context manager for temporary device changes." + ) + + @property + def batch_size(self) -> torch.Size: + return self._batch_size + + @batch_size.setter + def batch_size(self, new_size: torch.Size) -> None: + self._batch_size_setter(new_size) + + def _change_batch_size(self, new_size: torch.Size) -> None: + self._batch_size = new_size + + # Checks + def _check_is_shared(self) -> bool: + share_list = [_is_shared(value) for value in self.values()] + if any(share_list) and not all(share_list): + shared_str = ", ".join( + [f"{key}: {_is_shared(value)}" for key, value in self.items()] + ) + raise RuntimeError( + f"tensors must be either all shared or not, but mixed " + f"features is not allowed. " + f"Found: {shared_str}" + ) + return all(share_list) and len(share_list) > 0 + + def _check_device(self, *, raise_exception: bool = True) -> None | bool: + val = True + for value in self.values(): + if _is_tensor_collection(type(value)): + val &= value._check_device(raise_exception=raise_exception) + if not val: + return False + val &= self.device is None or (self.device == value.device) + if not val: + if raise_exception: + raise RuntimeError( + f"devices are incongruent, got value with device {value.device}, " + f"-- expected {self.device}." + ) + return False + return val + + @lock_blocked + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + return self._tensordict.popitem() + + def _set_str( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ) -> Self: + if inplace is not False: + best_attempt = inplace is BEST_ATTEMPT_INPLACE + inplace = self._convert_inplace(inplace, key) + if not validated: + value = self._validate_value( + value, check_shape=True, non_blocking=non_blocking + ) + if not inplace: + if self._is_locked and not ignore_lock: + raise RuntimeError(_LOCK_ERROR) + self._tensordict[key] = value + else: + try: + dest = self._get_str(key, default=NO_DEFAULT) + if best_attempt and _is_tensor_collection(type(dest)): + dest.update( + value, + inplace=True, + non_blocking=non_blocking, + ignore_lock=ignore_lock, + ) + else: + if dest is not value: + try: + dest.copy_(value, non_blocking=non_blocking) + except RuntimeError: + # if we're updating a param and the storages match, nothing needs to be done + if not ( + isinstance(dest, torch.Tensor) + and dest.data.untyped_storage().data_ptr() + == value.data.untyped_storage().data_ptr() + ): + raise + except KeyError as err: + raise err + except Exception as err: + raise ValueError( + f"Failed to update '{key}' in tensordict {self}" + ) from err + return self + + def _set_dict( + self, + d: dict[str, CompatibleType], + *, + validated: bool, + ): + if not validated: + raise RuntimeError("Not Implemented for non-validated inputs") + self._tensordict = d + + def _set_tuple( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + non_blocking: bool = False, + ) -> Self: + if len(key) == 1: + return self._set_str( + key[0], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + td = self._get_str(key[0], None) + if td is None: + td = self._create_nested_str(key[0]) + inplace = False + elif not _is_tensor_collection(type(td)): + raise KeyError( + f"The entry {key[0]} is already present in tensordict {self}." + ) + td._set_tuple( + key[1:], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + return self + + _SHARED_INPLACE_ERROR = ( + "You're attempting to update a leaf in-place with a shared " + "tensordict, but the new value does not match the previous. " + "If you're using NonTensorData, see the class documentation " + "to see how to properly pre-allocate memory in shared contexts." + ) + + def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool): + if not validated: + value = self._validate_value( + value, check_shape=False, non_blocking=non_blocking + ) + validated = True + tensor_in = self._get_str(key, NO_DEFAULT) + + if is_non_tensor(value) and not (self._is_shared or self._is_memmap): + if isinstance(idx, tuple) and len(idx) == 1: + idx = idx[0] + dest = tensor_in + if ( + isinstance(idx, torch.Tensor) + and idx.shape == () + and self.shape == () + and idx.dtype == torch.bool + and idx + ): + self._set_str( + key, + dest.squeeze(0), + validated=True, + inplace=False, + ignore_lock=True, + ) + is_diff = dest[idx].tolist() != value.tolist() + if is_diff: + dest_val = dest.maybe_to_stack() + dest_val[idx] = value + if dest_val is not dest: + self._set_str( + key, + dest_val, + validated=True, + inplace=False, + ignore_lock=True, + ) + # TODO: ultimately, we want to get rid of the above logic + # dest_val = dest.maybe_to_stack() + # dest_val[idx] = value + # if dest_val is not dest: + # self._set_str( + # key, + # dest_val, + # validated=True, + # inplace=False, + # ignore_lock=True, + # ) + return + + if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple): + warn( + "Multiple indexing can lead to unexpected behaviours when " + "setting items, for instance `td[idx1][idx2] = other` may " + "not write to the desired location if idx1 is a list/tensor." + ) + tensor_in = _sub_index(tensor_in, idx) + tensor_in.copy_(value, non_blocking=non_blocking) + else: + tensor_out = _set_item( + tensor_in, idx, value, validated=validated, non_blocking=non_blocking + ) + if tensor_in is not tensor_out: + if self._is_shared or self._is_memmap: + raise RuntimeError(self._SHARED_INPLACE_ERROR) + # this happens only when a NonTensorData becomes a NonTensorStack + # so it is legitimate (there is no in-place modification of a tensor + # that was expected to happen but didn't). + # For this reason we can ignore the locked attribute of the td. + self._set_str( + key, + tensor_out, + validated=True, + inplace=False, + ignore_lock=True, + non_blocking=non_blocking, + ) + + return self + + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool): + if len(key) == 1: + return self._set_at_str( + key[0], value, idx, validated=validated, non_blocking=non_blocking + ) + if key[0] not in self.keys(): + # this won't work + raise KeyError(f"key {key} not found in set_at_ with tensordict {self}.") + else: + td = self._get_str(key[0], NO_DEFAULT) + td._set_at_tuple( + key[1:], value, idx, validated=validated, non_blocking=non_blocking + ) + return self + + @lock_blocked + def del_(self, key: NestedKey) -> Self: + key = _unravel_key_to_tuple(key) + if len(key) > 1: + td, subkey = _get_leaf_tensordict(self, key) + td.del_(subkey) + return self + + del self._tensordict[key[0]] + return self + + @lock_blocked + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> Self: + # these checks are not perfect, tuples that are not tuples of strings or empty + # tuples could go through but (1) it will raise an error anyway and (2) + # those checks are expensive when repeated often. + if not isinstance(old_key, (str, tuple)): + raise TypeError( + f"Expected old_name to be a string or a tuple of strings but found {type(old_key)}" + ) + if not isinstance(new_key, (str, tuple)): + raise TypeError( + f"Expected new_name to be a string or a tuple of strings but found {type(new_key)}" + ) + old_key = unravel_key(old_key) + new_key = unravel_key(new_key) + if old_key == new_key: + if old_key not in self.keys(include_nested=isinstance(old_key, tuple)): + raise KeyError(f"Key {old_key} not found in tensordict.") + return self + if safe and (new_key in self.keys(include_nested=True)): + raise KeyError(f"key {new_key} already present in TensorDict.") + + if isinstance(new_key, str): + self._set_str( + new_key, + self.get(old_key, default=NO_DEFAULT), + inplace=False, + validated=True, + non_blocking=False, + ) + else: + self._set_tuple( + new_key, + self.get(old_key, default=NO_DEFAULT), + inplace=False, + validated=True, + non_blocking=False, + ) + new_key_tuple = (new_key,) if isinstance(new_key, str) else new_key + if not ( + isinstance(old_key, tuple) + and old_key[: len(new_key_tuple)] == new_key_tuple + ): + # eg: td.rename_(("0", "1", "2"), ("0", "1")), then ("0", "1", "2") should not be deleted + self.del_(old_key) + return self + + def _stack_onto_(self, list_item: list[CompatibleType], dim: int) -> TensorDict: + # if not isinstance(key, str): + # raise ValueError("_stack_onto_ expects string keys.") + for key in self.keys(): + vals = [item._get_str(key, None) for item in list_item] + if all(v is None for v in vals): + continue + dest = self._get_str(key, NO_DEFAULT) + new_dest = torch.stack( + vals, + dim=dim, + out=dest, + ) + if new_dest is not dest: + # This can happen with non-tensor data + self._set_str(key, new_dest, inplace=False, validated=True) + return self + + def entry_class(self, key: NestedKey) -> type: + return type(self.get(key)) + + def _stack_onto_at_( + self, + list_item: list[CompatibleType], + dim: int, + idx: IndexType, + ) -> TensorDict: + if not isinstance(idx, tuple): + idx = (idx,) + idx = convert_ellipsis_to_idx(idx, self.batch_size) + for key in self.keys(): + vals = [td._get_str(key, NO_DEFAULT) for td in list_item] + if all(v is None for v in vals): + continue + v = self._get_str(key, NO_DEFAULT) + v_idx = v[idx] + if v.data_ptr() != v_idx.data_ptr(): + raise IndexError( + f"Index {idx} is incompatible with stack(..., out=data) as the storages of the indexed tensors differ." + ) + torch.stack(vals, dim=dim, out=v_idx) + # raise ValueError( + # f"Cannot stack onto an indexed tensor with index {idx} " + # f"as its storage differs." + # ) + return self + + def _get_str(self, key, default, **kwargs): + try: + return self._tensordict[key] + except KeyError: + return self._default_get(key, default) + + def _get_tuple(self, key, default, **kwargs): + first = self._get_str(key[0], default, **kwargs) + if len(key) == 1 or first is default: + return first + try: + return first._get_tuple(key[1:], default=default, **kwargs) + except AttributeError as err: + if "has no attribute" in str(err): + raise ValueError( + f"Expected a TensorDictBase instance but got {type(first)} instead" + f" for key '{key[1:]}' in tensordict:\n{self}." + ) + + def share_memory_(self) -> Self: + if self.is_memmap(): + raise RuntimeError( + "memmap and shared memory are mutually exclusive features." + ) + if self.device is not None and self.device.type == "cuda": + # cuda tensors are shared by default + return self + for value in self.values(): + if ( + isinstance(value, Tensor) + and value.device.type == "cpu" + or _is_tensor_collection(type(value)) + ): + value.share_memory_() + self._is_shared = True + self.lock_() + return self + + def detach_(self) -> Self: + for value in self.values(): + value.detach_() + return self + + def _memmap_( + self, + *, + prefix: str | None, + copy_existing: bool, + executor, + futures, + inplace, + like, + share_non_tensor, + existsok, + robust_key, + ) -> Self: + if prefix is not None: + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + metadata = {} + if inplace and self._is_shared: + raise RuntimeError( + "memmap and shared memory are mutually exclusive features." + ) + + dest = self if inplace else self.empty(device=torch.device("cpu")) + + # We must set these attributes before memmapping because we need the metadata + # to match the tensordict content. + if inplace: + self._is_memmap = True + self._is_shared = False # since they are mutually exclusive + self._device = torch.device("cpu") + else: + dest._is_memmap = True + dest._is_shared = False # since they are mutually exclusive + + for key, value in self.items(): + type_value = type(value) + if _is_tensor_collection(type_value): + dest._tensordict[key] = value._memmap_( + prefix=prefix / key if prefix is not None else None, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if prefix is not None: + _update_metadata( + metadata=metadata, key=key, value=value, is_collection=True + ) + continue + else: + if executor is None: + _populate_memmap( + dest=dest, + value=value, + key=key, + copy_existing=copy_existing, + prefix=prefix, + like=like, + existsok=existsok, + robust_key=robust_key, + ) + else: + futures.append( + executor.submit( + _populate_memmap, + dest=dest, + value=value, + key=key, + copy_existing=copy_existing, + prefix=prefix, + like=like, + existsok=existsok, + robust_key=robust_key, + ) + ) + if prefix is not None: + _update_metadata( + metadata=metadata, key=key, value=value, is_collection=False + ) + + if prefix is not None: + if executor is None: + _save_metadata( + dest, + prefix, + metadata=metadata, + ) + else: + futures.append(executor.submit(_save_metadata, dest, prefix, metadata)) + dest._is_locked = True + dest._memmap_prefix = prefix + return dest + + @classmethod + def _load_memmap( + cls, + prefix: str, + metadata: dict, + device: torch.device | None = None, + out=None, + *, + robust_key, + ) -> Self: + if metadata.get("device", "None") == "None": + metadata["device"] = None + else: + metadata["device"] = torch.device(metadata["device"]) + metadata["shape"] = torch.Size(metadata.get("shape", ())) + + if out is None: + result = cls( + {}, + batch_size=metadata.pop("shape"), + device=metadata.pop("device") if device is None else device, + ) + else: + result = out + + paths = set() + for key, entry_metadata in metadata.items(): + if not isinstance(entry_metadata, dict): + # there can be other metadata + continue + type_value = entry_metadata.get("type") + if type_value is not None: + paths.add(key) + continue + dtype = entry_metadata.get("dtype") + shape = entry_metadata.get("shape") + from .utils import ( + _encode_key_for_filesystem, + _get_robust_key_setting_with_warning, + ) + + # Use smart warning for loading that only warns when encoding would differ + effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key) + + # Use encoded filename for file system operations + safe_key = _encode_key_for_filesystem(key, robust=effective_robust_key) + memmap_file = prefix / f"{safe_key}.memmap" + + # If robust encoding is requested but file doesn't exist, try legacy filename + if not memmap_file.exists() and effective_robust_key: + legacy_key = _encode_key_for_filesystem(key, robust=False) + legacy_file = prefix / f"{legacy_key}.memmap" + if legacy_file.exists(): + # Use legacy filename for backward compatibility + safe_key = legacy_key + memmap_file = legacy_file + + if not memmap_file.exists() or dtype is None or shape is None: + # invalid dict means + continue + try: + # this was absent in earlier versions of pytorch + is_fake = torch._guards.active_fake_mode() + except AttributeError: + # Let's just make sure that the private function is just not gone + if torch.__version__ > "2.3.0": + raise + is_fake = False + if (device is None or device != torch.device("meta")) and not is_fake: + if entry_metadata.get("is_nested", False): + # The shape is the shape of the shape, get the shape from it + shape = MemoryMappedTensor.from_filename( + (prefix / f"{safe_key}.memmap").with_suffix(".shape.memmap"), + shape=shape, + dtype=torch.long, + ) + else: + shape = torch.Size(shape) + tensor = MemoryMappedTensor.from_filename( + dtype=_STR_DTYPE_TO_DTYPE[dtype], + shape=shape, + filename=str(prefix / f"{safe_key}.memmap"), + ) + if device is not None: + tensor = tensor.to(device, non_blocking=True) + else: + tensor = torch.zeros( + torch.Size(shape), + device=device, + dtype=_STR_DTYPE_TO_DTYPE[dtype], + ) + result._set_str( + key, + tensor, + validated=True, + inplace=False, + non_blocking=False, + ) + # iterate over folders and load them + for path in prefix.iterdir(): + if path.is_dir() and path.parts[-1] in paths: + key = path.parts[-1] # path.parts[len(prefix.parts) :] + existing_elt = result._get_str(key, default=None) + if existing_elt is not None: + existing_elt.load_memmap_(path) + else: + result._set_str( + key, + TensorDict.load_memmap(path, device=device, non_blocking=True), + inplace=False, + validated=False, + ) + result._memmap_prefix = prefix + return result + + def _make_memmap_subtd(self, key): + """Creates a sub-tensordict given a tuple key.""" + result = self + for key_str in key: + result_tmp = result._get_str(key_str, default=None) + if result_tmp is None: + result_tmp = result.empty() + if result._memmap_prefix is not None: + result_tmp.memmap_(prefix=result._memmap_prefix / key_str) + metadata = _load_metadata(result._memmap_prefix) + _update_metadata( + metadata=metadata, + key=key_str, + value=result_tmp, + is_collection=True, + ) + _save_metadata( + result, prefix=result._memmap_prefix, metadata=metadata + ) + result._tensordict[key_str] = result_tmp + result = result_tmp + return result + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + if not self.is_memmap(): + raise RuntimeError( + "Can only make a memmap tensor within a memory-mapped tensordict." + ) + + key = unravel_key(key) + if isinstance(key, tuple): + last_node = self._make_memmap_subtd(key[:-1]) + last_key = key[-1] + else: + last_node = self + last_key = key + if last_key in last_node.keys(): + raise RuntimeError( + f"The key {last_key} already exists within the target tensordict. Delete that entry before " + f"overwriting it." + ) + if dtype is None: + dtype = torch.get_default_dtype() + if last_node._memmap_prefix is not None: + metadata = _load_metadata(last_node._memmap_prefix) + memmap_tensor = _populate_empty( + key=last_key, + dest=last_node, + prefix=last_node._memmap_prefix, + shape=shape, + dtype=dtype, + robust_key=robust_key, + ) + _update_metadata( + metadata=metadata, + key=last_key, + value=memmap_tensor, + is_collection=False, + ) + _save_metadata( + last_node, prefix=last_node._memmap_prefix, metadata=metadata + ) + else: + memmap_tensor = MemoryMappedTensor.empty(shape=shape, dtype=dtype) + + last_node._set_str( + last_key, memmap_tensor, validated=False, inplace=False, ignore_lock=True + ) + + return memmap_tensor + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + if not self.is_memmap(): + raise RuntimeError( + "Can only make a memmap tensor within a memory-mapped tensordict." + ) + + key = unravel_key(key) + if isinstance(key, tuple): + last_node = self._make_memmap_subtd(key[:-1]) + last_key = key[-1] + else: + last_node = self + last_key = key + if last_key in last_node.keys(): + raise RuntimeError( + f"The key {last_key} already exists within the target tensordict. Delete that entry before " + f"overwriting it." + ) + if dtype is None: + dtype = torch.get_default_dtype() + + if last_node._memmap_prefix is not None: + metadata = _load_metadata(last_node._memmap_prefix) + memmap_tensor = _populate_storage( + key=last_key, + dest=last_node, + prefix=last_node._memmap_prefix, + storage=storage, + shape=shape, + dtype=dtype, + robust_key=robust_key, + ) + _update_metadata( + metadata=metadata, + key=last_key, + value=memmap_tensor, + is_collection=False, + ) + _save_metadata( + last_node, prefix=last_node._memmap_prefix, metadata=metadata + ) + else: + memmap_tensor = MemoryMappedTensor.from_storage( + storage=storage, shape=shape, dtype=dtype + ) + + last_node._set_str( + last_key, memmap_tensor, validated=False, inplace=False, ignore_lock=True + ) + + return memmap_tensor + + def make_memmap_from_tensor( + self, + key: NestedKey, + tensor: torch.Tensor, + *, + copy_data: bool = True, + existsok: bool = True, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + if not self.is_memmap(): + raise RuntimeError( + "Can only make a memmap tensor within a memory-mapped tensordict." + ) + + key = unravel_key(key) + if isinstance(key, tuple): + last_node = self._make_memmap_subtd(key[:-1]) + last_key = key[-1] + else: + last_node = self + last_key = key + if last_key in last_node.keys(): + raise RuntimeError( + f"The key {last_key} already exists within the target tensordict. Delete that entry before " + f"overwriting it." + ) + + if last_node._memmap_prefix is not None: + metadata = _load_metadata(last_node._memmap_prefix) + memmap_tensor = _populate_memmap( + dest=last_node, + value=tensor, + key=last_key, + copy_existing=True, + prefix=last_node._memmap_prefix, + like=not copy_data, + existsok=existsok, + robust_key=robust_key, + ) + _update_metadata( + metadata=metadata, + key=last_key, + value=memmap_tensor, + is_collection=False, + ) + _save_metadata( + last_node, prefix=last_node._memmap_prefix, metadata=metadata + ) + else: + memmap_tensor = MemoryMappedTensor.from_tensor(tensor) + + last_node._set_str( + last_key, memmap_tensor, validated=False, inplace=False, ignore_lock=True + ) + + return memmap_tensor + + def where( + self, + condition: Tensor, + other: Tensor | TensorCollection, + *, + out: TensorCollection | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ) -> Self: + if _is_tensor_collection(type(other)): + + def func(tensor, _other, key): + if tensor is None: + if pad is not None: + tensor = _other + _other = torch.tensor( + pad, dtype=_other.dtype, device=_other.device + ) + else: + raise KeyError( + f"Key {key} not found and no pad value provided." + ) + cond = expand_as_right(~condition, tensor) + elif _other is None: + if pad is not None: + _other = torch.tensor( + pad, dtype=tensor.dtype, device=tensor.device + ) + else: + raise KeyError( + f"Key {key} not found and no pad value provided." + ) + cond = expand_as_right(condition, tensor) + else: + cond = expand_as_right(condition, tensor) + return torch.where( + condition=cond, + input=tensor, + other=_other, + ) + + result = self.empty() if out is None else out + other_keys = set(other.keys()) + # we turn into a list because out could be = to self! + for key in list(self.keys()): + tensor = self._get_str(key, default=NO_DEFAULT) + _other = other._get_str(key, default=None) + if _is_tensor_collection(type(tensor)): + _out = None if out is None else out._get_str(key, None) + if _other is None: + _other = tensor.empty() + val = tensor.where( + condition=condition, other=_other, out=_out, pad=pad + ) + else: + val = func(tensor, _other, key) + result._set_str( + key, val, inplace=False, validated=True, non_blocking=False + ) + other_keys.discard(key) + for key in other_keys: + tensor = None + _other = other._get_str(key, default=NO_DEFAULT) + if _is_tensor_collection(type(_other)): + try: + tensor = _other.empty() + except NotImplementedError: + # H5 tensordicts do not support select() + tensor = _other.to_tensordict().empty() + val = _other.where( + condition=~condition, other=tensor, out=None, pad=pad + ) + else: + val = func(tensor, _other, key) + result._set_str( + key, val, inplace=False, validated=True, non_blocking=False + ) + return result + else: + if out is None: + + def func(tensor): + return torch.where( + condition=expand_as_right(condition, tensor), + input=tensor, + other=other, + ) + + return self._fast_apply(func, propagate_lock=True) + else: + + def func(tensor, _out): + return torch.where( + condition=expand_as_right(condition, tensor), + input=tensor, + other=other, + out=_out, + update_batch_size=update_batch_size, + ) + + return self._fast_apply(func, out, propagate_lock=True) + + def masked_fill_(self, mask: Tensor, value: float | int | bool) -> Self: + for item in self.values(): + mask_expand = expand_as_right(mask, item) + item.masked_fill_(mask_expand, value) + return self + + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: + td_copy = self.clone() + return td_copy.masked_fill_(mask, value) + + def is_contiguous(self) -> bool: + return all([value.is_contiguous() for _, value in self.items()]) + + def _clone(self, recurse: bool = True) -> Self: + if recurse and self.device is not None: + return self._clone_recurse() + + result = self._new_unsafe( + source={key: _clone_value(value, recurse) for key, value in self.items()}, + batch_size=self.batch_size, + device=self.device, + names=self._maybe_names(), + ) + # If this is uncommented, a shallow copy of a shared/memmap will be shared and locked too + # This may be undesirable, not sure if this should be the default behaviour + # (one usually does a copy to modify it). + # if not recurse: + # self._maybe_set_shared_attributes(result) + return result + + def contiguous(self) -> Self: + source = {key: value.contiguous() for key, value in self.items()} + batch_size = self.batch_size + device = self.device + out = self._new_unsafe( + source=source, + batch_size=batch_size, + device=device, + names=self.names if self._has_names() else None, + ) + return out + + def empty( + self, recurse=False, *, batch_size=None, device=NO_DEFAULT, names=NO_DEFAULT + ) -> Self: + if not recurse: + return self._new_unsafe( + device=self._device if device is NO_DEFAULT else device, + batch_size=( + self._batch_size if batch_size is None else torch.Size(batch_size) + ), + source={}, + names=( + (self.names if self._has_names() else None) + if names is NO_DEFAULT + else names + ), + ) + return super().empty(recurse=recurse) + + def _select( + self, + *keys: NestedKey, + inplace: bool = False, + strict: bool = True, + set_shared: bool = True, + ) -> Self: + if inplace and self.is_locked: + raise RuntimeError(_LOCK_ERROR) + + source = {} + if len(keys): + keys_to_select = None + for key in keys: + if isinstance(key, str): + subkey = [] + else: + key, subkey = key[0], key[1:] + + val = self._get_str(key, default=_UNSET if not strict else NO_DEFAULT) + if val is _UNSET: + continue + source[key] = val + if len(subkey): + if keys_to_select is None: + # delay creation of defaultdict + keys_to_select = defaultdict(list) + keys_to_select[key].append(subkey) + + if keys_to_select is not None: + for key, val in keys_to_select.items(): + source[key] = source[key]._select( + *val, strict=strict, inplace=inplace, set_shared=set_shared + ) + + names = self._td_dim_names + if names is not None: + names = list(names) + + result = self._new_unsafe( + device=self.device, + batch_size=self.batch_size, + source=source, + # names=self.names if self._has_names() else None, + names=names, + ) + if inplace: + self._tensordict = result._tensordict + return self + # If this is uncommented, a shallow copy of a shared/memmap will be shared and locked too + # This may be undesirable, not sure if this should be the default behaviour + # (one usually does a copy to modify it). + # if set_shared: + # self._maybe_set_shared_attributes(result) + return result + + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> Self: + # faster than Base.exclude + if not len(keys): + return self.copy() if not inplace else self + if not inplace: + _tensordict = dict(self._tensordict) + else: + _tensordict = self._tensordict + keys_to_exclude = None + for key in keys: + key = unravel_key(key) + if isinstance(key, str): + _tensordict.pop(key, None) + else: + if keys_to_exclude is None: + # delay creation of defaultdict + keys_to_exclude = defaultdict(list) + if key[0] in self._tensordict: + keys_to_exclude[key[0]].append(key[1:]) + if keys_to_exclude is not None: + for key, cur_keys in keys_to_exclude.items(): + val = _tensordict.get(key) + if val is not None: + val = val._exclude( + *cur_keys, inplace=inplace, set_shared=set_shared + ) + if not inplace: + _tensordict[key] = val + if inplace: + return self + result = self._new_unsafe( + _tensordict, + batch_size=self.batch_size, + device=self.device, + names=self.names if self._has_names() else None, + ) + # If this is uncommented, a shallow copy of a shared/memmap will be shared and locked too + # This may be undesirable, not sure if this should be the default behaviour + # (one usually does a copy to modify it). + # if set_shared: + # self._maybe_set_shared_attributes(result) + return result + + # @cache + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _TensorDictKeysView: + if not include_nested and not leaves_only and is_leaf is None: + if not sort: + return _StringKeys(self._tensordict.keys()) + else: + + def keyfunc(x): + return ".".join(x) if isinstance(x, tuple) else x + + return sorted( + _StringKeys(self._tensordict.keys()), + key=keyfunc, + ) + else: + return self._nested_keys( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + @cache # noqa: B019 + def _nested_keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _TensorDictKeysView: + return _TensorDictKeysView( + self, + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + # some custom methods for efficiency + def items( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> Iterator[tuple[str, CompatibleType]]: + if not include_nested and not leaves_only: + if not sort: + return self._tensordict.items() + + def keyfunc(x): + return x[0] + + return sorted(self._tensordict.items(), key=keyfunc) + elif include_nested and leaves_only and not sort: + is_leaf = _default_is_leaf if is_leaf is None else is_leaf + result = [] + + def fast_iter(): + for key, val in self._tensordict.items(): + # We could easily make this faster, here we're iterating twice over the keys, + # but we could iterate just once. + # Ideally we should make a "dirty" list of items then call unravel_key on all of them. + if not is_leaf(type(val)): + for _key, _val in val.items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + ): + if isinstance(_key, str): + _key = (key, _key) + else: + _key = (key, *_key) + result.append((_key, _val)) + else: + result.append((key, val)) + return result + + return fast_iter() + else: + return super().items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + def values( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> Iterator[tuple[str, CompatibleType]]: + if not include_nested and not leaves_only: + if not sort: + return self._tensordict.values() + else: + + def keyfunc(x): + return x[0] + + return list(zip(*sorted(self._tensordict.items(), key=keyfunc)))[1] + else: + return TensorDictBase.values( + self, + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + +class _SubTensorDict(TensorDictBase): + """A TensorDict that only sees an index of the stored tensors.""" + + _lazy = True + _inplace_set = True + _safe = False + + def __init__( + self, + source: T, + idx: IndexType, + batch_size: Sequence[int] | None = None, + ) -> None: + if not _is_tensor_collection(type(source)): + raise TypeError( + f"Expected source to be a subclass of TensorDictBase, " + f"got {type(source)}" + ) + self._source = source + idx = ( + (idx,) + if not isinstance( + idx, + ( + tuple, + list, + ), + ) + else tuple(idx) + ) + if any(item is Ellipsis for item in idx): + idx = convert_ellipsis_to_idx(idx, self._source.batch_size) + self._batch_size = _getitem_batch_size(self._source.batch_size, idx) + self.idx = idx + + if batch_size is not None and batch_size != self.batch_size: + raise RuntimeError("batch_size does not match self.batch_size.") + + # These attributes should never be set + @property + def _is_shared(self): + return self._source._is_shared + + @property + def _is_memmap(self): + return self._source._is_memmap + + @staticmethod + def _convert_ellipsis(idx, shape): + if any(_idx is Ellipsis for _idx in idx): + new_idx = [] + cursor = -1 + for _idx in idx: + if _idx is Ellipsis: + if cursor == len(idx) - 1: + # then we can just skip + continue + n_upcoming = len(idx) - cursor - 1 + while cursor < len(shape) - n_upcoming: + cursor += 1 + new_idx.append(slice(None)) + else: + new_idx.append(_idx) + return tuple(new_idx) + return idx + + @property + def batch_size(self) -> torch.Size: + return self._batch_size + + @batch_size.setter + def batch_size(self, new_size: torch.Size) -> None: + self._batch_size_setter(new_size) + + @property + def names(self): + names = self._source._get_names_idx(self.idx) + if names is None: + return [None] * self.batch_dims + return names + + @names.setter + def names(self, value): + self._set_names(value) + + def _set_names(self, names: Sequence[str] | None): + if names is None: + names = [None] * self.batch_dims + if names[: self.batch_dims] == self.names: + return + raise RuntimeError( + "Names of a subtensordict cannot be modified. Instantiate it as a TensorDict first." + ) + + def _has_names(self): + return self._source._has_names() + + def _erase_names(self): + raise RuntimeError( + "Cannot erase names of a _SubTensorDict. Erase source TensorDict's names instead." + ) + + def _rename_subtds(self, names): + for key in self.keys(): + if _is_tensor_collection(self.entry_class(key)): + raise RuntimeError("Cannot rename nested sub-tensordict dimensions.") + + @property + def device(self) -> None | torch.device: + return self._source.device + + @device.setter + def device(self, value: DeviceType) -> None: + self._source.device = value + + def _preallocate(self, key: NestedKey, value: CompatibleType) -> Self: + return self._source.set(key, value) + + def _convert_inplace(self, inplace, key): + has_key = key in self.keys() + if inplace is not False: + if inplace is True and not has_key: # inplace could be None + raise KeyError( + _KEY_ERROR.format(key, type(self).__name__, sorted(self.keys())) + ) + inplace = has_key + if not inplace and has_key: + raise RuntimeError( + "Calling `_SubTensorDict.set(key, value, inplace=False)` is " + "prohibited for existing tensors. Consider calling " + "_SubTensorDict.set_(...) or cloning your tensordict first." + ) + elif not inplace and self.is_locked: + raise RuntimeError(_LOCK_ERROR) + return inplace + + from_dict_instance = TensorDict.from_dict_instance + + def _set_str( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ) -> Self: + inplace = self._convert_inplace(inplace, key) + # it is assumed that if inplace=False then the key doesn't exist. This is + # checked in set method, but not here. responsibility lies with the caller + # so that this method can have minimal overhead from runtime checks + parent = self._source + if not validated: + value = self._validate_value( + value, check_shape=True, non_blocking=non_blocking + ) + validated = True + if not inplace: + if _is_tensor_collection(type(value)): + # value has the shape of subtd[idx], so we want an expanded + # version value_expand such that value_expand[idx] has the + # shape of value + value_expand = _expand_to_match_shape( + parent.batch_size, + value, + self.batch_dims, + self.device, + index=self.idx, + ) + for _key, _tensor in value.items(): + value_expand._set_str( + _key, + _expand_to_match_shape( + parent.batch_size, + _tensor, + self.batch_dims, + self.device, + index=self.idx, + ), + inplace=inplace, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + else: + value_expand = torch.zeros( + ( + *parent.batch_size, + *_shape(value)[self.batch_dims :], + ), + dtype=value.dtype, + device=self.device, + ) + if self._is_shared: + value_expand.share_memory_() + elif self._is_memmap: + value_expand = MemoryMappedTensor.from_tensor(value_expand) + parent._set_str( + key, + value_expand, + inplace=False, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + + parent._set_at_str( + key, value, self.idx, validated=validated, non_blocking=non_blocking + ) + return self + + def _set_tuple( + self, + key: NestedKey, + value: dict[str, CompatibleType] | CompatibleType, + *, + inplace: bool, + validated: bool, + non_blocking: bool = False, + ) -> Self: + if len(key) == 1: + return self._set_str( + key[0], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + parent = self._source + td = parent._get_str(key[0], None) + if td is None: + td = parent.select() + parent._set_str( + key[0], td, inplace=False, validated=True, non_blocking=non_blocking + ) + _SubTensorDict(td, self.idx)._set_tuple( + key[1:], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + return self + + def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool): + tensor_in = self._get_str(key, NO_DEFAULT) + if not validated: + value = self._validate_value( + value, check_shape=False, non_blocking=non_blocking + ) + validated = True + if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple): + warn( + "Multiple indexing can lead to unexpected behaviours when " + "setting items, for instance `td[idx1][idx2] = other` may " + "not write to the desired location if idx1 is a list/tensor." + ) + tensor_in = _sub_index(tensor_in, idx) + tensor_in.copy_(value) + tensor_out = tensor_in + else: + tensor_out = _set_item( + tensor_in, idx, value, validated=validated, non_blocking=non_blocking + ) + # make sure that the value is updated + self._source._set_at_str( + key, tensor_out, self.idx, validated=validated, non_blocking=non_blocking + ) + return self + + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool): + if len(key) == 1: + return self._set_at_str( + key[0], value, idx, validated=validated, non_blocking=non_blocking + ) + if key[0] not in self.keys(): + # this won't work + raise KeyError(f"key {key} not found in set_at_ with tensordict {self}.") + else: + td = self._get_str(key[0], NO_DEFAULT) + td._set_at_tuple( + key[1:], value, idx, validated=validated, non_blocking=non_blocking + ) + return self + + # @cache # noqa: B019 + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _TensorDictKeysView: + return self._source.keys( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + def entry_class(self, key: NestedKey) -> type: + source_type = type(self._source.get(key)) + if _is_tensor_collection(source_type): + return type(self) + return source_type + + def _stack_onto_(self, list_item: list[CompatibleType], dim: int) -> _SubTensorDict: + self._source._stack_onto_at_(list_item, dim=dim, idx=self.idx) + return self + + @_as_context_manager() + def to(self, *args, **kwargs: Any) -> Self: + ( + device, + dtype, + non_blocking, + convert_to_format, + batch_size, + pin_memory, + num_threads, + inplace, + ) = _parse_to(*args, **kwargs) + result = self + if inplace: + raise TypeError( + "Cannot send a _SubTensorDict instance to device/dtype inplace." + ) + if device is not None and dtype is None and device == self.device: + return result + return self.to_tensordict().to(*args, **kwargs) + + def _change_batch_size(self, new_size: torch.Size) -> None: + self._batch_size = new_size + + def _get_non_tensor(self, key: NestedKey, default=NO_DEFAULT): + out = super()._get_non_tensor(key, default=default) + + if isinstance(out, _SubTensorDict) and is_non_tensor(out._source): + return out._source + return out + + def _get_str(self, key, default, **kwargs): + if key in self.keys() and _is_tensor_collection(self.entry_class(key)): + data = self._source._get_str(key, NO_DEFAULT, **kwargs) + if _pass_through(data): + return data[self.idx] + return _SubTensorDict(data, self.idx) + return self._source._get_at_str(key, self.idx, default=default, **kwargs) + + def _get_tuple(self, key, default, **kwargs): + return self._source._get_at_tuple(key, self.idx, default=default, **kwargs) + + @lock_blocked + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorCollection, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + **kwargs, + ) -> _SubTensorDict: + if input_dict_or_td is self: + # no op + return self + if is_leaf is None: + is_leaf = _is_leaf_nontensor + + if getattr(self._source, "_has_exclusive_keys", False): + raise RuntimeError( + "Cannot use _SubTensorDict.update with a LazyStackedTensorDict that has exclusive keys." + ) + if keys_to_update is not None: + if len(keys_to_update) == 0: + return self + keys_to_update = unravel_key_list(keys_to_update) + keys = set(self.keys(False)) + for key, value in input_dict_or_td.items(): + key = _unravel_key_to_tuple(key) + firstkey, subkey = key[0], key[1:] + if keys_to_update and not any( + firstkey == ktu if isinstance(ktu, str) else firstkey == ktu[0] + for ktu in keys_to_update + ): + continue + if clone and hasattr(value, "clone"): + value = value.clone() + elif clone: + value = tree_map(torch.clone, value) + # the key must be a string by now. Let's check if it is present + if firstkey in keys: + target_class = self.entry_class(firstkey) + if _is_tensor_collection(target_class): + target = self._source.get(firstkey)._get_sub_tensordict(self.idx) + if len(subkey): + sub_keys_to_update = _prune_selected_keys( + keys_to_update, firstkey + ) + target.update( + {subkey: value}, + inplace=False, + keys_to_update=sub_keys_to_update, + non_blocking=non_blocking, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + continue + elif isinstance(value, dict) or _is_tensor_collection(type(value)): + sub_keys_to_update = _prune_selected_keys( + keys_to_update, firstkey + ) + target.update( + value, + keys_to_update=sub_keys_to_update, + non_blocking=non_blocking, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + continue + raise ValueError( + f"Tried to replace a tensordict with an incompatible object of type {type(value)}" + ) + else: + self._set_tuple( + key, + value, + inplace=True, + validated=False, + non_blocking=non_blocking, + ) + else: + self._set_tuple( + key, + value, + inplace=BEST_ATTEMPT_INPLACE if inplace else False, + validated=False, + non_blocking=non_blocking, + ) + return self + + def update_( + self, + input_dict: dict[str, CompatibleType] | TensorCollection, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> _SubTensorDict: + return self.update_at_( + input_dict, + idx=self.idx, + discard_idx_attr=True, + clone=clone, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + ) + + def update_at_( + self, + input_dict: dict[str, CompatibleType] | TensorCollection, + idx: IndexType, + *, + discard_idx_attr: bool = False, + clone: bool = False, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> _SubTensorDict: + if keys_to_update is not None: + if len(keys_to_update) == 0: + return self + keys_to_update = unravel_key_list(keys_to_update) + for key, value in input_dict.items(): + key = _unravel_key_to_tuple(key) + firstkey, _ = key[0], key[1:] + if keys_to_update and not any( + firstkey == ktu if isinstance(ktu, str) else firstkey == ktu[0] + for ktu in keys_to_update + ): + continue + if not isinstance(value, tuple(_ACCEPTED_CLASSES)): + raise TypeError( + f"Expected value to be one of types {_ACCEPTED_CLASSES} " + f"but got {type(value)}" + ) + if clone: + value = value.clone() + if discard_idx_attr: + self._source._set_at_tuple( + key, + value, + idx, + non_blocking=non_blocking, + validated=False, + ) + else: + self._set_at_tuple( + key, value, idx, validated=False, non_blocking=non_blocking + ) + return self + + def get_parent_tensordict(self) -> Self: + if not isinstance(self._source, TensorDictBase): + raise TypeError( + f"_SubTensorDict was initialized with a source of type" + f" {type(self._source).__name__}, " + "parent tensordict not accessible" + ) + return self._source + + @lock_blocked + def del_(self, key: NestedKey) -> Self: + self._source = self._source.del_(key) + return self + + @lock_blocked + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + raise NotImplementedError( + f"popitem not implemented for class {type(self).__name__}." + ) + + def _clone(self, recurse: bool = True) -> _SubTensorDict: + """Clones the _SubTensorDict. + + Args: + recurse (bool, optional): if ``True`` (default), a regular + :class:`~.tensordict.TensorDict` instance will be created from the :class:`~.tensordict._SubTensorDict`. + Otherwise, another :class:`~.tensordict._SubTensorDict` with identical content + will be returned. + + Examples: + >>> data = TensorDict({"a": torch.arange(4).reshape(2, 2,)}, batch_size=[2, 2]) + >>> sub_data = data._get_sub_tensordict([0,]) + >>> print(sub_data) + _SubTensorDict( + fields={ + a: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + >>> # the data of both subtensordict is the same + >>> print(data.get("a").data_ptr(), sub_data.get("a").data_ptr()) + 140183705558208 140183705558208 + >>> sub_data_clone = sub_data.clone(recurse=True) + >>> print(sub_data_clone) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + >>. print(sub_data.get("a").data_ptr()) + 140183705558208 + >>> sub_data_clone = sub_data.clone(recurse=False) + >>> print(sub_data_clone) + _SubTensorDict( + fields={ + a: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + >>> print(sub_data.get("a").data_ptr()) + 140183705558208 + """ + if not recurse: + return _SubTensorDict( + source=self._source._clone(recurse=False), idx=self.idx + ) + return self.to_tensordict() + + def is_contiguous(self) -> bool: + return all(value.is_contiguous() for value in self.values()) + + def contiguous(self) -> Self: + return self._new_unsafe( + batch_size=self.batch_size, + source={key: value.contiguous() for key, value in self.items()}, + device=self.device, + names=self.names if self._has_names() else None, + ) + + def _select( + self, + *keys: NestedKey, + inplace: bool = False, + strict: bool = True, + set_shared: bool = True, + ) -> Self: + if inplace: + raise RuntimeError("Cannot call select inplace on a lazy tensordict.") + return self.to_tensordict()._select( + *keys, inplace=False, strict=strict, set_shared=set_shared + ) + + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> Self: + if inplace: + raise RuntimeError("Cannot call exclude inplace on a lazy tensordict.") + return self.to_tensordict()._exclude( + *keys, inplace=False, set_shared=set_shared + ) + + def expand(self, *args: int, inplace: bool = False) -> Self: + if len(args) == 1 and isinstance(args[0], Sequence): + shape = tuple(args[0]) + else: + shape = args + + def expand(x): + return x.expand((*shape, *x.shape[self.ndim :])) + + return self._fast_apply( + expand, + batch_size=shape, + propagate_lock=True, + ) + + @classmethod + def from_dict( + cls, + input_dict, + *others, + auto_batch_size: bool = False, + batch_size=None, + device=None, + batch_dims=None, + names=None, + ): + raise NotImplementedError(f"from_dict not implemented for {cls.__name__}.") + + def is_shared(self) -> bool: + return self._source.is_shared() + + def is_memmap(self) -> bool: + return self._source.is_memmap() + + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> _SubTensorDict: + self._source.rename_key_(old_key, new_key, safe=safe) + return self + + def pin_memory(self, *args, **kwargs) -> Self: + raise RuntimeError( + f"Cannot pin memory of a {type(self).__name__}. Call to_tensordict() before making this call." + ) + + def detach_(self) -> Self: + raise RuntimeError("Detaching a sub-tensordict in-place cannot be done.") + + def where( + self, + condition: Tensor, + other: Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ): + return self.to_tensordict().where( + condition=condition, + other=other, + out=out, + pad=pad, + update_batch_size=update_batch_size, + ) + + def masked_fill_(self, mask: Tensor, value: float | bool) -> Self: + for key, item in self.items(): + self.set_(key, torch.full_like(item, value)) + return self + + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: + td_copy = self.clone() + return td_copy.masked_fill_(mask, value) + + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + num_threads: int = 0, + ) -> Self: + raise RuntimeError( + "Converting a sub-tensordict values to memmap cannot be done." + ) + + def _memmap_( + self, + *, + prefix: str | None, + copy_existing: bool, + executor, + futures, + inplace, + like, + share_non_tensor, + existsok, + robust_key, + ) -> Self: + if prefix is not None: + + def save_metadata(prefix=prefix, self=self): + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + with open(prefix / "meta.json", "wb") as f: + from tensordict.utils import json_dumps + + metadata_json = json_dumps( + { + "_type": str(type(self)), + "index": _index_to_str(self.idx), + } + ) + if isinstance(metadata_json, str): + metadata_json = metadata_json.encode("utf-8") + f.write(metadata_json) + + if executor is None: + save_metadata() + else: + futures.append(executor.submit(save_metadata)) + + _source = self._source._memmap_( + prefix=prefix / "_source" if prefix is not None else None, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if not inplace: + result = _SubTensorDict(_source, idx=self.idx) + else: + result = self + return result + + @classmethod + def _load_memmap( + cls, + prefix: Path, + metadata: dict, + device: torch.device | None = None, + *, + robust_key, + ): + index = metadata["index"] + return _SubTensorDict( + TensorDict.load_memmap( + prefix / "_source", device=device, robust_key=robust_key + ), + _str_to_index(index), + ) + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for _SubTensorDict." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for _SubTensorDict." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_tensor( + self, key: NestedKey, tensor: torch.Tensor, *, copy_data: bool = True + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for _SubTensorDict." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def share_memory_(self) -> Self: + raise RuntimeError( + "Casting a sub-tensordict values to shared memory cannot be done." + ) + + @property + def is_locked(self) -> bool: + return self._source.is_locked + + @is_locked.setter + def is_locked(self, value) -> bool: + if value: + self.lock_() + else: + self.unlock_() + + @_as_context_manager("is_locked") + def lock_(self) -> Self: + # we can't lock sub-tensordicts because that would mean that the + # parent tensordict cannot be modified either. + if not self.is_locked: + raise RuntimeError( + "Cannot lock a _SubTensorDict. Lock the parent tensordict instead." + ) + return self + + @_as_context_manager("is_locked") + def unlock_(self) -> Self: + if self.is_locked: + raise RuntimeError( + "Cannot unlock a _SubTensorDict. Unlock the parent tensordict instead." + ) + return self + + def _remove_lock(self, lock_id): + raise RuntimeError( + "Cannot unlock a _SubTensorDict. Unlock the parent tensordict instead." + ) + + def _propagate_lock(self, lock_ids=None, *, is_compiling): + raise RuntimeError( + "Cannot lock a _SubTensorDict. Lock the parent tensordict instead." + ) + + def __del__(self): + pass + + def _create_nested_str(self, key): + # this may fail with a sub-sub tensordict + out = self._source.empty() + self._source._set_str( + key, out, inplace=False, validated=True, non_blocking=False + ) + # the id of out changes + return self._get_str(key, default=NO_DEFAULT) + + def _cast_reduction( + self, + *, + reduction_name, + dim=NO_DEFAULT, + keepdim=NO_DEFAULT, + tuple_ok=True, + **kwargs, + ): + try: + td = self.to_tensordict() + except Exception: + raise RuntimeError( + f"{reduction_name} requires this object to be cast to a regular TensorDict. " + f"If you need {type(self).__name__} to support {reduction_name}, help us by filing an issue" + f" on github!" + ) + return td._cast_reduction( + reduction_name=reduction_name, + dim=dim, + keepdim=keepdim, + tuple_ok=tuple_ok, + **kwargs, + ) + + # TODO: check these implementations + __eq__ = TensorDict.__eq__ + __ne__ = TensorDict.__ne__ + __ge__ = TensorDict.__ge__ + __gt__ = TensorDict.__gt__ + __le__ = TensorDict.__le__ + __lt__ = TensorDict.__lt__ + __setitem__ = TensorDict.__setitem__ + __xor__ = TensorDict.__xor__ + __or__ = TensorDict.__or__ + _check_device = TensorDict._check_device + _check_is_shared = TensorDict._check_is_shared + _to_module = TensorDict._to_module + _unbind = TensorDict._unbind + all = TensorDict.all + any = TensorDict.any + masked_select = TensorDict.masked_select + memmap_like = TensorDict.memmap_like + repeat_interleave = TensorDict.repeat_interleave + _repeat = TensorDict._repeat + reshape = TensorDict.reshape + split = TensorDict.split + + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorCollection, ...]: + splits = -(self.batch_size[dim] // -chunks) + return self.split(splits, dim) + + def _view(self, *args, **kwargs): + raise RuntimeError( + "Cannot call `view` on a sub-tensordict. Call `reshape` instead." + ) + + def _transpose(self, dim0, dim1): + raise RuntimeError( + "Cannot call `transpose` on a sub-tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _permute( + self, + *args, + **kwargs, + ): + raise RuntimeError( + "Cannot call `permute` on a sub-tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _squeeze(self, dim=None): + raise RuntimeError( + "Cannot call `squeeze` on a sub-tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _unsqueeze(self, dim: int): + raise RuntimeError( + "Cannot call `unsqueeze` on a sub-tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + _add_batch_dim = TensorDict._add_batch_dim + + _apply_nest = TensorDict._apply_nest + _multithread_apply_flat = TensorDict._multithread_apply_flat + _multithread_rebuild = TensorDict._multithread_rebuild + _convert_to_tensordict = TensorDict._convert_to_tensordict + + _get_names_idx = TensorDict._get_names_idx + + def _index_tensordict( + self, + index: IndexType, + new_batch_size: torch.Size | None = None, + names: Sequence[str] | None = None, + ) -> _SubTensorDict: + # we ignore the names and new_batch_size which are only provided for + # efficiency purposes + return self._get_sub_tensordict(index) + + def _remove_batch_dim(self, *args, **kwargs): + raise NotImplementedError + + def _maybe_remove_batch_dim(self, *args, **kwargs): + raise NotImplementedError + + +########################### +# Keys utils + + +class _TensorDictKeysView: + """A Key view for TensorDictBase instance. + + _TensorDictKeysView is returned when accessing tensordict.keys() and holds a + reference to the original TensorDict. This class enables us to support nested keys + when performing membership checks and when iterating over keys. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + + >>> td = TensorDict( + >>> {"a": TensorDict({"b": torch.rand(1, 2)}, [1, 2]), "c": torch.rand(1)}, + >>> [1], + >>> ) + + >>> assert "a" in td.keys() + >>> assert ("a",) in td.keys() + >>> assert ("a", "b") in td.keys() + >>> assert ("a", "c") not in td.keys() + + >>> assert set(td.keys()) == {("a", "b"), "c"} + """ + + def __init__( + self, + tensordict: T, + include_nested: bool, + leaves_only: bool, + is_leaf: Callable[[Type], bool] | None = None, + sort: bool = False, + ) -> None: + self.tensordict = tensordict + self.include_nested = include_nested + self.leaves_only = leaves_only + if is_leaf is None: + is_leaf = _default_is_leaf + self.is_leaf = is_leaf + self.sort = sort + + def __iter__(self) -> Iterator[str | tuple[str, ...]]: + def _iter(): + if not self.include_nested: + if self.leaves_only: + for key in self._keys(): + target_class = self.tensordict.entry_class(key) + if not self.is_leaf(target_class): + continue + yield key + else: + yield from self._keys() + else: + yield from ( + key if len(key) > 1 else key[0] + for key in self._iter_helper(self.tensordict) + ) + + if self.sort: + + def keyfunc(key): + return ".".join(key) if isinstance(key, tuple) else key + + yield from sorted( + _iter(), + key=keyfunc, + ) + else: + yield from _iter() + + def _iter_helper( + self, tensordict: T, prefix: tuple | None = None + ) -> Iterable[str | tuple[str, ...]]: + for key, value in self._items(tensordict): + full_key = self._combine_keys(prefix, key) + cls = type(value) + while cls is list: + # For lazy stacks + value = value[0] + cls = type(value) + is_tc = _is_tensor_collection(cls) + if self.include_nested and is_tc: + if not is_non_tensor(cls): + yield from self._iter_helper(value, prefix=full_key) + is_leaf = self.is_leaf(cls) + if not self.leaves_only or is_leaf: + yield full_key + + def _combine_keys(self, prefix: tuple | None, key: NestedKey) -> tuple: + if prefix is not None: + return prefix + (key,) + return (key,) + + def __len__(self) -> int: + return sum(1 for _ in self) + + def _items( + self, tensordict: TensorDictBase | None = None + ) -> Iterable[tuple[NestedKey, CompatibleType]]: + if tensordict is None: + tensordict = self.tensordict + if is_tensorclass(tensordict): + tensordict = tensordict._tensordict + if isinstance(tensordict, TensorDict): + return tensordict._tensordict.items() + from tensordict.nn import TensorDictParams + + if isinstance(tensordict, TensorDictParams): + assert tensordict._param_td is not None + return tensordict._param_td.items() + from tensordict._lazy import ( + _CustomOpTensorDict, + _iter_items_lazystack, + LazyStackedTensorDict, + ) + + if isinstance(tensordict, LazyStackedTensorDict): + return _iter_items_lazystack(tensordict, return_none_for_het_values=True) + if isinstance(tensordict, _CustomOpTensorDict): + # it's possible that a TensorDict contains a nested LazyStackedTensorDict, + # or _CustomOpTensorDict, so as we iterate through the contents we need to + # be careful to not rely on tensordict._tensordict existing. + return ( + (key, tensordict._get_str(key, NO_DEFAULT)) + for key in tensordict._source.keys() + ) + from tensordict.persistent import PersistentTensorDict + + if isinstance(tensordict, PersistentTensorDict): + return ( + (key, tensordict._get_str(key, NO_DEFAULT)) for key in tensordict.keys() + ) + raise NotImplementedError(type(tensordict)) + + def _keys(self) -> _TensorDictKeysView: + return self.tensordict._tensordict.keys() + + def __contains__(self, key: NestedKey) -> bool: + key = _unravel_key_to_tuple(key) + if not key: + raise TypeError(_NON_STR_KEY_ERR) + + if isinstance(key, str): + if key in self._keys(): + if self.leaves_only: + # TODO: make this faster for LazyStacked without compromising regular + return not _is_tensor_collection( + type(self.tensordict._get_str(key, NO_DEFAULT)) + ) + return True + return False + else: + # thanks to _unravel_key_to_tuple we know the key is a tuple + if len(key) == 1: + return key[0] in self._keys() + elif self.include_nested: + item_root = self.tensordict._get_str(key[0], default=None) + if item_root is not None: + entry_type = type(item_root) + if issubclass(entry_type, Tensor): + return False + # TODO: make this faster for LazyStacked without compromising regular + _is_tensordict = _is_tensor_collection(entry_type) + if _is_tensordict: + # # this will call _unravel_key_to_tuple many times + # return key[1:] in self.tensordict._get_str(key[0], NO_DEFAULT).keys(include_nested=self.include_nested) + # this won't call _unravel_key_to_tuple but requires to get the default which can be suboptimal + if len(key) >= 3: + leaf_td = item_root._get_tuple(key[1:-1], None) + if leaf_td is None or ( + not _is_tensor_collection(type(leaf_td)) + ): + return False + else: + leaf_td = item_root + return key[-1] in leaf_td.keys() + return False + # this is reached whenever there is more than one key but include_nested is False + if all(isinstance(subkey, str) for subkey in key): + raise TypeError(_NON_STR_KEY_TUPLE_ERR) + + def __repr__(self): + include_nested = f"include_nested={self.include_nested}" + leaves_only = f"leaves_only={self.leaves_only}" + return f"{type(self).__name__}({list(self)},\n{indent(include_nested, 4 * ' ')},\n{indent(leaves_only, 4 * ' ')})" + + +def _set_tensor_dict( # noqa: F811 + __dict__, + _parameters, + _buffers, + hooks, + module: torch.nn.Module, + name: str, + tensor: torch.Tensor, + inplace: bool, +) -> None: + """Simplified version of torch.nn.utils._named_member_accessor.""" + was_buffer = False + out = _parameters.pop(name, None) # type: ignore[assignment] + if out is None: + out = _buffers.pop(name, None) + was_buffer = out is not None + if out is None: + # dynamo doesn't like pop... + out = __dict__.pop(name) + if inplace: + # swap tensor and out after updating out + out_tmp = out.clone() + out.data.copy_(tensor.data) + tensor = out + out = out_tmp + + if isinstance(tensor, torch.nn.Parameter): + for hook in hooks: + output = hook(module, name, tensor) + if output is not None: + tensor = output + _parameters[name] = tensor + + if isinstance(tensor, UninitializedTensorMixin): + module.register_forward_pre_hook( + _add_batch_dim_pre_hook(), with_kwargs=True + ) + + elif was_buffer and isinstance(tensor, torch.Tensor): + _buffers[name] = tensor + else: + __dict__[name] = tensor + return out + + +def _index_to_str(index: IndexType) -> Any: + if isinstance(index, tuple): + return tuple(_index_to_str(elt) for elt in index) + if isinstance(index, slice): + return ("slice", {"start": index.start, "stop": index.stop, "step": index.step}) + if isinstance(index, range): + return ("range", {"start": index.start, "stop": index.stop, "step": index.step}) + if isinstance(index, Tensor): + return ("tensor", index.tolist(), str(index.device)) + return index + + +def _str_to_index(index: Any) -> IndexType: + if isinstance(index, tuple): + if not len(index): + return index + if index[0] == "slice": + index = index[1] + return slice(index["start"], index["stop"], index["step"]) + if index[0] == "range": + index = index[1] + return range(index["start"], index["stop"], index["step"]) + if index[0] == "tensor": + index, device = index[1:] + return torch.tensor(index, device=device) + return tuple(_index_to_str(elt) for elt in index) + return index + + +_register_tensor_class(TensorDict) +_register_tensor_class(_SubTensorDict) + + +def _save_metadata(data: TensorCollection, prefix: Path, metadata=None) -> None: + """Saves the metadata of a memmap tensordict on disk.""" + filepath = prefix / "meta.json" + if metadata is None: + metadata = {} + metadata.update( + { + "shape": list(data.shape), + "device": str(data.device), + "_type": str(type(data)), + } + ) + with open(filepath, "wb") as json_metadata: + from tensordict.utils import json_dumps + + json_str = json_dumps(metadata) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + json_metadata.write(json_str.encode("utf-8")) + else: + json_metadata.write(json_str) + + +# user did specify location and memmap is in wrong place, so we copy +def _populate_memmap( + *, dest, value, key, copy_existing, prefix, like, existsok, robust_key +): + from .utils import _encode_key_for_filesystem, _get_robust_key_setting_with_warning + + if prefix is None: + filename = None + else: + # Use smart warning that only warns when encoding would differ + effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key) + # Encode the key to make it filesystem-safe + safe_key = _encode_key_for_filesystem(key, robust=effective_robust_key) + filename = str(prefix / f"{safe_key}.memmap") + if value.is_nested: + shape = value._nested_tensor_size() + # Make the shape a memmap tensor too + if prefix is not None and filename is not None: + shape_filename = Path(filename) + shape_filename = shape_filename.with_suffix(".shape.memmap") + MemoryMappedTensor.from_tensor( + shape, + filename=shape_filename, + copy_existing=copy_existing, + existsok=existsok, + copy_data=True, + ) + else: + shape = None + memmap_tensor = MemoryMappedTensor.from_tensor( + value.data if value.requires_grad else value, + filename=filename, + copy_existing=copy_existing, + copy_data=not like, + shape=shape, + existsok=existsok, + ) + dest._tensordict[key] = memmap_tensor + return memmap_tensor + + +def _populate_empty( + *, + dest, + key, + shape, + dtype, + prefix, + robust_key, +): + from .utils import _encode_key_for_filesystem, _get_robust_key_setting_with_warning + + if prefix is None: + filename = None + else: + # Use smart warning that only warns when encoding would differ + effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key) + # Encode the key to make it filesystem-safe + safe_key = _encode_key_for_filesystem(key, robust=effective_robust_key) + filename = str(prefix / f"{safe_key}.memmap") + if isinstance(shape, torch.Tensor): + # Make the shape a memmap tensor too + if prefix is not None: + shape_filename = Path(filename) + shape_filename = shape_filename.with_suffix(".shape.memmap") + MemoryMappedTensor.from_tensor( + shape, + filename=shape_filename, + existsok=True, + copy_data=True, + ) + memmap_tensor = MemoryMappedTensor.empty( + shape=shape, + dtype=dtype, + filename=filename, + existsok=True, + ) + dest._tensordict[key] = memmap_tensor + return memmap_tensor + + +def _populate_storage( + *, + dest, + key, + shape, + dtype, + prefix, + storage, + robust_key, +): + from .utils import _encode_key_for_filesystem, _get_robust_key_setting_with_warning + + if prefix is None: + filename = None + else: + # Use smart warning that only warns when encoding would differ + effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key) + # Encode the key to make it filesystem-safe + safe_key = _encode_key_for_filesystem(key, robust=effective_robust_key) + filename = str(prefix / f"{safe_key}.memmap") + if isinstance(shape, torch.Tensor): + # Make the shape a memmap tensor too + if prefix is not None: + shape_filename = Path(filename) + shape_filename = shape_filename.with_suffix(".shape.memmap") + MemoryMappedTensor.from_tensor( + shape, + filename=shape_filename, + existsok=True, + copy_data=True, + ) + memmap_tensor = MemoryMappedTensor.from_storage( + storage=storage, + shape=shape, + dtype=dtype, + filename=filename, + ) + dest._tensordict[key] = memmap_tensor + return memmap_tensor + + +def _update_metadata(*, metadata, key, value, is_collection): + if not is_collection: + metadata[key] = { + "device": str(value.device), + "shape": ( + list(value.shape) + if not value.is_nested + else list(value._nested_tensor_size().shape) + ), + "dtype": str(value.dtype), + "is_nested": value.is_nested, + } + else: + metadata[key] = { + "type": type(value).__name__, + } + + +def from_module( + module, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, +): + """Copies the params and buffers of a module in a tensordict. + + Args: + module (nn.Module): the module to get the parameters from. + as_module (bool, optional): if ``True``, a :class:`~tensordict.nn.TensorDictParams` + instance will be returned which can be used to store parameters + within a :class:`torch.nn.Module`. Defaults to ``False``. + lock (bool, optional): if ``True``, the resulting tensordict will be locked. + Defaults to ``True``. + use_state_dict (bool, optional): if ``True``, the state-dict from the + module will be used and unflattened into a TensorDict with + the tree structure of the model. Defaults to ``False``. + + .. note:: + This is particularly useful when state-dict hooks have to be used. + + Examples: + >>> from torch import nn + >>> module = nn.TransformerDecoder( + ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), + ... num_layers=1) + >>> params = from_module(module) + >>> print(params["layers", "0", "linear1"]) + TensorDict( + fields={ + bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + return TensorDict.from_module( + module=module, as_module=as_module, lock=lock, use_state_dict=use_state_dict + ) + + +def from_modules( + *modules, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + lazy_stack: bool = False, + expand_identical: bool = False, +): + """Retrieves the parameters of several modules for ensebmle learning/feature of expects applications through vmap. + + Args: + modules (sequence of nn.Module): the modules to get the parameters from. + If the modules differ in their structure, a lazy stack is needed + (see the ``lazy_stack`` argument below). + + Keyword Args: + as_module (bool, optional): if ``True``, a :class:`~tensordict.nn.TensorDictParams` + instance will be returned which can be used to store parameters + within a :class:`torch.nn.Module`. Defaults to ``False``. + lock (bool, optional): if ``True``, the resulting tensordict will be locked. + Defaults to ``True``. + use_state_dict (bool, optional): if ``True``, the state-dict from the + module will be used and unflattened into a TensorDict with + the tree structure of the model. Defaults to ``False``. + + .. note:: + This is particularly useful when state-dict hooks have to be used. + + lazy_stack (bool, optional): whether parameters should be densly or + lazily stacked. Defaults to ``False`` (dense stack). + + .. note:: + ``lazy_stack`` and ``as_module`` are exclusive features. + + .. warning:: + There is a crucial difference between lazy and non-lazy outputs + in that non-lazy output will reinstantiate parameters with the + desired batch-size, while ``lazy_stack`` will just represent + the parameters as lazily stacked. This means that whilst the + original parameters can safely be passed to an optimizer + when ``lazy_stack=True``, the new parameters need to be passed + when it is set to ``True``. + + .. warning:: + Whilst it can be tempting to use a lazy stack to keep the + orignal parameter references, remember that lazy stack + perform a stack each time :meth:`~.get` is called. This will + require memory (N times the size of the parameters, more if a + graph is built) and time to be computed. + It also means that the optimizer(s) will contain more + parameters, and operations like :meth:`~torch.optim.Optimizer.step` + or :meth:`~torch.optim.Optimizer.zero_grad` will take longer + to be executed. In general, ``lazy_stack`` should be reserved + to very few use cases. + + expand_identical (bool, optional): if ``True`` and the same parameter (same + identity) is being stacked to itself, an expanded version of this parameter + will be returned instead. This argument is ignored when ``lazy_stack=True``. + + Examples: + >>> from torch import nn + >>> from tensordict import from_modules + >>> torch.manual_seed(0) + >>> empty_module = nn.Linear(3, 4, device="meta") + >>> n_models = 2 + >>> modules = [nn.Linear(3, 4) for _ in range(n_models)] + >>> params = from_modules(*modules) + >>> print(params) + TensorDict( + fields={ + bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + >>> # example of batch execution + >>> def exec_module(params, x): + ... with params.to_module(empty_module): + ... return empty_module(x) + >>> x = torch.randn(3) + >>> y = torch.vmap(exec_module, (0, None))(params, x) + >>> assert y.shape == (n_models, 4) + >>> # since lazy_stack = False, backprop leaves the original params untouched + >>> y.sum().backward() + >>> assert params["weight"].grad.norm() > 0 + >>> assert modules[0].weight.grad is None + + With ``lazy_stack=True``, things are slightly different: + + >>> params = TensorDict.from_modules(*modules, lazy_stack=True) + >>> print(params) + LazyStackedTensorDict( + fields={ + bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + exclusive_fields={ + }, + batch_size=torch.Size([2]), + device=None, + is_shared=False, + stack_dim=0) + >>> # example of batch execution + >>> y = torch.vmap(exec_module, (0, None))(params, x) + >>> assert y.shape == (n_models, 4) + >>> y.sum().backward() + >>> assert modules[0].weight.grad is not None + + + """ + return TensorDict.from_modules( + *modules, + lazy_stack=lazy_stack, + expand_identical=expand_identical, + lock=lock, + use_state_dict=use_state_dict, + as_module=as_module, + ) + + +def from_pytree( + pytree, + *, + batch_size: torch.Size | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, +): + """Converts a pytree to a TensorDict instance. + + This method is designed to keep the pytree nested structure as much as possible. + + Additional non-tensor keys are added to keep track of each level's identity, providing + a built-in pytree-to-tensordict bijective transform API. + + Accepted classes currently include lists, tuples, named tuples and dict. + + .. note:: + For dictionaries, non-NestedKey keys are registered separately as :class:`~tensordict.NonTensorData` + instances. + + .. note:: + Tensor-castable types (such as int, float or np.ndarray) will be converted to torch.Tensor instances. + Note that this transformation is surjective: transforming back the tensordict to a pytree will not + recover the original types. + + Examples: + >>> # Create a pytree with tensor leaves, and one "weird"-looking dict key + >>> class WeirdLookingClass: + ... pass + ... + >>> weird_key = WeirdLookingClass() + >>> # Make a pytree with tuple, lists, dict and namedtuple + >>> pytree = ( + ... [torch.randint(10, (3,)), torch.zeros(2)], + ... { + ... "tensor": torch.randn( + ... 2, + ... ), + ... "td": TensorDict({"one": 1}), + ... weird_key: torch.randint(10, (2,)), + ... "list": [1, 2, 3], + ... }, + ... {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()}, + ... ) + >>> # Build a TensorDict from that pytree + >>> td = from_pytree(pytree) + >>> # Recover the pytree + >>> pytree_recon = td.to_pytree() + >>> # Check that the leaves match + >>> def check(v1, v2): + >>> assert (v1 == v2).all() + >>> + >>> torch.utils._pytree.tree_map(check, pytree, pytree_recon) + >>> assert weird_key in pytree_recon[1] + + """ + return TensorDict.from_pytree( + pytree, + batch_size=batch_size, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + ) + + +def stack(input: Sequence[T], dim: int = 0, *, out=None) -> T: + """Stacks tensordicts into a single tensordict along the given dimension. + + This call is equivalent to calling :func:`torch.stack` but is compatible with torch.compile. + + """ + return torch.stack(input, dim=dim, out=out) + + +def lazy_stack(input: Sequence[T], dim: int = 0, *, out=None) -> T: + """Creates a lazy stack of tensordicts. + + See :meth:`~tensordict.LazyStackTensorDict.lazy_stack` for details. + """ + return TensorDict.lazy_stack(input, dim=dim, out=out) + + +def cat(input: Sequence[T], dim: int = 0, *, out=None) -> T: + """Concatenates tensordicts into a single tensordict along the given dimension. + + This call is equivalent to calling :func:`torch.cat` but is compatible with torch.compile. + + """ + return torch.cat(input, dim=dim, out=out) + + +def maybe_dense_stack(input: Sequence[T], dim: int = 0, *, out=None, **kwargs) -> T: + """Attempts to make a dense stack of tensordicts, and falls back on lazy stack when required.. + + See :meth:`~tensordict.LazyStackTensorDict.maybe_dense_stack` for details. + """ + return TensorDict.maybe_dense_stack(input, dim=dim, out=out, **kwargs) + + +def fromkeys(keys: List[NestedKey], value: Any = 0): + """Creates a tensordict from a list of keys and a single value. + + Args: + keys (list of NestedKey): An iterable specifying the keys of the new dictionary. + value (compatible type, optional): The value for all keys. Defaults to ``0``. + """ + return TensorDict.fromkeys(keys=keys, value=value) + + +def from_consolidated(filename): + """Reconstructs a tensordict from a consolidated file.""" + return TensorDict.from_consolidated(filename) + + +def load( + prefix: str | Path, + device: torch.device | None = None, + non_blocking: bool = False, + *, + out: TensorCollection | None = None, + robust_key: bool | None = None, +) -> Self: + """Loads a tensordict from disk. + + This class method is a proxy to :meth:`~.load_memmap`. + """ + return load_memmap( + prefix=prefix, + device=device, + non_blocking=non_blocking, + out=out, + robust_key=robust_key, + ) + + +def load_memmap( + prefix: str | Path, + device: torch.device | None = None, + non_blocking: bool = False, + *, + out: TensorCollection | None = None, + robust_key: bool | None = None, +) -> Self: + """Loads a memory-mapped tensordict from disk. + + Args: + prefix (str or Path to folder): the path to the folder where the + saved tensordict should be fetched. + device (torch.device or equivalent, optional): if provided, the + data will be asynchronously cast to that device. + Supports `"meta"` device, in which case the data isn't loaded + but a set of empty "meta" tensors are created. This is + useful to get a sense of the total model size and structure + without actually opening any file. + non_blocking (bool, optional): if ``True``, synchronize won't be + called after loading tensors on device. Defaults to ``False``. + out (TensorDictBase, optional): optional tensordict where the data + should be written. + robust_key (bool, optional): if ``True``, expects robust key encoding was used + when saving and decodes filenames accordingly. If ``False``, uses legacy + behavior. If ``None`` (default), emits a deprecation warning and falls + back to legacy behavior. Will default to ``True`` in v0.12. + + Examples: + >>> from tensordict import TensorDict, load_memmap + >>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0) + >>> td.memmap("./saved_td") + >>> td_load = TensorDict.load_memmap("./saved_td") + >>> assert (td == td_load).all() + + This method also allows loading nested tensordicts. + + Examples: + >>> nested = TensorDict.load_memmap("./saved_td/nested") + >>> assert nested["e"] == 0 + + A tensordict can also be loaded on "meta" device or, alternatively, + as a fake tensor. + + Examples: + >>> import tempfile + >>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) + >>> with tempfile.TemporaryDirectory() as path: + ... td.save(path) + ... td_load = load_memmap(path, device="meta") + ... print("meta:", td_load) + ... from torch._subclasses import FakeTensorMode + ... with FakeTensorMode(): + ... td_load = load_memmap(path) + ... print("fake:", td_load) + meta: TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=meta, + is_shared=False)}, + batch_size=torch.Size([]), + device=meta, + is_shared=False) + fake: TensorDict( + fields={ + a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False) + + """ + return TensorDict.load_memmap( + prefix=prefix, + device=device, + non_blocking=non_blocking, + out=out, + robust_key=robust_key, + ) + + +def save( + data: TensorCollection, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, +) -> None: + """Saves the tensordict to disk. + + This function is a proxy to :meth:`~.memmap`. + """ + return data.memmap( + prefix=prefix, + copy_existing=copy_existing, + num_threads=num_threads, + return_early=return_early, + share_non_tensor=share_non_tensor, + robust_key=robust_key, + ) + + +def memmap( + data: TensorCollection, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, +) -> Self: + """Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict. + + Args: + data (TensorDictBase): a data structure to save. + prefix (str): directory prefix where the memory-mapped tensors will + be stored. The directory tree structure will mimic the tensordict's. + copy_existing (bool): If False (default), an exception will be raised if an + entry in the tensordict is already a tensor stored on disk + with an associated file, but is not saved in the correct + location according to prefix. + If ``True``, any existing Tensor will be copied to the new location. + + Keyword Args: + num_threads (int, optional): the number of threads used to write the memmap + tensors. Defaults to `0`. + return_early (bool, optional): if ``True`` and ``num_threads>0``, + the method will return a future of the tensordict. + share_non_tensor (bool, optional): if ``True``, the non-tensor data will be + shared between the processes and writing operation (such as inplace update + or set) on any of the workers within a single node will update the value + on all other workers. If the number of non-tensor leaves is high (e.g., + sharing large stacks of non-tensor data) this may result in OOM or similar + errors. Defaults to ``False``. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + The TensorDict is then locked, meaning that any writing operations that + isn't in-place will throw an exception (eg, rename, set or remove an + entry). + Once the tensordict is unlocked, the memory-mapped attribute is turned to ``False``, + because cross-process identity is not guaranteed anymore. + + Returns: + A new tensordict with the tensors stored on disk if ``return_early=False``, + otherwise a :class:`~tensordict.utils.TensorDictFuture` instance. + + Note: + Serialising in this fashion might be slow with deeply nested tensordicts, so + it is not recommended to call this method inside a training loop. + """ + return data.memmap( + prefix=prefix, + copy_existing=copy_existing, + num_threads=num_threads, + return_early=return_early, + share_non_tensor=share_non_tensor, + robust_key=robust_key, + ) diff --git a/lib/python3.12/site-packages/tensordict/_tensorcollection.py b/lib/python3.12/site-packages/tensordict/_tensorcollection.py new file mode 100644 index 0000000000000000000000000000000000000000..d774bb9cc7e13eb55d7ca589c5a482894858e55d --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_tensorcollection.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from typing import TypeVar + +T = TypeVar("T", bound="TensorCollection") + + +class TensorCollection: + """A base class for TensorDictBase and TensorClass. + + This is an abstract base class that provides the foundation for tensor collections + like TensorDict and TensorClass. It serves as a common interface for all tensor + collection types in the tensordict library. + """ + + def __init__(self, *args, **kwargs): + """Initialize the TensorCollection. + + This is an abstract base class and should not be instantiated directly. + """ + raise NotImplementedError( + "TensorCollection is an abstract base class and cannot be instantiated directly." + ) diff --git a/lib/python3.12/site-packages/tensordict/_tensorcollection.pyi b/lib/python3.12/site-packages/tensordict/_tensorcollection.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a189b9066358ceaf0db5e4c7586f2c0e4e73883e --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_tensorcollection.pyi @@ -0,0 +1,1534 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import abc +import dataclasses +import enum +from abc import abstractmethod +from collections.abc import MutableMapping +from pathlib import Path +from typing import ( + Any, + Callable, + dataclass_transform, + Generator, + Iterator, + Literal, + OrderedDict, + overload, + Sequence, + TYPE_CHECKING, + TypeVar, +) + +import numpy as np + +import torch +from _typeshed import Incomplete +from tensordict._nestedkey import NestedKey as NestedKey +from tensordict._tensorcollection import TensorCollection +from tensordict.memmap import MemoryMappedTensor as MemoryMappedTensor +from tensordict.utils import ( + Buffer as Buffer, + cache as cache, + convert_ellipsis_to_idx as convert_ellipsis_to_idx, + DeviceType as DeviceType, + erase_cache as erase_cache, + implement_for as implement_for, + IndexType as IndexType, + infer_size_impl as infer_size_impl, + int_generator as int_generator, + is_namedtuple as is_namedtuple, + is_namedtuple_class as is_namedtuple_class, + is_non_tensor as is_non_tensor, + lazy_legacy as lazy_legacy, + lock_blocked as lock_blocked, + prod as prod, + set_lazy_legacy as set_lazy_legacy, + strtobool as strtobool, + TensorDictFuture as TensorDictFuture, + unravel_key as unravel_key, + unravel_key_list as unravel_key_list, +) +from torch import multiprocessing as mp, nn, Tensor + +class _NoDefault(enum.IntEnum): + ZERO = 0 + +NO_DEFAULT: Incomplete + +class _BEST_ATTEMPT_INPLACE: + def __bool__(self) -> bool: ... + +BEST_ATTEMPT_INPLACE: Incomplete +CompatibleType = Tensor + +T = TypeVar("T", bound="TensorCollection") + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + +class TensorCollection: + _autocast: bool = False + _nocast: bool = False + _frozen: bool = False + def __init__( + self, + *args, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: DeviceType | None = None, + names: Sequence[str] | None = None, + non_blocking: bool | None = None, + lock: bool = False, + **kwargs, + ) -> None: ... + @property + def is_meta(self) -> bool: ... + def __bool__(self) -> bool: ... + def __ne__(self, other: object) -> Self: ... + def __xor__(self, other: TensorCollection | float): ... + def __or__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __eq__(self, other: object) -> Self: ... + def __ge__(self, other: object) -> Self: ... + def __gt__(self, other: object) -> Self: ... + def __le__(self, other: object) -> Self: ... + def __lt__(self, other: object) -> Self: ... + def __deepcopy__(self, memodict={}): ... + def __iter__(self) -> Generator: ... + def __len__(self) -> int: ... + def __contains__(self, key: NestedKey) -> bool: ... + def __getitem__( + self, index: IndexType + ) -> Self | Tensor | TensorCollection | Any: ... + __getitems__ = __getitem__ + + def __setitem__(self, index: IndexType, value: Any) -> None: ... + def __delitem__(self, key: NestedKey) -> Self: ... + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: ... + def all(self, dim: int | None = None) -> bool | TensorCollection: ... + def any(self, dim: int | None = None) -> bool | TensorCollection: ... + def isfinite(self) -> Self: ... + def isnan(self) -> Self: ... + def isneginf(self) -> Self: ... + def isposinf(self) -> Self: ... + def isreal(self) -> Self: ... + @overload + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + ) -> Self: ... + @overload + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + @overload + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + @overload + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + ) -> Self: ... + @overload + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + @overload + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + @overload + def cummin(self, dim: int, *, return_indices: bool = True) -> Self: ... + @overload + def cummin( + self, dim: int, *, reduce: bool, return_indices: bool = True + ) -> Self | torch.Tensor: ... + def cummin( + self, dim: int, *, reduce: bool | None = None, return_indices: bool = True + ) -> Self | torch.Tensor: ... + @overload + def cummax(self, dim: int, *, return_indices: bool = True) -> Self: ... + @overload + def cummax( + self, dim: int, *, reduce: bool, return_indices: bool = True + ) -> Self | torch.Tensor: ... + def cummax( + self, dim: int, *, reduce: bool | None = None, return_indices: bool = True + ) -> Self | torch.Tensor: ... + @overload + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + @overload + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + @overload + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + @overload + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + @overload + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + ) -> Self: ... + @overload + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + @overload + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + @overload + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + ) -> Self: ... + @overload + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + @overload + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + ) -> Self: ... + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + reduce: bool, + ) -> Self | torch.Tensor: ... + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + def auto_batch_size_(self, batch_dims: int | None = None) -> Self: ... + def auto_device_(self) -> Self: ... + @classmethod + def from_dataclass( + cls, dataclass, *, auto_batch_size: bool = False, as_tensorclass: bool = False + ): ... + @classmethod + def from_any(cls, obj, *, auto_batch_size: bool = False): ... + @classmethod + def from_dict( + cls, + input_dict, + *, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: list[str] | None = None, + ): ... + def from_dict_instance( + self, + input_dict, + batch_size: Incomplete | None = None, + device: Incomplete | None = None, + batch_dims: Incomplete | None = None, + names: list[str] | None = None, + ): ... + @classmethod + def from_pytree( + cls, + pytree, + *, + batch_size: torch.Size | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + ): ... + def to_pytree(self): ... + @classmethod + def from_h5(cls, filename, mode: str = "r"): ... + @classmethod + def from_module( + cls, + module, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + ): ... + @classmethod + def from_modules( + cls, + *modules, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + lazy_stack: bool = False, + expand_identical: bool = False, + ): ... + def to_module( + self, + module: nn.Module, + *, + inplace: bool | None = None, + return_swap: bool = True, + swap_dest: Incomplete | None = None, + use_state_dict: bool = False, + non_blocking: bool = False, + memo: Incomplete | None = None, + ): ... + @property + def shape(self) -> torch.Size: ... + @shape.setter + def shape(self, value) -> torch.Size: ... + @property + def batch_size(self) -> torch.Size: ... + def size(self, dim: int | None = None) -> torch.Size | int: ... + @property + def data(self) -> Self: ... + @property + def grad(self) -> Self: ... + def data_ptr(self, *, storage: bool = False): ... + @grad.setter + def grad(self, grad) -> None: ... + def zero_grad(self, set_to_none: bool = True) -> Self: ... + @property + def dtype(self): ... + @property + def batch_dims(self) -> int: ... + def ndimension(self) -> int: ... + @property + def ndim(self) -> int: ... + def dim(self) -> int: ... + def numel(self) -> int: ... + @property + def depth(self) -> int: ... + @overload + def expand(self, *shape: int) -> Self: ... + @overload + def expand(self, shape: torch.Size) -> Self: ... + def expand_as(self, other: TensorCollection | torch.Tensor) -> Self: ... + def new_zeros( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): ... + def new_ones( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): ... + def new_empty( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): ... + def new_full( + self, + size: torch.Size, + fill_value, + *, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): ... + def new_tensor( + self, + data: torch.Tensor | TensorCollection, + *, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + pin_memory: bool | None = None, + ): ... + def unbind(self, dim: int) -> tuple[T, ...]: ... + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorCollection, ...]: ... + def unsqueeze(self, dim: int) -> Self: ... + def squeeze(self, dim: int | None = None) -> Self: ... + @overload + def reshape(self, *shape: int) -> Self: ... + @overload + def reshape(self, shape: list | tuple) -> Self: ... + def reshape(self, *args, **kwargs) -> Self: ... + def repeat_interleave( + self, + repeats: torch.Tensor | int, + dim: int | None = None, + *, + output_size: int | None = None, + ) -> Self: ... + def repeat(self, *repeats: int) -> Self: ... + def cat_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: ... + def stack_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: ... + def cat_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | list[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: ... + def stack_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | list[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: ... + @classmethod + def stack(cls, input, dim: int = 0, *, out: Incomplete | None = None): ... + @classmethod + def cat(cls, input, dim: int = 0, *, out: Incomplete | None = None): ... + @classmethod + def lazy_stack( + cls, input, dim: int = 0, *, out: Incomplete | None = None, **kwargs + ): ... + @classmethod + def maybe_dense_stack( + cls, input, dim: int = 0, *, out: Incomplete | None = None, **kwargs + ): ... + def split( + self, split_size: int | list[int], dim: int = 0 + ) -> list[TensorCollection]: ... + def gather(self, dim: int, index: Tensor, out: T | None = None) -> Self: ... + @overload + def view(self, *shape: int): ... + @overload + def view(self, dtype) -> Self: ... + @overload + def view(self, shape: torch.Size): ... + def view( + self, + *shape: int, + size: list | tuple | torch.Size | None = None, + batch_size: torch.Size | None = None, + ): ... + def transpose(self, dim0, dim1): ... + def swapaxes(self, axis0: int, axis1: int): ... + def swapdims(self, dim0: int, dim1: int): ... + def flip(self, dims: int | tuple[int, ...]): ... + def fliplr(self): ... + def flipud(self): ... + def roll( + self, shifts: int | tuple[int, ...], dims: int | tuple[int, ...] | None = None + ): ... + def rot90(self, k: int = 1, dims: tuple[int, int] = (0, 1)): ... + def narrow(self, dim: int, start: int, length: int): ... + def tile(self, dims: tuple[int, ...]): ... + def broadcast_to(self, shape: tuple[int, ...]): ... + def atleast_1d(self): ... + def atleast_2d(self): ... + def atleast_3d(self): ... + def movedim( + self, source: int | tuple[int, ...], destination: int | tuple[int, ...] + ): ... + def moveaxis( + self, source: int | tuple[int, ...], destination: int | tuple[int, ...] + ): ... + @overload + def permute(self, *dims: int): ... + @overload + def permute(self, dims: list | tuple): ... + @property + def names(self): ... + def refine_names(self, *names) -> Self: ... + def rename(self, *names, **rename_map): ... + def rename_(self, *names, **rename_map): ... + @property + def device(self) -> torch.device | None: ... + @device.setter + def device(self, value: DeviceType) -> torch.device | None: ... + def clear(self) -> Self: ... + def clear_refs_for_compile_(self) -> Self: ... + @classmethod + def fromkeys(cls, keys: list[NestedKey], value: Any = 0): ... + def popitem(self) -> tuple[NestedKey, CompatibleType]: ... + def clear_device_(self) -> Self: ... + def param_count(self, *, count_duplicates: bool = True) -> int: ... + def bytes(self, *, count_duplicates: bool = True) -> int: ... + def pin_memory( + self, num_threads: int | None = None, inplace: bool = False + ) -> Self: ... + def pin_memory_(self, num_threads: int | str = 0) -> Self: ... + def cpu(self, **kwargs) -> Self: ... + def cuda(self, device: int | None = None, **kwargs) -> Self: ... + @property + def is_cuda(self): ... + @property + def is_cpu(self): ... + def state_dict( + self, + destination: Incomplete | None = None, + prefix: str = "", + keep_vars: bool = False, + flatten: bool = False, + ) -> OrderedDict[str, Any]: ... + def load_state_dict( + self, + state_dict: OrderedDict[str, Any], + strict: bool = True, + assign: bool = False, + from_flatten: bool = False, + ) -> Self: ... + def is_shared(self) -> bool: ... + def is_memmap(self) -> bool: ... + def share_memory_(self) -> Self: ... + def densify(self, layout: torch.layout = ...): ... + @property + def saved_path(self): ... + def consolidate( + self, + filename: Path | str | None = None, + *, + num_threads: int = 0, + device: torch.device | None = None, + non_blocking: bool = False, + inplace: bool = False, + return_early: bool = False, + use_buffer: bool = False, + share_memory: bool = False, + pin_memory: bool = False, + metadata: bool = False, + ) -> Self: ... + @classmethod + def from_consolidated(cls, filename): ... + def is_consolidated(self): ... + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: ... + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def make_memmap_from_tensor( + self, + key: NestedKey, + tensor: torch.Tensor, + *, + copy_data: bool = True, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def save( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + def dumps( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + def memmap( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: ... + def memmap_like( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + existsok: bool = True, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + @classmethod + def load( + cls, prefix: str | Path, *args, robust_key: bool | None = None, **kwargs + ) -> Self: ... + def load_( + self, prefix: str | Path, *args, robust_key: bool | None = None, **kwargs + ): ... + @classmethod + def load_memmap( + cls, + prefix: str | Path, + device: torch.device | None = None, + non_blocking: bool = False, + *, + out: TensorCollection | None = None, + robust_key: bool | None = None, + ) -> Self: ... + def load_memmap_(self, prefix: str | Path, robust_key: bool | None = None): ... + def memmap_refresh_(self): ... + def entry_class(self, key: NestedKey) -> type: ... + def set( + self, + key: NestedKey, + item: CompatibleType, + inplace: bool = False, + *, + non_blocking: bool = False, + **kwargs: Any, + ) -> Self: ... + def set_non_tensor(self, key: NestedKey, value: Any): ... + def get_non_tensor(self, key: NestedKey, default=...): ... + def filter_non_tensor_data(self) -> Self: ... + def filter_empty_(self): ... + def set_at_( + self, + key: NestedKey, + value: CompatibleType, + index: IndexType, + *, + non_blocking: bool = False, + ) -> Self: ... + def set_( + self, key: NestedKey, item: CompatibleType, *, non_blocking: bool = False + ) -> Self: ... + @overload + def get(self, key): ... + @overload + def get(self, key, default): ... + def get(self, key: NestedKey, *args, **kwargs) -> CompatibleType: ... + @overload + def get_at(self, key, index): ... + @overload + def get_at(self, key, index, default): ... + def get_at( + self, + key: NestedKey, + *args, + **kwargs, + ) -> CompatibleType: ... + def get_item_shape(self, key: NestedKey): ... + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> Self: ... + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: ... + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + idx: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: ... + def replace(self, *args, **kwargs): ... + def create_nested(self, key): ... + def copy_(self, tensordict: T, non_blocking: bool = False) -> Self: ... + def copy_at_( + self, tensordict: T, idx: IndexType, non_blocking: bool = False + ) -> Self: ... + def is_empty(self) -> bool: ... + def setdefault( + self, key: NestedKey, default: CompatibleType, inplace: bool = False + ) -> CompatibleType: ... + def items( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Incomplete | None = None, + *, + sort: bool = False, + ) -> Iterator[tuple[str, CompatibleType]]: ... + def non_tensor_items(self, include_nested: bool = False): ... + def values( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Incomplete | None = None, + *, + sort: bool = False, + ) -> Iterator[CompatibleType]: ... + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[type], bool] | None = None, + *, + sort: bool = False, + ): ... + def pop(self, key: NestedKey, default: Any = ...) -> CompatibleType: ... + @property + def sorted_keys(self) -> list[NestedKey]: ... + def flatten(self, start_dim: int = 0, end_dim: int = -1): ... + def unflatten(self, dim, unflattened_size): ... + def _transform_keys( + self, key_transform: Callable[[NestedKey], NestedKey] + ) -> Self: ... + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> Self: ... + def del_(self, key: NestedKey) -> Self: ... + def gather_and_stack( + self, dst: int, group: "dist.ProcessGroup" | None = None + ) -> Self | None: ... + def send( + self, + dst: int, + *, + group: dist.ProcessGroup | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> None: ... + def recv( + self, + src: int, + *, + group: dist.ProcessGroup | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> int: ... + @classmethod + def from_remote_init( + cls: T, + src: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ) -> Self: ... + def init_remote( + self, + dst: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ): ... + def isend( + self, + dst: int, + *, + group: "dist.ProcessGroup" | None = None, # noqa: F821 + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> int: ... + def irecv( + self, + src: int, + *, + group: dist.ProcessGroup | None = None, + return_premature: bool = False, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: ... + def reduce( + self, + dst, + op: Incomplete | None = None, + async_op: bool = False, + return_premature: bool = False, + group: Incomplete | None = None, + ): ... + def apply_(self, fn: Callable, *others, **kwargs) -> Self: ... + def apply( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = ..., + names: Sequence[str] | None = ..., + inplace: bool = False, + default: Any = ..., + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorCollection | None = None, + **constructor_kwargs, + ) -> Self | None: ... + def named_apply( + self, + fn: Callable, + *others: T, + nested_keys: bool = False, + batch_size: Sequence[int] | None = None, + device: torch.device | None = ..., + names: Sequence[str] | None = ..., + inplace: bool = False, + default: Any = ..., + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorCollection | None = None, + **constructor_kwargs, + ) -> Self | None: ... + def map( + self, + fn: Callable[[TensorCollection], TensorCollection | None], + dim: int = 0, + num_workers: int | None = None, + *, + out: TensorCollection | None = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = False, + pbar: bool = False, + mp_start_method: str | None = None, + ): ... + def map_iter( + self, + fn: Callable[[TensorCollection], TensorCollection | None], + dim: int = 0, + num_workers: int | None = None, + *, + shuffle: bool = False, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = True, + pbar: bool = False, + mp_start_method: str | None = None, + ): ... + def record_stream(self, stream: torch.cuda.Stream): ... + def __add__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __iadd__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __abs__(self): ... + def __truediv__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __itruediv__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __mod__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __mul__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __imul__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __sub__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __isub__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __pow__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __ipow__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def abs(self) -> Self: ... + def abs_(self) -> Self: ... + def acos(self) -> Self: ... + def acos_(self) -> Self: ... + def exp(self) -> Self: ... + def exp_(self) -> Self: ... + def neg(self) -> Self: ... + def neg_(self) -> Self: ... + def reciprocal(self) -> Self: ... + def reciprocal_(self) -> Self: ... + def sigmoid(self) -> Self: ... + def sigmoid_(self) -> Self: ... + def sign(self) -> Self: ... + def sign_(self) -> Self: ... + def sin(self) -> Self: ... + def sin_(self) -> Self: ... + def sinh(self) -> Self: ... + def sinh_(self) -> Self: ... + def tan(self) -> Self: ... + def tan_(self) -> Self: ... + def tanh(self) -> Self: ... + def tanh_(self) -> Self: ... + def trunc(self) -> Self: ... + def trunc_(self) -> Self: ... + def lgamma(self) -> Self: ... + def lgamma_(self) -> Self: ... + def frac(self) -> Self: ... + def frac_(self) -> Self: ... + def expm1(self) -> Self: ... + def expm1_(self) -> Self: ... + def log(self) -> Self: ... + def log_(self) -> Self: ... + def log10(self) -> Self: ... + def log10_(self) -> Self: ... + def log1p(self) -> Self: ... + def log1p_(self) -> Self: ... + def log2(self) -> Self: ... + def log2_(self) -> Self: ... + def ceil(self) -> Self: ... + def ceil_(self) -> Self: ... + def floor(self) -> Self: ... + def floor_(self) -> Self: ... + def round(self) -> Self: ... + def round_(self) -> Self: ... + def erf(self) -> Self: ... + def erf_(self) -> Self: ... + def erfc(self) -> Self: ... + def erfc_(self) -> Self: ... + def asin(self) -> Self: ... + def asin_(self) -> Self: ... + def atan(self) -> Self: ... + def atan_(self) -> Self: ... + def cos(self) -> Self: ... + def cos_(self) -> Self: ... + def cosh(self) -> Self: ... + def cosh_(self) -> Self: ... + def add( + self, + other: TensorCollection | torch.Tensor, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: ... + def add_(self, other: TensorCollection | float, *, alpha: float | None = None): ... + def lerp( + self, + end: TensorCollection | torch.Tensor, + weight: TensorCollection | torch.Tensor | float, + ): ... + def lerp_( + self, end: TensorCollection | float, weight: TensorCollection | float + ): ... + def addcdiv( + self, + other1: TensorCollection | torch.Tensor, + other2: TensorCollection | torch.Tensor, + value: float | None = 1, + ): ... + def addcdiv_(self, other1, other2, *, value: float | None = 1): ... + def addcmul(self, other1, other2, *, value: float | None = 1): ... + def addcmul_(self, other1, other2, *, value: float | None = 1): ... + def sub( + self, + other: TensorCollection | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ): ... + def rsub( + self, + other: TensorCollection | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ): ... + def sub_(self, other: TensorCollection | float, alpha: float | None = None): ... + def mod(self, other: TensorCollection | torch.Tensor) -> Self: ... + def mul_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def mul( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def maximum_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def maximum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def minimum_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def minimum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def clamp( + self, + min: TensorCollection | torch.Tensor = None, + max: TensorCollection | torch.Tensor = None, + *, + out=None, + ) -> Self: ... + def logsumexp(self, dim=None, keepdim=False, *, out=None): ... + def clamp_max_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def clamp_max( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def clamp_min_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def clamp_min( + self, + other: TensorCollection | torch.Tensor, + default: str | CompatibleType | None = None, + ) -> Self: ... + def pow_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def pow( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def div_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def div( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def sqrt_(self) -> Self: ... + def sqrt(self) -> Self: ... + def __enter__(self): ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: types.TracebackType | None, + ): ... + def select( + self, *keys: NestedKey, inplace: bool = False, strict: bool = True + ) -> Self: ... + def exclude(self, *keys: NestedKey, inplace: bool = False) -> Self: ... + def to_tensordict(self, *, retain_none: bool | None = None) -> Self: ... + def clone(self, recurse: bool = True, **kwargs) -> Self: ... + def copy(self) -> Self: ... + def to_padded_tensor( + self, padding: float = 0.0, mask_key: NestedKey | None = None + ) -> Self: ... + def as_tensor(self) -> Self: ... + def to_lazystack(self, dim: int = 0) -> Self: ... + def to_mds( + self, + *, + out: str | tuple[str, str], + columns: dict[str, str] | None = None, + writer: "MDSWriter" | None = None, + ) -> None: ... + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + ) -> dict[str, Any]: ... + @classmethod + def from_list( + cls, + input, + *, + auto_batch_size: bool | None = None, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: list[str] | None = None, + lazy: bool | None = None, + ) -> Self: ... + def tolist( + self, + *, + convert_nodes: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ) -> list[Any]: ... + def numpy(self) -> np.ndarray | dict[str, Any]: ... + def to_namedtuple(self, dest_cls: type | None = None) -> Any: ... + @classmethod + def from_namedtuple(cls, named_tuple, *, auto_batch_size: bool = False): ... + def from_tuple( + cls, + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ): ... + def logical_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def bitwise_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + @classmethod + def from_struct_array( + cls, struct_array: np.ndarray, device: torch.device | None = None + ) -> Self: ... + def to_struct_array(self) -> np.ndarray: ... + def to_h5(self, filename, **kwargs) -> Any: ... + def empty( + self, + recurse: bool = False, + *, + batch_size: Incomplete | None = None, + device=..., + names: Incomplete | None = None, + ) -> Self: ... + def zero_(self) -> Self: ... + def fill_(self, key: NestedKey, value: float | bool) -> Self: ... + def masked_fill_(self, mask: Tensor, value: float | bool) -> Self: ... + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: ... + def where( + self, + condition, + other, + *, + out: Incomplete | None = None, + pad: Incomplete | None = None, + update_batch_size: bool = False, + ) -> Self: ... + def masked_select(self, mask: Tensor) -> Self: ... + def is_contiguous(self) -> bool: ... + def contiguous(self) -> Self: ... + def flatten_keys( + self, + separator: str = ".", + inplace: bool = False, + is_leaf: Callable[[type], bool] | None = None, + ) -> Self: ... + def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> Self: ... + def split_keys( + self, + *key_sets, + inplace: bool = False, + strict: bool = True, + reproduce_struct: bool = False, + ): ... + def separates( + self, + *keys: NestedKey, + default: Any = NO_DEFAULT, + strict: bool = True, + filter_empty: bool = True, + ) -> Self: ... + def norm( + self, + *, + out=None, + dtype: torch.dtype | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ): ... + def softmax(self, dim: int, dtype: torch.dtype | None = None): ... + @property + def is_locked(self) -> bool: ... + @is_locked.setter + def is_locked(self, value: bool) -> None: ... + def lock_(self) -> Self: ... + def unlock_(self) -> Self: ... + @overload + def to( + self, + device: int | device | None = ..., + dtype: torch.dtype | None = ..., + non_blocking: bool = ..., + inplace: bool = False, + ) -> Self: ... + @overload + def to(self, dtype: torch.dtype, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, tensor: Tensor, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, *, other: T, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, *, batch_size: torch.Size) -> Self: ... + def to(self, *args, **kwargs) -> Self: ... + def is_floating_point(self) -> bool: ... + def double(self): ... + def float(self): ... + def int(self): ... + def bool(self): ... + def half(self): ... + def type(self, dst_type): ... + @property + def requires_grad(self) -> bool: ... + def requires_grad_(self, requires_grad: bool = True) -> Self: ... + def detach_(self) -> Self: ... + def detach(self) -> Self: ... + def bfloat16(self) -> Self: ... + def complex128(self) -> Self: ... + def complex32(self) -> Self: ... + def complex64(self) -> Self: ... + def float16(self) -> Self: ... + def float32(self) -> Self: ... + def float64(self) -> Self: ... + def int16(self) -> Self: ... + def int32(self) -> Self: ... + def int64(self) -> Self: ... + def int8(self) -> Self: ... + def qint32(self) -> Self: ... + def qint8(self) -> Self: ... + def quint4x2(self) -> Self: ... + def quint8(self) -> Self: ... + def uint16(self) -> Self: ... + def uint32(self) -> Self: ... + def uint64(self) -> Self: ... + def uint8(self) -> Self: ... + +class NonTensorDataBase(TensorClass): ... +class NonTensorData(NonTensorDataBase): ... +class MetaData(NonTensorDataBase): ... +class NonTensorStack(TensorCollection): ... + +@dataclass_transform() +def tensorclass( + cls: T = None, + /, + *, + autocast: bool = False, + frozen: bool = False, + nocast: bool = False, + shadow: bool = False, + tensor_only: bool = False, +) -> Self | None: ... +def is_non_tensor(obj) -> bool: ... +def from_dataclass( + obj: Any, + *, + dest_cls: Type | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + batch_size: torch.Size | None = None, + frozen: bool = False, + autocast: bool = False, + nocast: bool = False, + inplace: bool = False, + shadow: bool = False, + tensor_only: bool = False, + device: torch.device | None = None, +) -> Any: ... diff --git a/lib/python3.12/site-packages/tensordict/_tensordict/__init__.py b/lib/python3.12/site-packages/tensordict/_tensordict/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a78ac2c10951ba3d716c4656c9f52c1d1777fe76 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_tensordict/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import warnings + +from tensordict.utils import ( # noqa + _unravel_key_to_tuple, + unravel_key, + unravel_key_list, + unravel_keys, +) + +warnings.warn( + "tensordict._tensordict will soon be removed in favour of tensordict._C.", + category=DeprecationWarning, +) + +__all__ = [ + "_unravel_key_to_tuple", + "unravel_key", + "unravel_key_list", + "unravel_keys", +] diff --git a/lib/python3.12/site-packages/tensordict/_tensordict/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/_tensordict/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bebf1fe1fae94017514def76c6d419a1b89e0c9 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/_tensordict/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/_torch_func.py b/lib/python3.12/site-packages/tensordict/_torch_func.py new file mode 100644 index 0000000000000000000000000000000000000000..5734ae6faea180d18e78055102e569014529adb7 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_torch_func.py @@ -0,0 +1,1036 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import contextlib +import functools +from functools import partial + +from typing import Any, Callable, Sequence, Tuple, TypeVar + +import torch +from tensordict._lazy import LazyStackedTensorDict +from tensordict._td import TensorDict +from tensordict.base import ( + _is_leaf_nontensor, + _is_tensor_collection, + NO_DEFAULT, + TensorDictBase, +) +from tensordict.persistent import PersistentTensorDict +from tensordict.utils import ( + _check_keys, + _ErrorInteceptor, + _is_tensorclass, + _pass_through, + _shape, + _zip_strict, + DeviceType, + is_tensorclass, + lazy_legacy, + set_lazy_legacy, +) +from torch import Tensor +from torch.nn.parameter import ( + UninitializedBuffer, + UninitializedParameter, + UninitializedTensorMixin, +) + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + +TD_HANDLED_FUNCTIONS: dict[Callable, Callable] = {} +LAZY_TD_HANDLED_FUNCTIONS: dict[Callable, Callable] = {} +T = TypeVar("T", bound="TensorDictBase") + +try: + from torch.utils._pytree import tree_leaves +except ImportError: + from torch.utils._pytree import tree_flatten + + def tree_leaves(pytree): + """Torch 2.0 compatible version of tree_leaves.""" + return tree_flatten(pytree)[0] + + +def implements_for_td(torch_function: Callable) -> Callable[[Callable], Callable]: + """Register a torch function override for TensorDict.""" + + @functools.wraps(torch_function) + def decorator(func: Callable) -> Callable: + TD_HANDLED_FUNCTIONS[torch_function] = func + return func + + return decorator + + +def implements_for_lazy_td(torch_function: Callable) -> Callable[[Callable], Callable]: + """Register a torch function override for TensorDict.""" + + @functools.wraps(torch_function) + def decorator(func: Callable) -> Callable: + LAZY_TD_HANDLED_FUNCTIONS[torch_function] = func + return func + + return decorator + + +@implements_for_td(torch.unbind) +def _unbind(td: T, *args: Any, **kwargs: Any) -> tuple[T, ...]: + return td.unbind(*args, **kwargs) + + +@implements_for_td(torch.unflatten) +def _unflatten(td: T, *args: Any, **kwargs: Any) -> tuple[T, ...]: + return td.unflatten(*args, **kwargs) + + +@implements_for_td(torch.flatten) +def _flatten(td: T, *args: Any, **kwargs: Any) -> tuple[T, ...]: + return td.flatten(*args, **kwargs) + + +@implements_for_td(torch.flip) +def _flip(td: T, dims: Sequence[int]) -> T: + return td.flip(dims) + + +@implements_for_td(torch.fliplr) +def _fliplr(td: T) -> T: + return td.fliplr() + + +@implements_for_td(torch.flipud) +def _flipud(td: T) -> T: + return td.flipud() + + +@implements_for_td(torch.roll) +def _roll( + td: T, shifts: int | Sequence[int], dims: int | Sequence[int] | None = None +) -> T: + return td.roll(shifts, dims) + + +@implements_for_td(torch.rot90) +def _rot90(td: T, k: int = 1, dims: Sequence[int] = (0, 1)) -> T: + return td.rot90(k, dims) + + +@implements_for_td(torch.narrow) +def _narrow(td: T, dim: int, start: int, length: int) -> T: + return td.narrow(dim, start, length) + + +@implements_for_td(torch.tile) +def _tile(td: T, dims: Sequence[int]) -> T: + return td.tile(dims) + + +@implements_for_td(torch.broadcast_to) +def _broadcast_to(td: T, shape: Sequence[int]) -> T: + return td.broadcast_to(shape) + + +@implements_for_td(torch.atleast_1d) +def _atleast_1d(td: T) -> T: + return td.atleast_1d() + + +@implements_for_td(torch.atleast_2d) +def _atleast_2d(td: T) -> T: + return td.atleast_2d() + + +@implements_for_td(torch.atleast_3d) +def _atleast_3d(td: T) -> T: + return td.atleast_3d() + + +@implements_for_td(torch.transpose) +def _transpose(td: T, *args: Any, **kwargs: Any) -> tuple[T, ...]: + return td.transpose(*args, **kwargs) + + +@implements_for_td(torch.swapaxes) +def _swapaxes(td: T, axis0: int, axis1: int) -> T: + return td.swapaxes(axis0, axis1) + + +@implements_for_td(torch.swapdims) +def _swapdims(td: T, dim0: int, dim1: int) -> T: + return td.swapdims(dim0, dim1) + + +@implements_for_td(torch.gather) +def _gather( + input: T, + dim: int, + index: Tensor, + *, + sparse_grad: bool = False, + out: T | None = None, +) -> T: + if sparse_grad: + raise NotImplementedError( + "sparse_grad=True not implemented for torch.gather(tensordict, ...)" + ) + # the index must have as many dims as the tensordict + if not len(index): + raise RuntimeError("Cannot use torch.gather with an empty index") + dim_orig = dim + if dim < 0: + dim = input.batch_dims + dim + if dim > input.batch_dims - 1 or dim < 0: + raise RuntimeError( + f"Cannot gather tensordict with shape {input.shape} along dim {dim_orig}." + ) + + def _gather_tensor(tensor, dest_container=None, dest_key=None): + if dest_container is not None: + dest = dest_container._get_str(dest_key, default=NO_DEFAULT) + else: + dest = None + index_expand = index + while index_expand.ndim < tensor.ndim: + index_expand = index_expand.unsqueeze(-1) + target_shape = list(tensor.shape) + target_shape[dim] = index_expand.shape[dim] + index_expand = index_expand.expand(target_shape) + out = torch.gather(tensor, dim, index_expand, out=dest) + return out + + if out is None: + if len(index.shape) == input.ndim: + names = input._maybe_names() + else: + names = None + device = input.device + return type(input)._new_unsafe( + { + key: _gather_tensor(value) + for key, value in input.items(is_leaf=_is_leaf_nontensor) + }, + batch_size=index.shape, + names=names, + device=device, + ) + for key, value in input.items(is_leaf=_is_leaf_nontensor): + _gather_tensor(value, out, key) + return out + + +@implements_for_td(torch.full_like) +def _full_like(td: T, fill_value: float, *args, **kwargs: Any) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def full_like(x): + return torch.full_like( + x, fill_value=fill_value, device=device, dtype=dtype, *args, **kwargs + ) + + td_clone = td._fast_apply( + full_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with full_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.zeros_like) +def _zeros_like(td: T, *args, **kwargs: Any) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def zeros_like(x): + return torch.zeros_like(x, device=device, dtype=dtype, *args, **kwargs) + + td_clone = td._fast_apply( + zeros_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with zeros_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.ones_like) +def _ones_like(td: T, *args, **kwargs: Any) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def ones_like(x): + return torch.ones_like(x, device=device, dtype=dtype, *args, **kwargs) + + td_clone = td._fast_apply( + ones_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with ones_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.rand_like) +def _rand_like(td: T, *args, **kwargs: Any) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def rand_like(x): + return torch.rand_like(x, device=device, dtype=dtype, *args, **kwargs) + + td_clone = td._fast_apply( + rand_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with rand_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.randn_like) +def _randn_like(td: T, *args, **kwargs: Any) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def randn_like(x): + return torch.randn_like(x, device=device, dtype=dtype, *args, **kwargs) + + td_clone = td._fast_apply( + randn_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with randn_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.empty_like) +def _empty_like(td: T, *args, **kwargs) -> T: + dtype = kwargs.pop("dtype", None) + device_nd = kwargs.pop("device", NO_DEFAULT) + if device_nd is NO_DEFAULT: + device = None + else: + device = device_nd + + def empty_like(x): + return torch.empty_like(x, device=device, dtype=dtype, *args, **kwargs) + + td_clone = td._fast_apply( + empty_like, + propagate_lock=True, + device=device_nd, + ) + if len(kwargs): + raise RuntimeError( + f"keyword arguments {list(kwargs.keys())} are not " + f"supported with empty_like with TensorDict" + ) + return td_clone + + +@implements_for_td(torch.clone) +def _clone(td: T, *args: Any, **kwargs: Any) -> T: + return td.clone(*args, **kwargs) + + +@implements_for_td(torch.squeeze) +def _squeeze(td: T, *args: Any, **kwargs: Any) -> T: + return td.squeeze(*args, **kwargs) + + +@implements_for_td(torch.unsqueeze) +def _unsqueeze(td: T, *args: Any, **kwargs: Any) -> T: + return td.unsqueeze(*args, **kwargs) + + +@implements_for_td(torch.masked_select) +def _masked_select(td: T, *args: Any, **kwargs: Any) -> T: + return td.masked_select(*args, **kwargs) + + +@implements_for_td(torch.permute) +def _permute(td: T, dims: Sequence[int]) -> T: + return td.permute(*dims) + + +@implements_for_td(torch.movedim) +def _movedim(td: T, source: int | Sequence[int], destination: int | Sequence[int]) -> T: + return td.movedim(source, destination) + + +@implements_for_td(torch.moveaxis) +def _moveaxis( + td: T, source: int | Sequence[int], destination: int | Sequence[int] +) -> T: + return td.moveaxis(source, destination) + + +@implements_for_td(torch.cat) +def _cat( + list_of_tensordicts: Sequence[T], + dim: int = 0, + device: DeviceType | None = None, + out: T | None = None, +) -> T: + if not len(list_of_tensordicts): + raise RuntimeError("list_of_tensordicts cannot be empty") + + batch_size = list(list_of_tensordicts[0].batch_size) + tdtype = type(list_of_tensordicts[0]) + if dim < 0: + dim = len(batch_size) + dim + if dim >= len(batch_size): + raise RuntimeError( + f"dim must be in the range 0 <= dim < len(batch_size), got dim" + f"={dim} and batch_size={batch_size}" + ) + batch_size[dim] = sum([td.batch_size[dim] for td in list_of_tensordicts]) + batch_size = TensorDict._parse_batch_size(None, batch_size) + + # check that all tensordict match + keys = _check_keys(list_of_tensordicts, strict=True) + if out is None: + out = {} + for key in keys: + items = [td._get_str(key, NO_DEFAULT) for td in list_of_tensordicts] + if not is_compiling(): + with _ErrorInteceptor( + key, "Attempted to concatenate tensors on different devices at key" + ): + out[key] = torch.cat(items, dim) + else: + out[key] = torch.cat(items, dim) + if device is None: + device = list_of_tensordicts[0].device + for td in list_of_tensordicts[1:]: + if device == td.device: + continue + else: + device = None + break + names = None + if list_of_tensordicts[0]._has_names(): + names = list_of_tensordicts[0].names + # if we have a TD subclass, use _new_unsafe bc we know it exists. Otherwise, use + # TensorDict's one + if issubclass(tdtype, TensorDict) or _is_tensorclass(tdtype): + clz = tdtype + else: + clz = TensorDict + return clz._new_unsafe(out, device=device, batch_size=batch_size, names=names) + else: + if out.batch_size != batch_size: + raise RuntimeError( + "out.batch_size and cat batch size must match, " + f"got out.batch_size={out.batch_size} and batch_size" + f"={batch_size}" + ) + + for key in keys: + with ( + _ErrorInteceptor( + key, "Attempted to concatenate tensors on different devices at key" + ) + if not is_compiling() + else contextlib.nullcontext() + ): + if isinstance(out, TensorDict): + torch.cat( + [td.get(key) for td in list_of_tensordicts], + dim, + out=out.get(key), + ) + else: + out.set_( + key, torch.cat([td.get(key) for td in list_of_tensordicts], dim) + ) + return out + + +@implements_for_lazy_td(torch.cat) +def _lazy_cat( + list_of_tensordicts: Sequence[LazyStackedTensorDict], + dim: int = 0, + out: LazyStackedTensorDict | None = None, +) -> LazyStackedTensorDict: + # why aren't they feeding you? + if not len(list_of_tensordicts): + raise RuntimeError("list_of_tensordicts cannot be empty") + + batch_size = list(list_of_tensordicts[0].batch_size) + if dim < 0: + dim = len(batch_size) + dim + if dim >= len(batch_size): + raise RuntimeError( + f"dim must be in the range 0 <= dim < len(batch_size), got dim" + f"={dim} and batch_size={batch_size}" + ) + stack_dim = list_of_tensordicts[0].stack_dim + if any((td.stack_dim != stack_dim) for td in list_of_tensordicts): + raise RuntimeError("cat lazy stacked tds must have same stack dim") + + batch_size[dim] = sum(td.batch_size[dim] for td in list_of_tensordicts) + batch_size = torch.Size(batch_size) + + new_dim = dim + if dim > stack_dim: + new_dim = dim - 1 + + if out is None: + out = [] + if dim == stack_dim: # if dim is stack, just add all to the same list + for lazy_td in list_of_tensordicts: + if lazy_td.batch_size[stack_dim] == 0: + continue + out += lazy_td.tensordicts + else: + for i in range(len(list_of_tensordicts[0].tensordicts)): + out.append( + torch.cat( + [lazy_td.tensordicts[i] for lazy_td in list_of_tensordicts], + new_dim, + ) + ) + return type(list_of_tensordicts[0])(*out, stack_dim=stack_dim) + else: + if not isinstance(out, LazyStackedTensorDict): + return _cat(list_of_tensordicts, dim=dim, out=out) + + if out.batch_size != batch_size: + raise RuntimeError( + "out.batch_size and cat batch size must match, " + f"got out.batch_size={out.batch_size} and batch_size" + f"={batch_size}" + ) + if out.stack_dim != dim: + index_base = (slice(None),) * out.stack_dim + for i, sub_dest in enumerate(out.tensordicts): + index = index_base + (i,) + tds_to_cat = [_td[index] for _td in list_of_tensordicts] + torch.cat(tds_to_cat, dim, out=sub_dest) + else: + init_idx = 0 + for td_in in list_of_tensordicts: + sub_dest = out.tensordicts[init_idx : init_idx + td_in.shape[dim]] + init_idx += init_idx + td_in.shape[dim] + LazyStackedTensorDict.maybe_dense_stack(sub_dest, out.stack_dim).update( + td_in, inplace=True + ) + + return out + + +@implements_for_td(torch.stack) +def _stack( + list_of_tensordicts: Sequence[TensorDictBase], + dim: int = 0, + device: DeviceType | None = None, + out: T | None = None, + strict: bool = False, + contiguous: bool = False, + maybe_dense_stack: bool | None = None, +) -> T: + if not len(list_of_tensordicts): + raise RuntimeError("list_of_tensordicts cannot be empty") + if maybe_dense_stack is None: + maybe_dense_stack = lazy_legacy() + td_types = [type(td) for td in list_of_tensordicts] + is_tc = any(_is_tensorclass(td_type) for td_type in td_types) + if all(_pass_through(td) for td in list_of_tensordicts): + return type(list_of_tensordicts[0])._stack_non_tensor( + list_of_tensordicts, dim=dim + ) + list_of_tensordicts_orig = list_of_tensordicts + if is_tc: + list_of_tensordicts = [tc._tensordict for tc in list_of_tensordicts] + clz = type(list_of_tensordicts_orig[0]) + elif issubclass(td_types[0], TensorDict): + clz = td_types[0] + else: + clz = TensorDict + + batch_size = list_of_tensordicts[0].batch_size + if dim < 0: + dim = len(batch_size) + dim + 1 + + # check that all tensordict match + # Read lazy_legacy + _lazy_legacy = lazy_legacy() + + if not _lazy_legacy: + for i, td in enumerate(list_of_tensordicts[1:]): + if td.batch_size != list_of_tensordicts[0].batch_size: + if not maybe_dense_stack: + raise RuntimeError( + "stacking tensordicts requires them to have congruent batch sizes, " + f"got td[{i + 1}].batch_size={td.batch_size} and td[0].batch_size=" + f"{list_of_tensordicts[0].batch_size}" + ) + elif td.batch_dims == list_of_tensordicts[0].batch_dims: + from tensordict import lazy_stack + + return lazy_stack( + list_of_tensordicts_orig, + dim=dim, + out=out, + ) + else: + raise RuntimeError( + "Lazy stacking of tensordicts requires them to have congruent batch dimensions, " + f"got td[{i + 1}].batch_dims={td.batch_dims} and td[0].batch_dims=" + f"{list_of_tensordicts[0].batch_dims}" + ) + + if out is None: + # We need to handle tensordicts with exclusive keys and tensordicts with + # mismatching shapes. + # The first case is handled within _check_keys which fails if keys + # don't match exactly. + # The second requires a check over the tensor shapes. + device = list_of_tensordicts[0].device + if contiguous or not _lazy_legacy: + try: + keys = _check_keys(list_of_tensordicts, strict=True) + except KeyError: + if not _lazy_legacy and not contiguous: + if maybe_dense_stack: + with set_lazy_legacy(True): + return _stack( + list_of_tensordicts_orig, + dim=dim, + maybe_dense_stack=maybe_dense_stack, + ) + else: + raise RuntimeError( + "The sets of keys in the tensordicts to stack are exclusive. " + "Consider using `LazyStackedTensorDict.maybe_dense_stack` instead." + ) + raise + + if all(_tensordict._lazy for _tensordict in list_of_tensordicts): + # Let's try to see if all tensors have the same shape + # If so, we can assume that we can densly stack the sub-tds + leaves = [tree_leaves(td) for td in list_of_tensordicts] + unique_leaves_len = len({len(leaf) for leaf in leaves}) == 1 + unique_leaves_len &= ( + len({len(td.tensordicts) for td in list_of_tensordicts}) == 1 + ) + if unique_leaves_len: + # If the number of sub-tensordicts is the same, we can stack them internally + # and lazy stack them on the outside. + lazy_stack_dim = list_of_tensordicts[0].stack_dim + if dim <= lazy_stack_dim: + lazy_stack_dim += 1 + else: + dim = dim - 1 + result = LazyStackedTensorDict( + *[ + _stack( + list(subtds), + dim=dim, + maybe_dense_stack=maybe_dense_stack, + ) + for subtds in _zip_strict( + *[td.tensordicts for td in list_of_tensordicts] + ) + ], + stack_dim=lazy_stack_dim, + ) + if is_tc: + return clz._from_tensordict(result) + return result + else: + from tensordict import lazy_stack + + return lazy_stack(list_of_tensordicts_orig, dim=dim, out=out) + + out = {} + for key in keys: + out[key] = [] + is_not_init = None + tensor_shape = None + is_tensor = None + for _tensordict in list_of_tensordicts: + tensor = _tensordict._get_str(key, default=NO_DEFAULT) + if is_tensor is None: + tensor_cls = type(tensor) + # is_tensor = ( + # not _is_tensor_collection(tensor_cls) + # ) or is_tensorclass(tensor_cls) + # TODO: make sense of this, dynamo cannot pass through stack (and it's unsafe) + # only tensors should be tensors + is_tensor = not _is_tensor_collection(tensor_cls) + if is_not_init is None: + is_not_init = isinstance(tensor, UninitializedTensorMixin) + if not is_not_init: + new_tensor_shape = _shape(tensor) + if tensor_shape is not None: + if len(new_tensor_shape) != len(tensor_shape) or not all( + s1 == s2 and s1 != -1 + for s1, s2 in _zip_strict(_shape(tensor), tensor_shape) + ): + # Nested tensors will require a lazy stack + if maybe_dense_stack: + with set_lazy_legacy(True): + return _stack( + list_of_tensordicts_orig, + dim=dim, + maybe_dense_stack=maybe_dense_stack, + ) + else: + raise RuntimeError( + f"The shapes of the tensors to stack is incompatible: {new_tensor_shape} vs {tensor_shape} for key {key}." + ) + else: + tensor_shape = new_tensor_shape + + out[key].append(tensor) + out[key] = (out[key], is_not_init, is_tensor) + + def stack_fn(key, values, is_not_init, is_tensor): + if is_not_init: + return _stack_uninit_params(values, dim) + if is_tensor: + return torch.stack(values, dim) + with ( + _ErrorInteceptor( + key, "Attempted to stack tensors on different devices at key" + ) + if not is_compiling() + else contextlib.nullcontext() + ): + return _stack(values, dim, maybe_dense_stack=maybe_dense_stack) + + out = { + key: stack_fn(key, values, is_not_init, is_tensor) + for key, (values, is_not_init, is_tensor) in out.items() + } + # Add names if all tensordicts have the same `names` + names = list_of_tensordicts[0]._maybe_names() + if names is not None: + if all(td._maybe_names() == names for td in list_of_tensordicts[1:]): + names.insert(dim, None) + else: + names = None + + result = clz._new_unsafe( + out, + batch_size=LazyStackedTensorDict._compute_batch_size( + batch_size, dim, len(list_of_tensordicts) + ), + device=device, + names=names, + ) + if is_tc and not is_tensorclass(result): + return clz._from_tensordict(result) + return result + else: + out = LazyStackedTensorDict( + *list_of_tensordicts, + stack_dim=dim, + ) + if is_tc and not is_tensorclass(out): + return clz._from_tensordict(out) + return out + else: + keys = _check_keys(list_of_tensordicts) + batch_size = list(batch_size) + batch_size.insert(dim, len(list_of_tensordicts)) + batch_size = torch.Size(batch_size) + + if out.batch_size != batch_size: + raise RuntimeError( + "out.batch_size and stacked batch size must match, " + f"got out.batch_size={out.batch_size} and batch_size" + f"={batch_size}" + ) + + out_keys = set(out.keys()) + if strict: + in_keys = set(keys) + if len(out_keys - in_keys) > 0: + raise RuntimeError( + "The output tensordict has keys that are missing in the " + "tensordict that has to be written: {out_keys - in_keys}. " + "As per the call to `stack(..., strict=True)`, this " + "is not permitted." + ) + elif len(in_keys - out_keys) > 0: + raise RuntimeError( + "The resulting tensordict has keys that are missing in " + f"its destination: {in_keys - out_keys}. As per the call " + "to `stack(..., strict=True)`, this is not permitted." + ) + + try: + out._stack_onto_(list_of_tensordicts, dim) + except KeyError as err: + raise err + return out + + +@implements_for_td(torch.split) +def _split( + td: TensorDict, split_size_or_sections: int | list[int], dim: int = 0 +) -> list[TensorDictBase]: + return td.split(split_size_or_sections, dim) + + +@implements_for_td(torch.where) +def where(condition, input, other, *, out=None): + """Return a ``TensorDict`` of elements selected from either input or other, depending on condition. + + Args: + condition (BoolTensor): When ``True`` (nonzero), yield ``input``, otherwise yield ``other``. + input (TensorDictBase or Scalar): value (if ``input`` is a scalar) or values selected at indices where condition is ``True``. + other (TensorDictBase or Scalar): value (if ``other`` is a scalar) or values selected at indices where condition is ``False``. + out (Tensor, optional): the output ``TensorDictBase`` instance. + + """ + if isinstance(out, PersistentTensorDict): + raise RuntimeError( + "Cannot use a persistent tensordict as output of torch.where." + ) + return input.where(condition, other, out=out) + + +def _stack_uninit_params(list_of_params, dim: int = 0, out=None): + if out is not None: + raise NotImplementedError + if dim > 0: + raise NotImplementedError + from tensordict.utils import ( + _BatchedUninitializedBuffer, + _BatchedUninitializedParameter, + ) + + if isinstance(list_of_params[0], UninitializedParameter): + out = _BatchedUninitializedParameter( + requires_grad=list_of_params[0].requires_grad, + device=list_of_params[0].device, + dtype=list_of_params[0].dtype, + ) + elif isinstance(list_of_params[0], UninitializedBuffer): + out = _BatchedUninitializedBuffer( + requires_grad=list_of_params[0].requires_grad, + device=list_of_params[0].device, + dtype=list_of_params[0].dtype, + ) + out.batch_size = torch.Size([len(list_of_params)]) + return out + + +@implements_for_td(torch.var) +def _var(td: T, *args: Any, **kwargs: Any) -> T: + return td.var(*args, **kwargs) + + +@implements_for_td(torch.std) +def _std(td: T, *args: Any, **kwargs: Any) -> T: + return td.std(*args, **kwargs) + + +@implements_for_td(torch.mean) +def _mean(td: T, *args: Any, **kwargs: Any) -> T: + return td.mean(*args, **kwargs) + + +@implements_for_td(torch.sum) +def _sum(td: T, *args: Any, **kwargs: Any) -> T: + return td.sum(*args, **kwargs) + + +@implements_for_td(torch.maximum) +def _maximum(td: T, other: T, *args: Any, **kwargs: Any) -> T: + return td.maximum(other, *args, **kwargs) + + +@implements_for_td(torch.nn.functional.l1_loss) +def _l1_loss(input: T, target: T, *args: Any, **kwargs: Any) -> T: + reduction = kwargs.pop("reduction", "mean") + if reduction == "none": + batch_size = None + elif reduction == "mean": + batch_size = () + elif reduction == "sum": + batch_size = () + else: + raise ValueError( + f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'." + ) + return input.apply( + partial(torch.nn.functional.l1_loss, reduction=reduction), + target, + *args, + **kwargs, + batch_size=batch_size, + ) + + +@implements_for_td(torch.nn.functional.mse_loss) +def _mse_loss(input: T, target: T, *args: Any, **kwargs: Any) -> T: + reduction = kwargs.pop("reduction", "mean") + if reduction == "none": + batch_size = None + elif reduction == "mean": + batch_size = () + elif reduction == "sum": + batch_size = () + else: + raise ValueError( + f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'." + ) + return input.apply( + partial(torch.nn.functional.mse_loss, reduction=reduction), + target, + *args, + **kwargs, + batch_size=batch_size, + ) + + +@implements_for_td(torch.nn.functional.smooth_l1_loss) +def _smooth_l1_loss(input: T, target: T, *args: Any, **kwargs: Any) -> T: + reduction = kwargs.pop("reduction", "mean") + if reduction == "none": + batch_size = None + elif reduction == "mean": + batch_size = () + elif reduction == "sum": + batch_size = () + else: + raise ValueError( + f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', 'sum'." + ) + return input.apply( + partial(torch.nn.functional.smooth_l1_loss, reduction=reduction), + target, + *args, + **kwargs, + batch_size=batch_size, + ) + + +@implements_for_td(torch.autograd.grad) +def _grad( + outputs: TensorDictBase | Tuple[TensorDictBase, ...], + inputs: TensorDictBase | Tuple[TensorDictBase, ...], + grad_outputs: TensorDictBase | Tuple[TensorDictBase, ...] | None = None, + **kwargs: Any, +) -> TensorDict: + from tensordict.base import _NESTED_TENSORS_AS_LISTS + + if isinstance(outputs, tuple) and len(outputs) > 1: + raise ValueError( + "torch.autograd.grad for TensorDict only supports a single output" + ) + elif isinstance(outputs, tuple): + outputs = outputs[0] + if not isinstance(outputs, TensorDictBase): + raise ValueError( + "torch.autograd.grad for TensorDict only supports TensorDictBase as output" + ) + + if isinstance(inputs, tuple) and len(inputs) > 1: + raise ValueError( + "torch.autograd.grad for TensorDict only supports a single input" + ) + elif isinstance(inputs, tuple): + inputs = inputs[0] + if not isinstance(inputs, TensorDictBase): + raise ValueError( + "torch.autograd.grad for TensorDict only supports TensorDictBase as input" + ) + + if ( + grad_outputs is not None + and isinstance(grad_outputs, tuple) + and len(grad_outputs) > 1 + ): + raise ValueError( + "torch.autograd.grad for TensorDict only supports a single grad_output" + ) + elif isinstance(grad_outputs, tuple): + grad_outputs = grad_outputs[0] + if grad_outputs is not None and not isinstance(grad_outputs, TensorDictBase): + raise ValueError( + "torch.autograd.grad for TensorDict only supports TensorDictBase as grad_output" + ) + + if grad_outputs is not None: + tup_grad_outputs = tuple( + grad_outputs._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + ) + else: + tup_grad_outputs = None + + tup_outputs = tuple( + outputs._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + ) + + keys, all_inputs = inputs._items_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + + all_grads = torch.autograd.grad(tup_outputs, all_inputs, tup_grad_outputs, **kwargs) + + pairs = dict(_zip_strict(keys, all_grads)) + + def pop(name, val): + return pairs.pop(name, None) + + # rebuild the tensordict + return inputs._fast_apply( + pop, + named=True, + nested_keys=True, + propagate_lock=True, + ) diff --git a/lib/python3.12/site-packages/tensordict/_unbatched.py b/lib/python3.12/site-packages/tensordict/_unbatched.py new file mode 100644 index 0000000000000000000000000000000000000000..ae2e461a5e6f73504332aa2bc697a1505d153e99 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_unbatched.py @@ -0,0 +1,257 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +from functools import wraps +from typing import Any, Callable, TYPE_CHECKING + +import torch +from tensordict._tensorcollection import TensorCollection +from tensordict.base import TensorDictBase + +from tensordict.tensorclass import ( + _arg_to_tensordict, + _from_tensordict_with_copy, + _TD_PASS_THROUGH, + TD_HANDLED_FUNCTIONS, + TensorClass, +) +from tensordict.utils import ( + _getitem_batch_size, + _is_tensorclass, + _maybe_correct_neg_dim, + IndexType, + unravel_key, +) +from torch import Tensor + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +def _arg_to_tensordict_unbatched(arg, batch_size): + if _is_tensorclass(type(arg)): + arg = arg._tensordict.empty() + arg.batch_size = batch_size + return arg + elif isinstance(arg, (tuple, list)) and all( + _is_tensorclass(type(item)) for item in arg + ): + arg_list = [] + for item in arg: + item = item._tensordict.empty() + item.batch_size = batch_size + arg_list.append(item) + + return type(arg)(arg_list) + return arg + + +def _bypass(func): + @wraps(func) + def bypassed_func(self, *args, **kwargs): + meta_tensor = torch.zeros( + self.batch_size, dtype=self.dtype, device=torch.device("meta") + ) + name = func.__name__ + r = getattr(meta_tensor, name)(*args, **kwargs) + self_copy = self.copy() + self_copy.batch_size = r.shape + return self_copy + + return bypassed_func + + +_TORCH_SHAPE_OPS = ( + torch.gather, + torch.unbind, + torch.cat, + torch.stack, + torch.unflatten, + torch.flatten, + torch.split, + torch.squeeze, + torch.unsqueeze, +) + + +class UnbatchedTensor(TensorClass): + """A TensorClass that represents a tensor whose shape is ignored during shape operations. + + This class allows tensors to be stored in a TensorDict without enforcing batch size consistency. + Shape operations (e.g., reshape, unsqueeze, squeeze) on the TensorDict will return the same UnbatchedTensor instance, + while other operations (e.g., apply, key manipulation, pointwise arithmetic) may modify the underlying tensor content. + + Example: + >>> td = TensorDict(a=UnbatchedTensor(torch.randn(3, 4)), b=torch.randn(2, 3), batch_size=(2,)) + >>> td_reshaped = td.reshape((1, 2)) + >>> td_reshaped["a"] is td["a"] + True + + Note that accessing an UnbatchedTensor using `get()` and `__getitem__()` will return different results. + `get()` returns the UnbatchedTensor instance, while `__getitem__()` returns the underlying tensor content. + + Example: + >>> td.get("a") + + >>> td["a"] + tensor([[...]]) + + """ + + data: torch.Tensor | TensorDictBase + _pass_through = True + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + if func not in _TD_PASS_THROUGH or not all( + issubclass(t, (Tensor, cls, TensorDictBase)) for t in types + ): + return NotImplemented + + if kwargs is None: + kwargs = {} + + # get the output type from the arguments / keyword arguments + if len(args) > 0: + tensorclass_instance = args[0] + else: + tensorclass_instance = kwargs.get("input", kwargs["tensors"]) + if isinstance(tensorclass_instance, (tuple, list)): + tensorclass_instance = tensorclass_instance[0] + + if func not in _TORCH_SHAPE_OPS: + args = tuple(_arg_to_tensordict(arg) for arg in args) + kwargs = {key: _arg_to_tensordict(value) for key, value in kwargs.items()} + result = TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + else: + # Get a brute force batch size + args = tuple( + _arg_to_tensordict_unbatched(arg, tensorclass_instance.batch_size) + for arg in args + ) + kwargs = { + key: _arg_to_tensordict_unbatched( + value, tensorclass_instance.batch_size + ) + for key, value in kwargs.items() + } + example_td = TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + result = tensorclass_instance.copy() + result.batch_size = example_td.batch_size + return result + + if isinstance(result, (list, tuple)): + return type(result)( + _from_tensordict_with_copy(tensorclass_instance, tensordict_result) + for tensordict_result in result + ) + return _from_tensordict_with_copy(tensorclass_instance, result) + + def chunk(self, chunks: int, dim: int | None = None): + self_copy = self.copy() + if dim is None: + dim = 0 + dim = _maybe_correct_neg_dim(dim, self.batch_size) + self_copy.batch_size = ( + self.batch_size[:dim] + + (self.batch_size[dim] // chunks,) + + self.batch_size[dim + 1 :] + ) + return self_copy + + def split(self, split_size: int | list[int], dim: int | None = None): + self_copy = self.copy() + if dim is None: + dim = 0 + dim = _maybe_correct_neg_dim(dim, self.batch_size) + chunks = ( + len(split_size) + if isinstance(split_size, (list, tuple)) + else -(self.batch_size[dim] // -split_size) + ) + self_copy.batch_size = ( + self.batch_size[:dim] + + (self.batch_size[dim] // chunks,) + + self.batch_size[dim + 1 :] + ) + return self_copy + + def __getitem__(self, index: IndexType) -> Self | Tensor | TensorCollection | Any: + if isinstance(index, (tuple, str)) and unravel_key(index): + raise ValueError( + "TensorClass fields must be accessed as attributes, not items." + ) + self_copy = self.copy() + self_copy.batch_size = _getitem_batch_size(self.batch_size, index) + return self_copy + + @property + def batch_size(self): + return self._batch_size + + @batch_size.setter + def batch_size(self, batch_size): + self.__dict__["_batch_size"] = torch.Size(batch_size) + + shape = batch_size + + def unbind(self, dim: int): + return tuple( + self[(slice(None),) * dim + (0,)] for _ in range(self.batch_size[dim]) + ) + + @_bypass + def reshape(self, *shape): ... + + @_bypass + def view(self, *shape): ... + + def unsqueeze(self, dim: int): + shape = list(self.batch_size) + shape.insert(dim, 0) + self_copy = self.copy() + self_copy.batch_size = shape + return self_copy + + def transpose(self, dim0, dim1): + batch_size = list(self.batch_size) + batch_size[dim1], batch_size[dim0] = batch_size[dim0], batch_size[dim1] + self_copy = self.copy() + self_copy.batch_size = batch_size + return self_copy + + def permute(self, *dims): + if len(dims) == 1 and not isinstance(dims[0], int): + return self.permute(*dims[0]) + batch_size = list(self.batch_size) + batch_size = [batch_size[d] for d in dims] + self_copy = self.copy() + self_copy.batch_size = batch_size + return self_copy + + @classmethod + def _stack_non_tensor( + cls, list_of_non_tensor, dim: int = 0, raise_if_non_unique=False + ): + result = list_of_non_tensor[0].copy() + batch_size = list(result.batch_size) + batch_size.insert(dim, len(list_of_non_tensor)) + result.batch_size = torch.Size(batch_size) + return result + + @_bypass + def unflatten(self, dim, unflattened_size): ... + + @_bypass + def flatten(self, start_dim: int = 0, end_dim=-1): ... diff --git a/lib/python3.12/site-packages/tensordict/_version.py b/lib/python3.12/site-packages/tensordict/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..5604a398d7d29db5c8cae5d2168412b8a66beeff --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/_version.py @@ -0,0 +1,34 @@ +# file generated by setuptools-scm +# don't change, don't track in version control + +__all__ = [ + "__version__", + "__version_tuple__", + "version", + "version_tuple", + "__commit_id__", + "commit_id", +] + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Tuple + from typing import Union + + VERSION_TUPLE = Tuple[Union[int, str], ...] + COMMIT_ID = Union[str, None] +else: + VERSION_TUPLE = object + COMMIT_ID = object + +version: str +__version__: str +__version_tuple__: VERSION_TUPLE +version_tuple: VERSION_TUPLE +commit_id: COMMIT_ID +__commit_id__: COMMIT_ID + +__version__ = version = '0.11.0' +__version_tuple__ = version_tuple = (0, 11, 0) + +__commit_id__ = commit_id = None diff --git a/lib/python3.12/site-packages/tensordict/base.py b/lib/python3.12/site-packages/tensordict/base.py new file mode 100644 index 0000000000000000000000000000000000000000..26230e8fc7fed5fb39d2a391ec34e733e10fabf0 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/base.py @@ -0,0 +1,16108 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import abc +import collections +import concurrent.futures +import contextlib +import enum +import gc +import importlib +import importlib.util + +# JSON backend is now handled by utils.json_dumps +import json +import math +import os.path +import queue +import sys +import uuid +import warnings +import weakref +from collections import UserDict +from collections.abc import MutableMapping + +from concurrent.futures import Future, ThreadPoolExecutor, wait +from copy import copy +from functools import wraps +from pathlib import Path +from textwrap import indent +from threading import Thread +from types import ModuleType +from typing import ( + Any, + Callable, + Dict, + Generator, + Iterator, + List, + Literal, + Mapping, + OrderedDict, + overload, + Sequence, + Tuple, + Type, + TYPE_CHECKING, + TypeVar, +) + +import numpy as np + +import torch + +from tensordict._contextlib import LAST_OP_MAPS +from tensordict._datasets import to_mds +from tensordict._nestedkey import NestedKey +from tensordict._tensorcollection import TensorCollection +from tensordict.memmap import MemoryMappedTensor +from tensordict.utils import ( + _as_context_manager, + _CloudpickleWrapper, + _convert_list_to_stack, + _DTYPE_TO_STR_DTYPE, + _GENERIC_NESTED_ERR, + _is_dataclass as is_dataclass, + _is_list_tensor_compatible, + _is_non_tensor, + _is_number, + _is_tensorclass, + _KEY_ERROR, + _lock_warn, + _make_dtype_promotion, + _maybe_correct_neg_dim, + _parse_to, + _pass_through, + _pass_through_cls, + _pin_mem, + _PIN_MEM_TIMEOUT, + _prefix_last_key, + _proc_init, + _prune_selected_keys, + _rebuild_njt_from_njt, + _set_max_batch_size, + _shape, + _split_tensordict, + _STR_DTYPE_TO_DTYPE, + _td_fields, + _unravel_key_to_tuple, + _zip_strict, + cache, + capture_non_tensor_stack, + convert_ellipsis_to_idx, + DeviceType, + erase_cache, + expand_as_right, + implement_for, + IndexType, + infer_size_impl, + int_generator, + is_namedtuple, + is_namedtuple_class, + is_non_tensor, + lazy_legacy, + LinkedList, + list_to_stack, + lock_blocked, + prod, + set_capture_non_tensor_stack, + set_lazy_legacy, + strtobool, + TensorDictFuture, + unravel_key, + unravel_key_list, +) +from torch import multiprocessing as mp, nn, Tensor + +from torch.nn.parameter import Parameter, UninitializedTensorMixin +from torch.utils._pytree import tree_map + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + +try: + from torch import _foreach_copy_ +except ImportError: + _foreach_copy_ = None + +try: + from torch.nn.parameter import Buffer +except ImportError: + from tensordict.utils import Buffer + +_has_h5 = importlib.util.find_spec("h5py") is not None + +try: + from torch._utils import _get_available_device_type, _get_device_module +except ImportError: + from torch._utils import _get_available_device_type + + def _get_device_module(device_type: str) -> ModuleType | None: + device_module = getattr(torch, device_type, None) + if device_module is None: + raise RuntimeError( + f"Device '{device_type}' does not have a corresponding module registered as 'torch.{device_type}'." + ) + return device_module + + +# NO_DEFAULT is a sentinel value used to detect when a default argument was not provided. +# Using None is not an option since `td.get(key)` (returning None) is a valid usage. +# When passed to methods like `get()`, it signals "raise KeyError if key is missing" +# rather than returning a default value. +class _NoDefault(enum.IntEnum): + ZERO = 0 + + +NO_DEFAULT = _NoDefault.ZERO + +# _UNSET is a sentinel used in pop() to detect if a key exists without try/except. +# Unlike NO_DEFAULT (which triggers KeyError in get()), _UNSET can be returned by get() +# to indicate a missing key. This makes pop() compatible with torch.compile. +_UNSET = object() + +T = TypeVar("T", bound="TensorCollection") + + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +class _BEST_ATTEMPT_INPLACE: + def __bool__(self): + # we use an exception to exit when running `inplace = BEST_ATTEMPT_INPLACE if inplace else False` + # more than once + raise NotImplementedError + + +BEST_ATTEMPT_INPLACE = _BEST_ATTEMPT_INPLACE() + +# some complex string used as separator to concatenate and split keys in +# distributed frameworks -- make a python<3.10 specific version +if sys.version_info < (3, 10): + from typing import Union + + CompatibleType = Union[Tensor, TensorCollection] +else: + CompatibleType = Tensor | TensorCollection + +_STR_MIXED_INDEX_ERROR = "Received a mixed string-non string index. Only string-only or string-free indices are supported." + +_HEURISTIC_EXCLUDED = (Tensor, tuple, list, set, dict, np.ndarray) + +if "TD_GET_DEFAULTS_TO_NONE" in os.environ: + _GET_DEFAULTS_TO_NONE = strtobool(os.environ["TD_GET_DEFAULTS_TO_NONE"]) +else: + _GET_DEFAULTS_TO_NONE = True + + +def set_get_defaults_to_none(set_to_none: bool = True): + """Sets the default of `get` to `None` and silences deprecation warnings during calls to `get` that result in a `KeyError`. + + This can also be controlled via the environment variable ``TD_GET_DEFAULTS_TO_NONE``. + + Args: + set_to_none (bool): whether the default of `get` should be `None`, or should `get` raise a `KeyError` if + no `default` is passed and the key is absent from the `TensorDict`. + Defaults to `True`. + + """ + global _GET_DEFAULTS_TO_NONE + _GET_DEFAULTS_TO_NONE = bool(set_to_none) + + +def get_defaults_to_none(set_to_none: bool = True): + """Returns the status of `get` default value.""" + return _GET_DEFAULTS_TO_NONE + + +class _RecordDeviceTransfer: + """A class that records the device transfers during a TensorDict initialization.""" + + def __init__(self): + self.marked = False + self._has_transfer = False + + def mark(self): + if self.marked: + raise RuntimeError("Can only mark one TensorDict at a time.") + self.marked = True + self._has_transfer = False + + def unmark(self): + self.marked = False + self._has_transfer = False + + def record_transfer(self, device): + # Adds a device to all markers + self._has_transfer = True + + def has_transfer(self): + return self._has_transfer + + +_device_recorder = _RecordDeviceTransfer() + + +def _maybe_broadcast_other(op: str, n_other: int = 1) -> Callable[[Callable], Callable]: + """Ensures that elementwise ops are broadcast when an nd tensor is passed.""" + + def wrap_func(func): + @wraps(func) + def new_func(self, *others, **kwargs): + others, args = others[:n_other], others[n_other:] + need_broadcast = False + for other in others: + if other is None: + continue + if (isinstance(other, torch.Tensor) and other.ndim) or ( + _is_tensor_collection(type(other)) + and other.ndim + and other.shape != self.shape + ): + need_broadcast = True + break + if not need_broadcast: + return func(self, *others, *args, **kwargs) + others_map = [] + shape = self.shape + self_expand = self + shapes = [shape, *[other.shape for other in others if other is not None]] + shape = torch.broadcast_shapes(*shapes) + if shape != self_expand.shape: + self_expand = self_expand.expand(shape) + for other in others: + if other is None: + others_map.append(other) + continue + # broadcast dims + if shape != other.shape: + other = other.expand(shape) + others_map.append(other) + if any(isinstance(other, torch.Tensor) for other in others_map): + return self_expand._fast_apply( + lambda x: getattr(x, op)( + *[ + expand_as_right(other, x) if other is not None else None + for other in others_map + ], + *args, + **kwargs, + ) + ) + return getattr(self_expand, op)(*others_map, *args, **kwargs) + + return new_func + + return wrap_func + + +class TensorDictBase(MutableMapping, TensorCollection): + """TensorDictBase is an abstract parent class for TensorDicts, a torch.Tensor data container.""" + + _safe: bool = False + _lazy: bool = False + _inplace_set: bool = False + is_meta: bool = False + _is_locked: bool = False + _cache: bool | None = None + _is_non_tensor: bool = False + _memmap_prefix = None + _stream: torch.cuda.Stream | None = None + + @classmethod + def _new_unsafe(cls, *args, **kwargs) -> "TensorDictBase": + # This to make sure all TensorDictBase subclasses have a proper fallback if they don't have a _new_unsafe + # In other words, only TensorDict subclasses will have their type preserved, others will become TensorDict + # instances (note that TensorDictBase should not be directly subclassed outside of this codebase, as it is + # highly abstract). + from tensordict._td import TensorDict + + return TensorDict._new_unsafe(*args, **kwargs) + + def __bool__(self) -> bool: + raise RuntimeError("Converting a tensordict to boolean value is not permitted") + + def __abs__(self) -> Self: + """Returns a new TensorDict instance with absolute values of all tensors. + + Returns: + A new TensorDict instance with the same key set as the original, + but with all tensors having their absolute values computed. + + .. seealso:: :meth:`~.abs` + + """ + return self.abs() + + def __neg__(self) -> Self: + """Returns a new TensorDict instance with negated values of all tensors. + + Returns: + A new TensorDict instance with the same key set as the original, + but with all tensors having their values negated. + + .. seealso:: :meth:`~.neg` + + """ + return self.neg() + + @abc.abstractmethod + def __ne__(self, other: object) -> Self: + """NOT operation over two tensordicts, for evey key. + + The two tensordicts must have the same key set. + + Args: + other (TensorDictBase, dict, or float): the value to compare against. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + @abc.abstractmethod + def __xor__(self, other: TensorCollection | torch.Tensor | float): + """XOR operation over two tensordicts, for evey key. + + The two tensordicts must have the same key set. + + Args: + other (TensorDictBase, dict, or float): the value to compare against. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + def __rxor__(self, other: TensorCollection | torch.Tensor | float): + """XOR operation over two tensordicts, for evey key. + + Wraps `__xor__` as it is assumed to be commutative. + """ + return self.__xor__(other) + + @abc.abstractmethod + def __or__(self, other: TensorCollection | torch.Tensor) -> Self: + """OR operation over two tensordicts, for evey key. + + The two tensordicts must have the same key set. + + Args: + other (TensorDictBase, dict, or float): the value to compare against. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + def __ror__(self, other: TensorCollection | torch.Tensor) -> Self: + """Right-side OR operation over two tensordicts, for evey key. + + This is a wrapper around `__or__` since it is assumed to be commutative. + """ + return self | other + + def __invert__(self) -> Self: + """Returns a new TensorDict instance with all tensors inverted (i.e., bitwise NOT operation). + + Returns: + A new TensorDict instance with the same key set as the original, + but with all tensors having their bits inverted. + """ + keys, vals = self._items_list(True, True) + vals = [~v for v in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def __and__(self, other: TensorCollection | torch.Tensor | float) -> Self: + """Returns a new TensorDict instance with all tensors performing a logical or bitwise AND operation with the given value. + + Args: + other: The value to perform the AND operation with. + + Returns: + A new TensorDict instance with the same key set as the original, + but with all tensors having performed a AND operation with the given value. + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list(True, True, sorting_keys=keys) + vals = [(v1 & v2) for v1, v2 in zip(vals, other_val)] + else: + vals = [(v & other) for v in vals] + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + __rand__ = __and__ + + @abc.abstractmethod + def __eq__(self, other: object) -> Self: + """Compares two tensordicts against each other, for every key. The two tensordicts must have the same key set. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + @abc.abstractmethod + def __ge__(self, other: object) -> Self: + """Compares two tensordicts against each other using the "greater or equal" operator, for every key. The two tensordicts must have the same key set. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + @abc.abstractmethod + def __gt__(self, other: object) -> Self: + """Compares two tensordicts against each other using the "greater than" operator, for every key. The two tensordicts must have the same key set. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + @abc.abstractmethod + def __le__(self, other: object) -> Self: + """Compares two tensordicts against each other using the "lower or equal" operator, for every key. The two tensordicts must have the same key set. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + @abc.abstractmethod + def __lt__(self, other: object) -> Self: + """Compares two tensordicts against each other using the "lower than" operator, for every key. The two tensordicts must have the same key set. + + Returns: + a new TensorDict instance with all tensors are boolean + tensors of the same shape as the original tensors. + + """ + raise NotImplementedError + + def __repr__(self) -> str: + try: + fields = _td_fields(self) + field_str = indent(f"fields={{{fields}}}", 4 * " ") + batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ") + device_str = indent(f"device={self.device}", 4 * " ") + is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ") + string = ",\n".join([field_str, batch_size_str, device_str, is_shared_str]) + except AttributeError: + # When using torch.compile, an exception may be raised with a tensordict object + # that has no attribute (no _tensordict or no _batch_size). + # To get the proper erro message and not an attribute error raised during __repr__, + # we simply default to '...' when trying to print the TD content. + string = "..." + return f"{type(self).__name__}(\n{string})" + + def __iter__(self) -> Generator: + """Iterates over the first shape-dimension of the tensordict.""" + if not self.batch_dims: + raise StopIteration + yield from self.unbind(0) + + def __len__(self) -> int: + """Returns the length of first dimension, if there is, otherwise 0.""" + batch_size = self.batch_size + if not batch_size: + return 0 + return batch_size[0] + + def __deepcopy__( + self, memo: Dict[Any, Any] + ) -> "tensordict.TensorDict": # noqa # type: ignore + return self.clone() + + def __contains__(self, key: NestedKey) -> bool: # type: ignore + if isinstance(key, str): + return key in self.keys() + if isinstance(key, tuple): + key = unravel_key(key) + if not key: + raise RuntimeError( + "key must be a NestedKey (a str or a possibly tuple of str)." + ) + return key in self.keys(True, is_leaf=_is_leaf_nontensor) + raise RuntimeError( + "key must be a NestedKey (a str or a possibly tuple of str)." + ) + + def __getitem__(self, index: IndexType) -> Self | Tensor | TensorCollection | Any: + """Indexes all tensors according to the provided index. + + The index can be a (nested) key or any valid shape index given the + tensordict batch size. + + If the index is a nested key and the result is a :class:`~tensordict.NonTensorData` + object, the content of the non-tensor is returned. + + Examples: + >>> td = TensorDict({"root": torch.arange(2), ("nested", "entry"): torch.arange(2)}, [2]) + >>> td["root"] + torch.tensor([0, 1]) + >>> td["nested", "entry"] + torch.tensor([0, 1]) + >>> td[:1] + TensorDict( + fields={ + nested: TensorDict( + fields={ + entry: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([1]), + device=None, + is_shared=False), + root: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([1]), + device=None, + is_shared=False) + """ + istuple = isinstance(index, tuple) + if istuple or isinstance(index, str): + # _unravel_key_to_tuple will return an empty tuple if the index isn't a NestedKey + idx_unravel = _unravel_key_to_tuple(index) + if idx_unravel: + return self._get_tuple_maybe_non_tensor(idx_unravel, NO_DEFAULT) + + if (istuple and not index) or (not istuple and index is Ellipsis): + # empty tuple returns self + return self + if not istuple: + if isinstance(index, int): + return self._index_tensordict(index) + # we only want tuple indices + index = (index,) + # # convert range/np.ndarray to tensor: this is not cheap + # index = tuple( + # torch.tensor(idx) if isinstance(idx, (np.ndarray, range)) else idx + # for idx in index + # ) + if istuple and any(idx is Ellipsis for idx in index): + index = convert_ellipsis_to_idx(index, self.batch_size) + if all(isinstance(idx, slice) and idx == slice(None) for idx in index): + return self + + return self._index_tensordict(index) + + # this is necessary for data collectors for instance, otherwise indexing + # will always be achieved one element at a time. + __getitems__ = __getitem__ + + def _get_sub_tensordict(self, idx: IndexType) -> Self: + """Returns a _SubTensorDict with the desired index.""" + from tensordict._td import _SubTensorDict + + return _SubTensorDict(source=self, idx=idx) + + @abc.abstractmethod + def __setitem__( + self, + index: IndexType, + value: Any, + ) -> None: + raise NotImplementedError + + def __delitem__(self, key: NestedKey) -> Self: + return self.del_(key) + + def __getstate__(self) -> dict[str, Any]: + result = dict(self.__dict__) + for key in ( + "_last_op", + "_cache", + "__lock_parents_weakrefs", + ): + result.pop(key, None) + return result + + def __setstate__(self, state: dict[str, Any]) -> None: + for key, value in state.items(): + setattr(self, key, value) + self._cache = None + self._last_op = None + if self._is_locked: + # this can cause avoidable overhead, as we will be locking the leaves + # then locking their parent, and the parent of the parent, every + # time re-locking tensordicts that have already been locked. + # To avoid this, we should lock only at the root, but it isn't easy + # to spot what the root is... + self._is_locked = False + self.lock_() + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + from tensordict._torch_func import TD_HANDLED_FUNCTIONS + + if kwargs is None: + kwargs = {} + if func not in TD_HANDLED_FUNCTIONS or not all( + issubclass(t, (Tensor, TensorDictBase)) or _is_tensorclass(t) for t in types + ): + return NotImplemented + return TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + + @abc.abstractmethod + def all(self, dim: int | None = None) -> bool | TensorCollection: + """Checks if all values are True/non-null in the tensordict. + + Args: + dim (int, optional): if ``None``, returns a boolean indicating + whether all tensors return `tensor.all() == True` + If integer, all is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + + """ + raise NotImplementedError + + @abc.abstractmethod + def any(self, dim: int | None = None) -> bool | TensorCollection: + """Checks if any value is True/non-null in the tensordict. + + Args: + dim (int, optional): if ``None``, returns a boolean indicating + whether all tensors return `tensor.any() == True`. + If integer, all is called upon the dimension specified if + and only if this dimension is compatible with + the tensordict shape. + + """ + raise NotImplementedError + + def isfinite(self) -> Self: + """Returns a new tensordict with boolean elements representing if each element is finite or not. + + Real values are finite when they are not NaN, negative infinity, or infinity. Complex values are finite when both their real and imaginary parts are finite. + + """ + keys, vals = self._items_list(True, True) + vals = [val.isfinite() for val in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def isnan(self) -> Self: + """Returns a new tensordict with boolean elements representing if each element of input is NaN or not. + + Complex values are considered NaN when either their real and/or imaginary part is NaN. + + """ + keys, vals = self._items_list(True, True) + vals = [val.isnan() for val in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def isneginf(self) -> Self: + """Tests if each element of input is negative infinity or not.""" + keys, vals = self._items_list(True, True) + vals = [val.isneginf() for val in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def isposinf(self) -> Self: + """Tests if each element of input is negative infinity or not.""" + keys, vals = self._items_list(True, True) + vals = [val.isposinf() for val in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def isreal(self) -> Self: + """Returns a new tensordict with boolean elements representing if each element of input is real-valued or not.""" + keys, vals = self._items_list(True, True) + vals = [val.isreal() for val in vals] + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + @overload + def amin( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + ) -> Self: ... + + @overload + def amin( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + + def amin( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the minimum values of all elements in the input tensordict. + + Same as :meth:`~.min` with ``return_indices=False``. + """ + return self._cast_reduction( + reduction_name="amin", + dim=dim, + keepdim=keepdim, + further_reduce=reduce, + tuple_ok=False, + values_only=True, + call_on_nested=False, + ) + + @overload + def min( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + + @overload + def min( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + + def min( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the minimum values of all elements in the input tensordict. + + Args: + dim (int, optional): if ``None``, returns a dimensionless + tensordict containing the min value of all leaves (if this can be computed). + If integer, `min` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + return_argmins (bool, optional): :func:`~torch.min` returns a named tuple with values and indices + when the ``dim`` argument is passed. The ``TensorDict`` equivalent of this is to return a tensorclass + with entries ``"values"`` and ``"indices"`` with idendical structure within. Defaults to ``True``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.min(dim=0) + min( + indices=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + vals=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.min() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.min(reduce=True) + tensor(-2.9953) + + """ + result = self._cast_reduction( + reduction_name="min", + dim=dim, + keepdim=keepdim, + further_reduce=reduce, + tuple_ok=False, + values_only=not return_indices, + call_on_nested=False, + ) + if dim is not NO_DEFAULT and return_indices: + # Split the tensordict + from torch.return_types import min + + values_dict = {} + indices_dict = {} + for key in result.keys(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + if key[-1] == "values": + values_dict[key] = key[:-1] + else: + indices_dict[key] = key[:-1] + return min( + result.split_keys(values_dict, indices_dict)[:2], + ) + return result + + @overload + def amax( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + ) -> Self: ... + + @overload + def amax( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + + def amax( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the maximum values of all elements in the input tensordict. + + Same as :meth:`~.max` with ``return_indices=False``. + """ + return self._cast_reduction( + reduction_name="amax", + dim=dim, + keepdim=keepdim, + further_reduce=reduce, + tuple_ok=False, + values_only=True, + call_on_nested=False, + ) + + @overload + def max( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + + @overload + def max( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + + def max( + self, + dim: int | NO_DEFAULT = NO_DEFAULT, + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the maximum values of all elements in the input tensordict. + + Args: + dim (int, optional): if ``None``, returns a dimensionless + tensordict containing the max value of all leaves (if this can be computed). + If integer, `max` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + return_argmins (bool, optional): :func:`~torch.max` returns a named tuple with values and indices + when the ``dim`` argument is passed. The ``TensorDict`` equivalent of this is to return a tensorclass + with entries ``"values"`` and ``"indices"`` with idendical structure within. Defaults to ``True``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.max(dim=0) + max( + indices=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + vals=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.max() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.max(reduce=True) + tensor(3.2942) + + """ + result = self._cast_reduction( + reduction_name="max", + dim=dim, + keepdim=keepdim, + further_reduce=reduce, + tuple_ok=False, + values_only=not return_indices, + call_on_nested=False, + ) + if dim is not NO_DEFAULT and return_indices: + # Split the tensordict + from torch.return_types import max + + values_dict = {} + indices_dict = {} + for key in result.keys(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + if key[-1] == "values": + values_dict[key] = key[:-1] + else: + indices_dict[key] = key[:-1] + return max( + result.split_keys(values_dict, indices_dict)[:2], + ) + return result + + @overload + def cummin( + self, + dim: int, + *, + return_indices: bool = True, + ) -> Self: ... + + @overload + def cummin( + self, + dim: int, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + + def cummin( + self, + dim: int, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the cumulative minimum values of all elements in the input tensordict. + + Args: + dim (int): integer representing the dimension along which to perform the cummin operation. + + Keyword Args: + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + return_argmins (bool, optional): :func:`~torch.cummin` returns a named tuple with values and indices + when the ``dim`` argument is passed. The ``TensorDict`` equivalent of this is to return a tensorclass + with entries ``"values"`` and ``"indices"`` with idendical structure within. Defaults to ``True``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.cummin(dim=0) + cummin( + indices=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + vals=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.cummin(reduce=True, dim=0) + torch.return_types.cummin(...) + + """ + result = self._cast_reduction( + reduction_name="cummin", + dim=dim, + further_reduce=reduce, + tuple_ok=False, + values_only=not return_indices, + call_on_nested=False, + batch_size=self.batch_size, + ) + if isinstance(result, (torch.Tensor, torch.return_types.cummin)): + return result + if dim is not NO_DEFAULT and return_indices: + # Split the tensordict + from torch.return_types import cummin + + values_dict = {} + indices_dict = {} + for key in result.keys(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + if key[-1] == "values": + values_dict[key] = key[:-1] + else: + indices_dict[key] = key[:-1] + return cummin( + result.split_keys(values_dict, indices_dict)[:2], + ) + return result + + @overload + def cummax( + self, + dim: int, + *, + return_indices: bool = True, + ) -> Self: ... + + @overload + def cummax( + self, + dim: int, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + + def cummax( + self, + dim: int, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the cumulative maximum values of all elements in the input tensordict. + + Args: + dim (int): integer representing the dimension along which to perform the cummax operation. + + Keyword Args: + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + return_argmins (bool, optional): :func:`~torch.cummax` returns a named tuple with values and indices + when the ``dim`` argument is passed. The ``TensorDict`` equivalent of this is to return a tensorclass + with entries ``"values"`` and ``"indices"`` with idendical structure within. Defaults to ``True``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.cummax(dim=0) + cummax( + indices=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + vals=TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False), + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.cummax(reduce=True, dim=0) + torch.return_types.cummax(...) + + """ + result = self._cast_reduction( + reduction_name="cummax", + dim=dim, + further_reduce=reduce, + tuple_ok=False, + values_only=not return_indices, + call_on_nested=False, + batch_size=self.batch_size, + ) + if isinstance(result, (torch.Tensor, torch.return_types.cummin)): + return result + if dim is not NO_DEFAULT and return_indices: + # Split the tensordict + from torch.return_types import cummax + + values_dict = {} + indices_dict = {} + for key in result.keys(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + if key[-1] == "values": + values_dict[key] = key[:-1] + else: + indices_dict[key] = key[:-1] + return cummax( + result.split_keys(values_dict, indices_dict)[:2], + ) + return result + + @overload + def mean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + + @overload + def mean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + + @overload + def mean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + + def mean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the mean value of all elements in the input tensordict. + + Args: + dim (int, tuple of int, str, optional): if ``None``, returns a dimensionless + tensordict containing the mean value of all leaves (if this can be computed). + If integer or tuple of integers, `mean` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to dtype before the operation is performed. + This is useful for preventing data type overflows. Default: ``None``. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + key_transform (Callable[[NestedKey], NestedKey], optional): A function to transform key names. + If provided, all keys in the result will be transformed using this function. + For string keys, the function receives a string. For tuple keys, it receives a tuple. + Only applied when ``reduce=False``. Default: ``None``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.mean(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.mean() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.mean(reduce=True) + tensor(-0.0547) + >>> td.mean(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.mean(reduce=True, dim="feature") + tensor([[1., 1., 1., 1.], + [1., 1., 1., 1.], + [1., 1., 1., 1.]]) + >>> td.mean(reduce=True, dim=0) + tensor([[1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.]]) + >>> # Using key_transform to add prefix to keys + >>> td.mean(key_transform=lambda key: f"avg_{key}") + TensorDict( + fields={ + avg_a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + avg_b: TensorDict( + fields={ + avg_c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + avg_d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + + """ + # if dim is NO_DEFAULT and not keepdim: + # dim = None + # keepdim = False + result = self._cast_reduction( + reduction_name="mean", + dim=dim, + keepdim=keepdim, + dtype=dtype, + further_reduce=reduce, + ) + if key_transform is not None and not reduce: + result = result._transform_keys(key_transform) + return result + + @overload + def nanmean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + + @overload + def nanmean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + + def nanmean( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the mean of all non-NaN elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the mean value of all leaves (if this can be computed). + If integer or tuple of integers, `mean` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to dtype before the operation is performed. + This is useful for preventing data type overflows. Default: ``None``. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.nanmean(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.nanmean() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.nanmean(reduce=True) + tensor(-0.0547) + >>> td.nanmean(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.nanmean(reduce=True, dim="feature") + tensor([[1., 1., 1., 1.], + [1., 1., 1., 1.], + [1., 1., 1., 1.]]) + >>> td.nanmean(reduce=True, dim=0) + tensor([[1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.]]) + + """ + return self._cast_reduction( + reduction_name="nanmean", + keepdim=keepdim, + dim=dim, + dtype=dtype, + further_reduce=reduce, + ) + + @overload + def prod( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + + @overload + def prod( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + + def prod( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the produce of values of all elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the prod value of all leaves (if this can be computed). + If integer or tuple of integers, `prod` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to dtype before the operation is performed. + This is useful for preventing data type overflows. Default: ``None``. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.prod(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.prod() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.prod(reduce=True) + tensor(-0.) + >>> td.prod(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.prod(reduce=True, dim="feature") + tensor([[1., 1., 1., 1.], + [1., 1., 1., 1.], + [1., 1., 1., 1.]]) + >>> td.prod(reduce=True, dim=0) + tensor([[1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1.]]) + + """ + result = self._cast_reduction( + reduction_name="prod", + dim=dim, + keepdim=False, + tuple_ok=False, + dtype=dtype, + further_reduce=reduce, + ) + if keepdim: + if isinstance(dim, tuple): + dim = dim[0] + if dim not in (None, NO_DEFAULT): + result = result.unsqueeze(dim) + else: + result = result.reshape([1 for _ in self.shape]) + return result + + @overload + def sum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + + @overload + def sum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + + @overload + def sum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + + def sum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the sum value of all elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the sum value of all leaves (if this can be computed). + If integer or tuple of integers, `sum` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to dtype before the operation is performed. + This is useful for preventing data type overflows. Default: ``None``. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + key_transform (Callable[[NestedKey], NestedKey], optional): A function to transform key names. + If provided, all keys in the result will be transformed using this function. + For string keys, the function receives a string. For tuple keys, it receives a tuple. + Only applied when ``reduce=False``. Default: ``None``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.sum(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.sum() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.sum(reduce=True) + tensor(-0.) + >>> td.sum(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.sum(reduce=True, dim="feature") + tensor([[15., 15., 15., 15.], + [15., 15., 15., 15.], + [15., 15., 15., 15.]]) + >>> td.sum(reduce=True, dim=0) + tensor([[9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.]]) + + """ + result = self._cast_reduction( + reduction_name="sum", + dim=dim, + keepdim=keepdim, + dtype=dtype, + further_reduce=reduce, + ) + if key_transform is not None and not reduce: + result = result._transform_keys(key_transform) + return result + + @overload + def nansum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + + @overload + def nansum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + + def nansum( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the sum of all non-NaN elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the sum value of all leaves (if this can be computed). + If integer or tuple of integers, `sum` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to dtype before the operation is performed. + This is useful for preventing data type overflows. Default: ``None``. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.nansum(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.nansum() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.nansum(reduce=True) + tensor(-0.) + >>> td.nansum(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.nansum(reduce=True, dim="feature") + tensor([[15., 15., 15., 15.], + [15., 15., 15., 15.], + [15., 15., 15., 15.]]) + >>> td.nansum(reduce=True, dim=0) + tensor([[9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.], + [9., 9., 9., 9., 9.]]) + + """ + return self._cast_reduction( + reduction_name="nansum", + dim=dim, + keepdim=keepdim, + dtype=dtype, + further_reduce=reduce, + ) + + @overload + def std( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + ) -> Self: ... + + @overload + def std( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + + @overload + def std( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + + def std( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the standard deviation value of all elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the sum value of all leaves (if this can be computed). + If integer or tuple of integers, `std` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to Bessel's correction, correction=1. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.std(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.std() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.std(reduce=True) + tensor(1.0006) + >>> td.std(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.std(reduce=True, dim="feature") + tensor([[0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.]]) + >>> td.std(reduce=True, dim=0) + tensor([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + + """ + result = self._cast_reduction( + reduction_name="std", + dim=dim, + keepdim=keepdim, + correction=correction, + further_reduce=reduce, + ) + if key_transform is not None and not reduce: + result = result._transform_keys(key_transform) + return result + + @overload + def var( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + ) -> Self: ... + + @overload + def var( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + + @overload + def var( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + + def var( + self, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + correction: int = 1, + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the variance value of all elements in the input tensordict. + + Args: + dim (int, tuple of int, optional): if ``None``, returns a dimensionless + tensordict containing the sum value of all leaves (if this can be computed). + If integer or tuple of integers, `var` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to Bessel's correction, correction=1. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.var(dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.var() + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.var(reduce=True) + tensor(1.0006) + >>> td.var(dim="feature") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> td = TensorDict( + ... a=torch.ones(3, 4, 5), + ... b=TensorDict( + ... c=torch.ones(3, 4, 5), + ... d=torch.ones(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.var(reduce=True, dim="feature") + tensor([[0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.]]) + >>> td.var(reduce=True, dim=0) + tensor([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + + """ + result = self._cast_reduction( + reduction_name="var", + dim=dim, + keepdim=keepdim, + correction=correction, + further_reduce=reduce, + ) + if key_transform is not None and not reduce: + result = result._transform_keys(key_transform) + return result + + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + interpolation: str = "linear", + ) -> Self: ... + + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + interpolation: str = "linear", + reduce: bool, + ) -> Self | torch.Tensor: ... + + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + interpolation: str = "linear", + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: ... + + def quantile( + self, + q: float | torch.Tensor, + dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT, + keepdim: bool = NO_DEFAULT, + *, + interpolation: str = "linear", + reduce: bool | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self | torch.Tensor: # noqa: D417 + """Returns the q-th quantile of all elements in the input tensordict. + + Args: + q (float or torch.Tensor): quantile to compute, must be between 0 and 1. + dim (int, tuple of int, str, optional): if ``None``, returns a dimensionless + tensordict containing the quantile value of all leaves (if this can be computed). + If integer or tuple of integers, `quantile` is called upon the dimension specified if + and only if this dimension is compatible with the tensordict + shape. + Only the `"feature"` string is currently permitted. Using `dim="feature"` will + achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the + shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict + with the same structure as ``self`` with reduced feature dimensions will be returned. + keepdim (bool): whether the output tensor has dim retained or not. + + Keyword Args: + interpolation (str): interpolation method to use when the desired quantile lies + between two data points. Options are 'linear', 'lower', 'higher', 'midpoint', and 'nearest'. + Defaults to 'linear'. + reduce (bool, optional): if ``True``, the reduction will occur across all TensorDict values + and a single reduced tensor will be returned. + Defaults to ``False``. + key_transform (Callable[[NestedKey], NestedKey], optional): A function to transform key names. + If provided, all keys in the result will be transformed using this function. + For string keys, the function receives a string. For tuple keys, it receives a tuple. + Only applied when ``reduce=False``. Default: ``None``. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict( + ... a=torch.randn(3, 4, 5), + ... b=TensorDict( + ... c=torch.randn(3, 4, 5, 6), + ... d=torch.randn(3, 4, 5), + ... batch_size=(3, 4, 5), + ... ), + ... batch_size=(3, 4) + ... ) + >>> td.quantile(0.5, dim=0) # median along dim 0 + TensorDict( + fields={ + a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + >>> td.quantile(0.5) # median of all elements + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.quantile(0.5, reduce=True) # single median value + tensor(0.1234) + >>> td.quantile(0.5, dim="feature") # median along feature dimensions + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> # Multiple quantiles + >>> td.quantile(torch.tensor([0.25, 0.5, 0.75]), dim=0) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + + """ + result = self._cast_reduction( + reduction_name="quantile", + dim=dim, + keepdim=keepdim, + q=q, + interpolation=interpolation, + further_reduce=reduce, + ) + if key_transform is not None and not reduce: + result = result._transform_keys(key_transform) + return result + + @abc.abstractmethod + def _cast_reduction( + self, + *, + reduction_name, + dim=NO_DEFAULT, + keepdim=NO_DEFAULT, + dtype, + tuple_ok=True, + further_reduce: bool, + **kwargs, + ): + raise NotImplementedError + + def auto_batch_size_( + self, batch_dims: int | None = None, keep_compliant_size: bool = False + ) -> Self: + """Sets the maximum batch-size for the tensordict, up to an optional batch_dims. + + Args: + batch_dims (int, optional): if provided, the batch-size will be at + most ``batch_dims`` long. + keep_compliant_size (bool, optional): if `True`, a sub-tensordict with a compliant + size will not see its shape be changed in case `batch_dims` is passed. + If `False`, all contained tensordicts will have a `batch_dims` that matches + `batch_dims`. + Defaults to `False`. + + Returns: + self + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[]) + >>> td.auto_batch_size_() + >>> print(td.batch_size) + torch.Size([3, 4]) + >>> td.auto_batch_size_(batch_dims=1) + >>> print(td.batch_size) + torch.Size([3]) + + """ + _set_max_batch_size(self, batch_dims, keep_compliant_size=keep_compliant_size) + return self + + def auto_device_(self) -> Self: + """Automatically sets the device, if it is unique. + + Returns: self with the edited ``device`` attribute. + + """ + devices = { + value.device + for value in self.values(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + if value.device is not None + } + if len(devices) == 1: + self.clear_device_() + self._set_device(list(devices)[0]) + else: + self.clear_device_() + return self + + @classmethod + def from_list( + cls, + input, + *, + auto_batch_size: bool | None = None, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: list[str] | None = None, + lazy: bool | None = None, + ) -> Self: + if lazy is None: + stack = cls.maybe_dense_stack + elif lazy: + stack = cls.lazy_stack + else: + stack = torch.stack + if batch_size is not None: + if isinstance(batch_size, int): + batch_size = torch.Size([batch_size]) + if batch_size[0] != len(input): + raise ValueError( + f"The provided batch size ({batch_size}) does not match the length of the list ({len(input)})." + ) + bsz = batch_size[1:] + else: + bsz = None + if batch_dims is not None: + batch_dims -= 1 + if names is not None: + names = names[1:] + if cls is TensorDictBase: + from tensordict import TensorDict + + cls = TensorDict + input = [ + ( + cls.from_dict( + d, + auto_batch_size=auto_batch_size, + batch_size=bsz, + batch_dims=batch_dims, + device=device, + names=names, + ) + if not is_tensor_collection(d) + else d + ) + for d in input + ] + return stack(input) + + @classmethod + @abc.abstractmethod + def from_dict( + cls, + input_dict, + *, + auto_batch_size: bool | None = None, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: List[str] | None = None, + ): + """Returns a TensorDict created from a dictionary or another :class:`~.tensordict.TensorDict`. + + If ``batch_size`` is not specified, returns the maximum batch size possible. + + This function works on nested dictionaries too, or can be used to determine the + batch-size of a nested tensordict. + + Args: + input_dict (dictionary, optional): a dictionary to use as a data source + (nested keys compatible). + + Keyword Args: + auto_batch_size (bool, optional): if ``True``, the batch size will be computed automatically. + Defaults to ``False``. + batch_size (iterable of int, optional): a batch size for the tensordict. + device (torch.device or compatible type, optional): a device for the TensorDict. + batch_dims (int, optional): the ``batch_dims`` (ie number of leading dimensions + to be considered for ``batch_size``). Exclusinve with ``batch_size``. + Note that this is the __maximum__ number of batch dims of the tensordict, + a smaller number is tolerated. + names (list of str, optional): the dimension names of the tensordict. + + Examples: + >>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)} + >>> print(TensorDict.from_dict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # nested dict: the nested TensorDict can have a different batch-size + >>> # as long as its leading dims match. + >>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}} + >>> print(TensorDict.from_dict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # we can also use this to work out the batch sie of a tensordict + >>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, []) + >>> print(TensorDict.from_dict(input_td)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + """ + raise NotImplementedError + + @classmethod + def _from_dict_validated(cls, *args, **kwargs): + """A faster version of from_dict when the values have been validated. + + By default, falls back on :meth:`~.from_dict`. + """ + return cls.from_dict(*args, **kwargs) + + @abc.abstractmethod + def from_dict_instance( + self, + input_dict, + *others, + auto_batch_size: bool | None = None, + batch_size=None, + device=None, + batch_dims=None, + names: List[str] | None = None, + ): + """Instance method version of :meth:`~tensordict.TensorDict.from_dict`. + + Unlike :meth:`~tensordict.TensorDict.from_dict`, this method will + attempt to keep the tensordict types within the existing tree (for + any existing leaf). + + Examples: + >>> from tensordict import TensorDict, tensorclass + >>> import torch + >>> + >>> @tensorclass + >>> class MyClass: + ... x: torch.Tensor + ... y: int + >>> + >>> td = TensorDict({"a": torch.randn(()), "b": MyClass(x=torch.zeros(()), y=1)}) + >>> print(td.from_dict_instance(td.to_dict())) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: MyClass( + x=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + y=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(td.from_dict(td.to_dict())) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + raise NotImplementedError + + @classmethod + def from_pytree( + cls, + pytree, + *, + batch_size: torch.Size | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + ): + """Converts a pytree to a TensorDict instance. + + This method is designed to keep the pytree nested structure as much as possible. + + Additional non-tensor keys are added to keep track of each level's identity, providing + a built-in pytree-to-tensordict bijective transform API. + + Accepted classes currently include lists, tuples, named tuples and dict. + + .. note:: + For dictionaries, non-NestedKey keys are registered separately as :class:`~tensordict.NonTensorData` + instances. + + .. note:: + Tensor-castable types (such as int, float or np.ndarray) will be converted to torch.Tensor instances. + Note that this transformation is surjective: transforming back the tensordict to a pytree will not + recover the original types. + + Examples: + >>> # Create a pytree with tensor leaves, and one "weird"-looking dict key + >>> class WeirdLookingClass: + ... pass + ... + >>> weird_key = WeirdLookingClass() + >>> # Make a pytree with tuple, lists, dict and namedtuple + >>> pytree = ( + ... [torch.randint(10, (3,)), torch.zeros(2)], + ... { + ... "tensor": torch.randn( + ... 2, + ... ), + ... "td": TensorDict({"one": 1}), + ... weird_key: torch.randint(10, (2,)), + ... "list": [1, 2, 3], + ... }, + ... {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()}, + ... ) + >>> # Build a TensorDict from that pytree + >>> td = TensorDict.from_pytree(pytree) + >>> # Recover the pytree + >>> pytree_recon = td.to_pytree() + >>> # Check that the leaves match + >>> def check(v1, v2): + >>> assert (v1 == v2).all() + >>> + >>> torch.utils._pytree.tree_map(check, pytree, pytree_recon) + >>> assert weird_key in pytree_recon[1] + + """ + if is_tensor_collection(pytree): + return pytree + if isinstance(pytree, (torch.Tensor,)): + return pytree + + from tensordict._td import TensorDict + + result = None + if is_namedtuple(pytree): + result = TensorDict.from_namedtuple(named_tuple=pytree) + if batch_dims is not None: + result.batch_size = batch_size + result["_pytree_type"] = type(pytree) + elif isinstance(pytree, (list, tuple)): + source = {str(i): cls.from_pytree(elt) for i, elt in enumerate(pytree)} + source["_pytree_type"] = type(pytree) + result = TensorDict(source, batch_size=batch_size) + elif isinstance(pytree, dict): + source = {} + for key, item in pytree.items(): + if isinstance(key, NestedKey): + source[key] = cls.from_pytree(item) + else: + subs_key = "" + str(uuid.uuid1()) + source[subs_key] = TensorDict( + {"value": cls.from_pytree(item), "key": key} + ) + source["_pytree_type"] = type(pytree) + result = TensorDict(source, batch_size=batch_size) + if result is not None: + if auto_batch_size: + result.auto_batch_size_(batch_dims) + return result + if isinstance(pytree, (int, float, np.ndarray)): + return torch.as_tensor(pytree) + raise NotImplementedError(f"Unknown type {type(pytree)}.") + + def to_pytree(self): + """Converts a tensordict to a PyTree. + + If the tensordict was not created from a pytree, this method just returns ``self`` without modification. + + See :meth:`~.from_pytree` for more information and examples. + + """ + _pytree_type = self._get_str("_pytree_type", default=None) + if _pytree_type is None: + return self + _pytree_type = _pytree_type.data + items = {key: val for (key, val) in self.items() if key != "_pytree_type"} + items = { + key: val if not is_tensor_collection(val) else val.to_pytree() + for key, val in items.items() + } + if _pytree_type in (list, tuple): + return _pytree_type((items[str(i)] for i in range(len(items)))) + if _pytree_type is dict: + items = dict( + ( + ( + (val["key"], val["value"]) + if key.startswith("") + else (key, val) + ) + for (key, val) in items.items() + ) + ) + return items + if is_namedtuple_class(_pytree_type): + from tensordict._td import TensorDict + + return TensorDict(items).to_namedtuple(dest_cls=_pytree_type) + raise NotImplementedError(f"unknown type {_pytree_type}") + + @classmethod + def from_h5( + cls, + filename, + *, + mode: str = "r", + auto_batch_size: bool = False, + batch_dims: int | None = None, + batch_size: torch.Size | None = None, + ): + """Creates a PersistentTensorDict from a h5 file. + + Args: + filename (str): The path to the h5 file. + + Keyword Arguments: + mode (str, optional): Reading mode. Defaults to ``"r"``. + auto_batch_size (bool, optional): If ``True``, the batch size will be computed automatically. + Defaults to ``False``. + batch_dims (int, optional): If auto_batch_size is ``True``, defines how many dimensions the output + tensordict should have. Defaults to ``None`` (full batch-size at each level). + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to ``None``. + + Returns: + A PersistentTensorDict representation of the input h5 file. + + Examples: + >>> td = TensorDict.from_h5("path/to/file.h5") + >>> print(td) + PersistentTensorDict( + fields={ + key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + from tensordict.persistent import PersistentTensorDict + + result = PersistentTensorDict.from_h5( + filename, mode=mode, batch_size=batch_size + ) + if auto_batch_size: + if batch_size is not None: + raise TypeError(cls._CONFLICTING_BATCH_SIZES.format("from_h5")) + result.auto_batch_size_(batch_dims=batch_dims) + return result + + # Module interaction + @classmethod + def from_module( + cls, + module, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + ): + """Copies the params and buffers of a module in a tensordict. + + Args: + module (nn.Module): the module to get the parameters from. + as_module (bool, optional): if ``True``, a :class:`~tensordict.nn.TensorDictParams` + instance will be returned which can be used to store parameters + within a :class:`torch.nn.Module`. Defaults to ``False``. + lock (bool, optional): if ``True``, the resulting tensordict will be locked. + Defaults to ``True``. + use_state_dict (bool, optional): if ``True``, the state-dict from the + module will be used and unflattened into a TensorDict with + the tree structure of the model. Defaults to ``False``. + + .. note:: + This is particularly useful when state-dict hooks have to be used. + + Examples: + >>> from torch import nn + >>> module = nn.TransformerDecoder( + ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), + ... num_layers=1 + ... ) + >>> params = TensorDict.from_module(module) + >>> print(params["layers", "0", "linear1"]) + TensorDict( + fields={ + bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + ... + + @classmethod + def from_modules( + cls, + *modules, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + lazy_stack: bool = False, + expand_identical: bool = False, + ): + """Retrieves the parameters of several modules for ensebmle learning/feature of expects applications through vmap. + + Args: + modules (sequence of nn.Module): the modules to get the parameters from. + If the modules differ in their structure, a lazy stack is needed + (see the ``lazy_stack`` argument below). + + Keyword Args: + as_module (bool, optional): if ``True``, a :class:`~tensordict.nn.TensorDictParams` + instance will be returned which can be used to store parameters + within a :class:`torch.nn.Module`. Defaults to ``False``. + lock (bool, optional): if ``True``, the resulting tensordict will be locked. + Defaults to ``True``. + use_state_dict (bool, optional): if ``True``, the state-dict from the + module will be used and unflattened into a TensorDict with + the tree structure of the model. Defaults to ``False``. + + .. note:: + This is particularly useful when state-dict hooks have to be used. + + lazy_stack (bool, optional): whether parameters should be densly or + lazily stacked. Defaults to ``False`` (dense stack). + + .. note:: + ``lazy_stack`` and ``as_module`` are exclusive features. + + .. warning:: + There is a crucial difference between lazy and non-lazy outputs + in that non-lazy output will reinstantiate parameters with the + desired batch-size, while ``lazy_stack`` will just represent + the parameters as lazily stacked. This means that whilst the + original parameters can safely be passed to an optimizer + when ``lazy_stack=True``, the new parameters need to be passed + when it is set to ``True``. + + .. warning:: + Whilst it can be tempting to use a lazy stack to keep the + orignal parameter references, remember that lazy stack + perform a stack each time :meth:`~.get` is called. This will + require memory (N times the size of the parameters, more if a + graph is built) and time to be computed. + It also means that the optimizer(s) will contain more + parameters, and operations like :meth:`~torch.optim.Optimizer.step` + or :meth:`~torch.optim.Optimizer.zero_grad` will take longer + to be executed. In general, ``lazy_stack`` should be reserved + to very few use cases. + + expand_identical (bool, optional): if ``True`` and the same parameter (same + identity) is being stacked to itself, an expanded version of this parameter + will be returned instead. This argument is ignored when ``lazy_stack=True``. + + Examples: + >>> from torch import nn + >>> from tensordict import TensorDict + >>> torch.manual_seed(0) + >>> empty_module = nn.Linear(3, 4, device="meta") + >>> n_models = 2 + >>> modules = [nn.Linear(3, 4) for _ in range(n_models)] + >>> params = TensorDict.from_modules(*modules) + >>> print(params) + TensorDict( + fields={ + bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + >>> # example of batch execution + >>> def exec_module(params, x): + ... with params.to_module(empty_module): + ... return empty_module(x) + >>> x = torch.randn(3) + >>> y = torch.vmap(exec_module, (0, None))(params, x) + >>> assert y.shape == (n_models, 4) + >>> # since lazy_stack = False, backprop leaves the original params untouched + >>> y.sum().backward() + >>> assert params["weight"].grad.norm() > 0 + >>> assert modules[0].weight.grad is None + + With ``lazy_stack=True``, things are slightly different: + + >>> params = TensorDict.from_modules(*modules, lazy_stack=True) + >>> print(params) + LazyStackedTensorDict( + fields={ + bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + exclusive_fields={ + }, + batch_size=torch.Size([2]), + device=None, + is_shared=False, + stack_dim=0) + >>> # example of batch execution + >>> y = torch.vmap(exec_module, (0, None))(params, x) + >>> assert y.shape == (n_models, 4) + >>> y.sum().backward() + >>> assert modules[0].weight.grad is not None + + + """ + param_list = [ + cls.from_module(module, use_state_dict=use_state_dict) for module in modules + ] + if lazy_stack: + from tensordict._lazy import LazyStackedTensorDict + + for param in param_list: + if any( + isinstance(tensor, UninitializedTensorMixin) + for tensor in param.values(True, True) + ): + raise RuntimeError( + "lasy_stack=True is not compatible with lazy modules." + ) + params = LazyStackedTensorDict.lazy_stack(param_list) + elif expand_identical: + from tensordict._torch_func import _stack_uninit_params + + # Check the keys + # If not expand_identical, `stack` takes care of that check but + # here we use apply which will ignore keys that are in one TD but not another + sets = [set(param.keys(True, True)) for param in param_list] + for set_ in sets[1:]: + if set_ != sets[0]: + raise ValueError( + f"All key sets must match. " + f"Got {set_.symmetric_difference(sets[0])} in one but not another." + ) + + def maybe_stack(*params): + param = params[0] + if isinstance(param, UninitializedTensorMixin): + return _stack_uninit_params(params, 0) + if len(set(params)) == 1: + return param.expand((len(params), *param.shape)) + result = torch.stack(params) + if isinstance(param, nn.Parameter): + return nn.Parameter(result.detach(), param.requires_grad) + return Buffer(result) + + params = param_list[0]._fast_apply( + maybe_stack, + *param_list[1:], + batch_size=torch.Size([len(param_list), *param_list[0].batch_size]), + ) + else: + with set_lazy_legacy(False), torch.no_grad(): + params = torch.stack(param_list) + + # Make sure params are params, buffers are buffers + def make_param(param, orig_param): + if isinstance(param, UninitializedTensorMixin): + return param + if isinstance(orig_param, nn.Parameter): + return nn.Parameter(param.detach(), orig_param.requires_grad) + return Buffer(param) + + params = params._fast_apply(make_param, param_list[0], propagate_lock=True) + if as_module: + from tensordict.nn import TensorDictParams + + params = TensorDictParams(params, no_convert=True) + if lock: + params.lock_() + return params + + @_as_context_manager() + def to_module( + self, + module: nn.Module, + *, + inplace: bool | None = None, + return_swap: bool = True, + swap_dest=None, + use_state_dict: bool = False, + non_blocking: bool = False, + memo=None, # deprecated + ): + """Writes the content of a TensorDictBase instance onto a given nn.Module attributes, recursively. + + ``to_module`` can also be used a context manager to temporarily populate a module with a collection of + parameters/buffers (see example below). + + Args: + module (nn.Module): a module to write the parameters into. + + Keyword Args: + inplace (bool, optional): if ``True``, the parameters or tensors + in the module are updated in-place. Defaults to ``False``. + return_swap (bool, optional): if ``True``, the old parameter configuration + will be returned. Defaults to ``False``. + swap_dest (TensorDictBase, optional): if ``return_swap`` is ``True``, + the tensordict where the swap should be written. + use_state_dict (bool, optional): if ``True``, state-dict API will be + used to load the parameters (including the state-dict hooks). + Defaults to ``False``. + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Examples: + >>> from torch import nn + >>> module = nn.TransformerDecoder( + ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), + ... num_layers=1) + >>> params = TensorDict.from_module(module) + >>> params.data.zero_() + >>> params.to_module(module) + >>> assert (module.layers[0].linear1.weight == 0).all() + + Using a tensordict as a context manager can be useful to make functional calls: + Examples: + >>> from tensordict import from_module + >>> module = nn.TransformerDecoder( + ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), + ... num_layers=1) + >>> params = TensorDict.from_module(module) + >>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances + >>> with params.to_module(module): + ... # Call the module with zeroed params + ... y = module(*inputs) + >>> # The module is repopulated with its original params + >>> assert (TensorDict.from_module(module) != 0).any() + + Returns: + A tensordict containing the values from the module if ``return_swap`` is ``True``, ``None`` otherwise. + + """ + if memo is not None: + raise RuntimeError("memo cannot be passed to the public to_module anymore.") + hooks = getattr( + torch.nn.modules.module, "_global_parameter_registration_hooks", {} + ) + memo = {"hooks": tuple(hooks.values())} + return self._to_module( + module=module, + inplace=inplace, + return_swap=return_swap, + swap_dest=swap_dest, + memo=memo, + use_state_dict=use_state_dict, + non_blocking=non_blocking, + ) + + @abc.abstractmethod + def _to_module( + self, + module, + *, + inplace: bool | None = None, + return_swap: bool = True, + swap_dest=None, + memo=None, + use_state_dict: bool = False, + non_blocking: bool = False, + ): + raise NotImplementedError + + # Shape functionality + @property + def shape(self) -> torch.Size: + """See :obj:`~tensordict.TensorDictBase.batch_size`.""" + return self.batch_size + + @shape.setter + def shape(self, value): + self.batch_size = value + + @property + @abc.abstractmethod + def batch_size(self) -> torch.Size: + """Shape (or batch_size) of a TensorDict. + + The shape of a tensordict corresponds to the common first ``N`` + dimensions of the tensors it contains, where ``N`` is an arbitrary + number. The batch-size contrasts with the "feature size" which repesents + the semantically relevant shapes of a tensor. For instance, a batch of videos + may have shape ``[B, T, C, W, H]``, where ``[B, T]`` is the batch-size (batch and + time dimensions) and ``[C, W, H]`` are the feature dimensions (channels and spacial + dimensions). + + The ``TensorDict`` shape is controlled by the user upon + initialization (ie, it is not inferred from the tensor shapes). + + The ``batch_size`` can be edited dynamically if the new size is compatible + with the TensorDict content. For instance, setting the batch size to + an empty value is always allowed. + + Returns: + a :obj:`~torch.Size` object describing the TensorDict batch size. + + Examples: + >>> data = TensorDict({ + ... "key 0": torch.randn(3, 4), + ... "key 1": torch.randn(3, 5), + ... "nested": TensorDict({"key 0": torch.randn(3, 4)}, batch_size=[3, 4])}, + ... batch_size=[3]) + >>> data.batch_size = () # resets the batch-size to an empty value + """ + raise NotImplementedError + + def size(self, dim: int | None = None) -> torch.Size | int: + """Returns the size of the dimension indicated by ``dim``. + + If ``dim`` is not specified, returns the ``batch_size`` attribute of the TensorDict. + + """ + if dim is None: + return self.batch_size + return self.batch_size[dim] + + @property + def data(self) -> Self: + """Returns a tensordict containing the .data attributes of the leaf tensors.""" + return self._data() + + @data.setter + def data(self, value: Self): + self._data_setter(value) + + @property + def grad(self) -> Self: + """Returns a tensordict containing the .grad attributes of the leaf tensors.""" + return self._grad() + + def data_ptr(self, *, storage: bool = False): + """Returns the data_ptr of the tensordict leaves. + + This can be useful to check if two tensordicts share the same ``data_ptr()``. + + Keyword Args: + storage (bool, optional): if ``True``, `tensor.untyped_storage().data_ptr()` will be called + instead. Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) + >>> assert (td0.data_ptr() == td.data_ptr()).all() + + .. note:: :class:`~tensordict.LazyStackedTensorDict` instances will be displayed as nested tensordicts to + reflect the true ``data_ptr()`` of their leaves: + + >>> td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) + >>> td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) + >>> td = TensorDict.lazy_stack([td0, td1]) + >>> td.data_ptr() + TensorDict( + fields={ + 0: TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False), + 1: TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False) + + """ + if storage: + + def func(x): + return x.untyped_storage().data_ptr() + + else: + + def func(x): + return x.data_ptr() + + from tensordict import TensorDict + + return TensorDict( + { + key: func(val) + for key, val in self.items(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + }, + device=torch.device("cpu"), + ) + + @grad.setter + def grad(self, grad): + def set_grad(x, grad): + if x.grad is None: + x.grad = grad + else: + x.grad.copy_(grad) + + self._fast_apply(set_grad, grad) + + def zero_grad(self, set_to_none: bool = True) -> Self: + """Zeros all the gradients of the TensorDict recursively. + + Args: + set_to_none (bool, optional): if ``True``, tensor.grad will be ``None``, + otherwise ``0``. + Defaults to ``True``. + + """ + if set_to_none: + for val in self._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + val.grad = None + return self + for val in self._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + val.grad.zero_() + return self + + @cache # noqa + def _dtype(self): + dtype = None + for val in self.values(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + val_dtype = getattr(val, "dtype", None) + if dtype is None and val_dtype is not None: + dtype = val_dtype + elif dtype is not None and val_dtype is not None and dtype != val_dtype: + return None + return dtype + + @property + def dtype(self): + """Returns the dtype of the values in the tensordict, if it is unique.""" + return self._dtype() + + def _batch_size_setter(self, new_batch_size: torch.Size) -> None: + if new_batch_size == self.batch_size: + return + if self._lazy: + raise RuntimeError( + f"Received a new batch size {new_batch_size} with an existing batch_size {self.batch_size} in a lazy TD. " + "Modifying the batch size of a lazy representation of a " + "tensordict is not permitted. Consider instantiating the " + "tensordict first by calling `td = td.to_tensordict()` before " + "resetting the batch size." + ) + if not isinstance(new_batch_size, torch.Size): + new_batch_size = torch.Size(new_batch_size) + for key, value in self.items(): + if _is_tensor_collection(type(value)): + if len(value.batch_size) < len(new_batch_size): + # document as edge case + value.batch_size = new_batch_size + self._set_str( + key, value, inplace=True, validated=True, non_blocking=False + ) + self._check_new_batch_size(new_batch_size) + has_names = self._has_names() + if has_names: + # if the tensordict has dim names and the new batch-size has more dims, + # we can simply add empty names after the current ones. + # Otherwise, we discard the extra existing names. + names = self.names + self._erase_names() + self._change_batch_size(new_batch_size) + if has_names: + # if the tensordict has dim names and the new batch-size has more dims, + # we can simply add empty names after the current ones. + # Otherwise, we discard the extra existing names. + if len(names) < len(new_batch_size): + self._set_names(names + [None] * (len(new_batch_size) - len(names))) + else: + self._set_names(names[: self.batch_dims]) + + @abc.abstractmethod + def _set_names(self, names: Sequence[str] | None) -> None: ... + + @property + def batch_dims(self) -> int: + """Length of the tensordict batch size. + + Returns: + int describing the number of dimensions of the tensordict. + + """ + return len(self.batch_size) + + def ndimension(self) -> int: + """See :meth:`~.batch_dims`.""" + return self.batch_dims + + @property + def ndim(self) -> int: + """See :meth:`~.batch_dims`.""" + return self.batch_dims + + def dim(self) -> int: + """See :meth:`~.batch_dims`.""" + return self.batch_dims + + def numel(self) -> int: + """Total number of elements in the batch. + + Lower-bounded to 1, as a stack of two tensordict with empty shape will + have two elements, therefore we consider that a tensordict is at least + 1-element big. + """ + return max(1, self.batch_size.numel()) + + @property + def depth(self) -> int: + """Returns the depth - maximum number of levels - of a tensordict. + + The minimum depth is 0 (no nested tensordict). + """ + return self._depth() + + @cache # noqa: B019 + def _depth(self): + depth = 0 + for key in self.keys(True, True, is_leaf=_is_leaf_nontensor): + if isinstance(key, tuple): + depth = max(depth, len(key) - 1) + return depth + + @overload + def expand(self, *shape: int) -> Self: ... + + @overload + def expand(self, shape: torch.Size) -> Self: ... + + @abc.abstractmethod + def expand(self, *args: int | torch.Size) -> Self: + """Expands each tensor of the tensordict according to the :func:`~torch.expand` function, ignoring the feature dimensions. + + Supports iterables to specify the shape. + + Examples: + >>> td = TensorDict({ + ... 'a': torch.zeros(3, 4, 5), + ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) + >>> td_expand = td.expand(10, 3, 4) + >>> assert td_expand.shape == torch.Size([10, 3, 4]) + >>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5]) + + """ + raise NotImplementedError + + def expand_as(self, other: TensorCollection | torch.Tensor) -> Self: + """Broadcasts the shape of the tensordict to the shape of `other` and expands it accordingly. + + If the input is a tensor collection (tensordict or tensorclass), + the leaves will be expanded on a one-to-one basis. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> td0 = TensorDict({ + ... "a": torch.ones(3, 1, 4), + ... "b": {"c": torch.ones(3, 2, 1, 4)}}, + ... batch_size=[3], + ... ) + >>> td1 = TensorDict({ + ... "a": torch.zeros(2, 3, 5, 4), + ... "b": {"c": torch.zeros(2, 3, 2, 6, 4)}}, + ... batch_size=[2, 3], + ... ) + >>> expanded = td0.expand_as(td1) + >>> assert (expanded==1).all() + >>> print(expanded) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([2, 3, 5, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([2, 3, 2, 6, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2, 3]), + device=None, + is_shared=False)}, + batch_size=torch.Size([2, 3]), + device=None, + is_shared=False) + + """ + if _is_tensor_collection(type(other)): + + def expand_as(x, y): + return x.expand_as(y) + + return self.apply(expand_as, other, batch_size=other.batch_size) + return self.expand(other.shape) + + def new_zeros( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): # noqa: D417 + """Returns a TensorDict of size ``size`` filled with 0. + + By default, the returned TensorDict has the same ``torch.dtype`` and ``torch.device`` as this tensordict. + + Args: + size (int...): a list, tuple, or torch.Size of integers defining the shape of the output tensor. + + Keyword Args: + dtype (torch.dtype, optional): the desired type of returned tensordict. + Default: if ``None``, the `torch.dtype` will be unchanged. + device (torch.device, optional): the desired device of returned tensordict. + Default: if ``None``, the ``torch.device`` will be unchanged. + requires_grad (bool, optional): If autograd should record operations on the + returned tensors. Default: ``False``. + layout (torch.layout, optional): the desired layout of returned TensorDict values. + Default: ``torch.strided``. + pin_memory (bool, optional): If set, returned tensor would be allocated in the + pinned memory. Works only for CPU tensors. Default: ``False``. + empty_lazy (bool, optional): If `True`, lazy stacks will be emptied of their content. + This can be useful whenever the content of a lazy stack is likely to change + during filling of the new tensordict. This argument is propagated to sub-tensordicts. + Defaults to ``False``. + + """ + kwargs = {} + if pin_memory is not None: + kwargs = {"pin_memory": pin_memory} + return self._new_impl( + size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + layout=layout, + funcname="new_zeros", + empty_lazy=empty_lazy, + **kwargs, + ) + + def new_ones( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): # noqa: D417 + """Returns a TensorDict of size ``size`` filled with 1. + + By default, the returned TensorDict has the same ``torch.dtype`` and ``torch.device`` as this tensordict. + + Args: + size (int...): a list, tuple, or torch.Size of integers defining the shape of the output tensor. + + Keyword Args: + dtype (torch.dtype, optional): the desired type of returned tensordict. + Default: if ``None``, the `torch.dtype` will be unchanged. + device (torch.device, optional): the desired device of returned tensordict. + Default: if ``None``, the ``torch.device`` will be unchanged. + requires_grad (bool, optional): If autograd should record operations on the + returned tensors. Default: ``False``. + layout (torch.layout, optional): the desired layout of returned TensorDict values. + Default: ``torch.strided``. + pin_memory (bool, optional): If set, returned tensor would be allocated in the + pinned memory. Works only for CPU tensors. Default: ``False``. + empty_lazy (bool, optional): If `True`, lazy stacks will be emptied of their content. + This can be useful whenever the content of a lazy stack is likely to change + during filling of the new tensordict. This argument is propagated to sub-tensordicts. + Defaults to ``False``. + + """ + kwargs = {} + if pin_memory is not None: + kwargs = {"pin_memory": pin_memory} + return self._new_impl( + size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + layout=layout, + funcname="new_ones", + empty_lazy=empty_lazy, + **kwargs, + ) + + def new_empty( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): # noqa: D417 + """Returns a TensorDict of size ``size`` with emtpy tensors. + + By default, the returned TensorDict has the same ``torch.dtype`` and ``torch.device`` as this tensordict. + + Args: + size (int...): a list, tuple, or torch.Size of integers defining the shape of the output tensor. + + Keyword Args: + dtype (torch.dtype, optional): the desired type of returned tensordict. + Default: if ``None``, the `torch.dtype` will be unchanged. + device (torch.device, optional): the desired device of returned tensordict. + Default: if ``None``, the ``torch.device`` will be unchanged. + requires_grad (bool, optional): If autograd should record operations on the + returned tensors. Default: ``False``. + layout (torch.layout, optional): the desired layout of returned TensorDict values. + Default: ``torch.strided``. + pin_memory (bool, optional): If set, returned tensor would be allocated in the + pinned memory. Works only for CPU tensors. Default: ``False``. + empty_lazy (bool, optional): If `True`, lazy stacks will be emptied of their content. + This can be useful whenever the content of a lazy stack is likely to change + during filling of the new tensordict. This argument is propagated to sub-tensordicts. + Defaults to ``False``. + + """ + kwargs = {} + if pin_memory is not None: + kwargs = {"pin_memory": pin_memory} + return self._new_impl( + size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + layout=layout, + funcname="new_empty", + empty_lazy=empty_lazy, + **kwargs, + ) + + def new_full( + self, + size: torch.Size, + fill_value, + *, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + pin_memory: bool | None = None, + empty_lazy: bool = False, + ): # noqa: D417 + """Returns a TensorDict of size ``size`` filled with 1. + + By default, the returned TensorDict has the same ``torch.dtype`` and ``torch.device`` as this tensordict. + + Args: + size (sequence of int): a list, tuple, or torch.Size of integers defining the shape of the output tensor. + fill_value (scalar): the number to fill the output tensor with. + + Keyword Args: + dtype (torch.dtype, optional): the desired type of returned tensordict. + Default: if ``None``, the `torch.dtype` will be unchanged. + device (torch.device, optional): the desired device of returned tensordict. + Default: if ``None``, the ``torch.device`` will be unchanged. + requires_grad (bool, optional): If autograd should record operations on the + returned tensors. Default: ``False``. + layout (torch.layout, optional): the desired layout of returned TensorDict values. + Default: ``torch.strided``. + pin_memory (bool, optional): If set, returned tensor would be allocated in the + pinned memory. Works only for CPU tensors. Default: ``False``. + empty_lazy (bool, optional): If `True`, lazy stacks will be emptied of their content. + This can be useful whenever the content of a lazy stack is likely to change + during filling of the new tensordict. This argument is propagated to sub-tensordicts. + Defaults to ``False``. + + """ + kwargs = {} + if pin_memory is not None: + kwargs = {"pin_memory": pin_memory} + return self._new_impl( + size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + layout=layout, + funcname="new_full", + fill_value=fill_value, + empty_lazy=empty_lazy, + **kwargs, + ) + + def _new_impl( + self, + size: torch.Size, + *, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + layout: torch.layout = torch.strided, + funcname: str, + empty_lazy: bool = False, + **kwargs, + ): + if isinstance(size, int): + size = (size,) + elif len(size) == 1 and not isinstance(size[0], int): + size = size[0] + + ndim = self.ndim + if device is not NO_DEFAULT: + kwargs["device"] = device + + def func(tensor, size=size): + feature_shape = tensor.shape[ndim:] + size = torch.Size((*size, *feature_shape)) + kwargs_copy = kwargs + if empty_lazy and is_tensor_collection(tensor): + kwargs_copy = dict(kwargs) + kwargs_copy["empty_lazy"] = empty_lazy + return getattr(tensor, funcname)( + size, + dtype=dtype, + requires_grad=requires_grad, + layout=layout, + **kwargs_copy, + ) + + names = self._maybe_names() + if names: + if len(size) > self.ndim: + names = [None] * (len(size) - self.ndim) + list(names) + elif self.ndim > len(size): + names = names[-len(size) :] + return self._fast_apply( + func, + call_on_nested=True, + device=device, + batch_size=size, + names=names, + ) + + def new_tensor( + self, + data: torch.Tensor | TensorCollection, + *, + dtype: torch.dtype = None, + device: DeviceType = NO_DEFAULT, + requires_grad: bool = False, + pin_memory: bool | None = None, + ) -> Self: # noqa: D417 + """Returns a new TensorDict with data as the tensor ``data``. + + By default, the returned TensorDict values have the same ``torch.dtype`` and ``torch.device`` as this tensor. + + The ``data`` can also be a tensor collection (``TensorDict`` or ``tensorclass``), in which case + the ``new_tensor`` method iterates over the tensor pairs of ``self`` and ``data``. + + Args: + data (torch.Tensor or TensorDictBase): the data to be copied. + + Keyword Args: + dtype (torch.dtype, optional): the desired type of returned tensordict. + Default: if ``None``, the `torch.dtype` will be unchanged. + device (torch.device, optional): the desired device of returned tensordict. + Default: if ``None``, the ``torch.device`` will be unchanged. + requires_grad (bool, optional): If autograd should record operations on the + returned tensors. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in the + pinned memory. Works only for CPU tensors. Default: ``False``. + + """ + kwargs = {} + if device is not NO_DEFAULT: + kwargs["device"] = device + else: + device = kwargs["device"] = data.device + if pin_memory is not None: + kwargs["pin_memory"] = pin_memory + + def func(x, tensor_b=None): + if tensor_b is None: + tensor_b = data + return x.new_tensor( + tensor_b, + dtype=dtype, + requires_grad=requires_grad, + **kwargs, + ) + + if _is_tensor_collection(type(data)): + return self._fast_apply( + func, + data, + call_on_nested=True, + device=device, + batch_size=data.shape, + ) + return self._fast_apply( + func, + call_on_nested=True, + device=device, + batch_size=data.shape, + ) + + def unbind(self, dim: int) -> tuple[T, ...]: + """Returns a tuple of indexed tensordicts, unbound along the indicated dimension. + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(12).reshape(3, 4), + ... }, batch_size=[3, 4]) + >>> td0, td1, td2 = td.unbind(0) + >>> td0['x'] + tensor([0, 1, 2, 3]) + >>> td1['x'] + tensor([4, 5, 6, 7]) + + """ + dim = _maybe_correct_neg_dim(dim, self.batch_size) + results = self._unbind(dim) + if self._is_memmap or self._is_shared: + for result in results: + result.lock_() + return results + + @abc.abstractmethod + def _unbind(self, dim: int) -> tuple[T, ...]: + raise NotImplementedError + + def tensor_split( + self, + indices_or_sections: int | list[int] | tuple[int, ...] | torch.Tensor, + dim=0, + ) -> tuple[TensorDictBase, ...]: + """Splits a TensorDict into multiple sub-tensordicts, all of which are views of input, along dimension dim according to the indices or number of sections specified by indices_or_sections. + + Args: + indices_or_sections (int or List(int) or tuple(int) or 1D tensor of ints): If `indices_or_sections` is an integer + `n` or a zero dimensional long tensordict with value `n`, input is split into `n` sections along dimension `dim`. + If input is divisible by `n` along dimension `dim`, each section will be of equal size, `input.size(dim) / n`. + If input is not divisible by `n`, the sizes of the first `int(input.size(dim) % n)` sections will have + size `int(input.size(dim) / n) + 1`, and the rest will have size `int(input.size(dim) / n)`. + If `indices_or_sections` is a list or tuple of ints, or a one-dimensional long tensor, then input is split + along dimension `dim` at each of the indices in the list, tuple or tensor. + For instance, `indices_or_sections=[2, 3]` and `dim=0` would result in the tensors `input[:2]`, `input[2:3]`, and `input[3:]`. + If `indices_or_sections` is a tensor, it must be a zero-dimensional or one-dimensional long tensor on the CPU. + dim (int, optional): dimension along which to split the tensor. Default: 0 + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 4, 2), + ... }, batch_size=[3, 4]) + >>> td0, td1 = td.tensor_split(dim=-1, indices_or_sections=2) + >>> td0['x'] + tensor([[[ 0, 1], + [ 2, 3]], + [[ 8, 9], + [10, 11]], + [[16, 17], + [18, 19]]]) + + """ + if not isinstance(indices_or_sections, (int, list, tuple, torch.Tensor)): + raise ValueError( + "indices_or_sections must be an integer, a list of integers, or a 1D tensor of integers" + ) + + batch_size = self.batch_size + dim = _maybe_correct_neg_dim(dim, batch_size) + + if self.ndim == 0: + msg = "tensor_split: received a rank zero tensor, but expected a tensor of rank one or greater!" + raise ValueError(msg) + + # Case 0 -- indices_or_sections is an integer or a scalar tensor n and a is split along dim into n parts of equal-ish length + if isinstance(indices_or_sections, int): + sections: int = indices_or_sections # type: ignore[assignment] + + if sections <= 0: + msg = f"tensor_split: number of sections must be greater than 0, but was {sections}" + raise ValueError(msg) + + dim_size = self.shape[dim] + min_split_size = math.floor(dim_size / sections) + num_splits_one_extra = dim_size % sections + + split_sizes = [] + for split_idx in range(sections): + split_size = ( + min_split_size + 1 + if (split_idx < num_splits_one_extra) + else min_split_size + ) + split_sizes.append(split_size) + + return tuple(self.split(split_sizes, dim=dim)) + # Case 1 -- indices_or_sections is a sequence of integers or a 1D tensor describing the splits + else: + indices = indices_or_sections + indices = [0] + list(indices) + [self.shape[dim]] + split_sizes = [indices[i + 1] - indices[i] for i in range(len(indices) - 1)] + return tuple(self.split(split_sizes, dim=dim)) + + @abc.abstractmethod + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorCollection, ...]: + """Splits a tensordict into the specified number of chunks, if possible. + + Each chunk is a view of the input tensordict. + + Args: + chunks (int): number of chunks to return + dim (int, optional): dimension along which to split the + tensordict. Default is 0. + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 4, 2), + ... }, batch_size=[3, 4]) + >>> td0, td1 = td.chunk(dim=-1, chunks=2) + >>> td0['x'] + tensor([[[ 0, 1], + [ 2, 3]], + [[ 8, 9], + [10, 11]], + [[16, 17], + [18, 19]]]) + + """ + raise NotImplementedError + + @overload + def unsqueeze(self, dim: int) -> Self: ... + + @_as_context_manager() + def unsqueeze(self, *args, **kwargs): + """Unsqueezes all tensors for a dimension comprised in between `-td.batch_dims` and `td.batch_dims` and returns them in a new tensordict. + + Args: + dim (int): dimension along which to unsqueeze + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 4, 2), + ... }, batch_size=[3, 4]) + >>> td = td.unsqueeze(-2) + >>> td.shape + torch.Size([3, 1, 4]) + >>> td.get("x").shape + torch.Size([3, 1, 4, 2]) + + This operation can be used as a context manager too. Changes to the original + tensordict will occur out-place, i.e. the content of the original tensors + will not be altered. This also assumes that the tensordict is not locked + (otherwise, unlocking the tensordict is necessary). + + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 4, 2), + ... }, batch_size=[3, 4]) + >>> with td.unsqueeze(-2) as tds: + ... tds.set("y", torch.zeros(3, 1, 4)) + >>> assert td.get("y").shape == [3, 4] + + """ + _lazy_legacy = lazy_legacy() + + if _lazy_legacy: + return self._legacy_unsqueeze(*args, **kwargs) + else: + result = self._unsqueeze(*args, **kwargs) + if result._is_memmap or result._is_shared: + result.lock_() + return result + + @abc.abstractmethod + def _unsqueeze(self, dim: int): + raise NotImplementedError + + def _legacy_unsqueeze(self, dim: int) -> Self: + if dim < 0: + dim = self.batch_dims + dim + 1 + + if (dim > self.batch_dims) or (dim < 0): + raise RuntimeError( + f"unsqueezing is allowed for dims comprised between " + f"`-td.batch_dims` and `td.batch_dims` only. Got " + f"dim={dim} with a batch size of {self.batch_size}." + ) + from tensordict._lazy import _UnsqueezedTensorDict + + return _UnsqueezedTensorDict( + source=self, + custom_op="unsqueeze", + inv_op="squeeze", + custom_op_kwargs={"dim": dim}, + inv_op_kwargs={"dim": dim}, + ) + + @overload + def squeeze(self, dim: int | None = None) -> Self: ... + + @_as_context_manager() + def squeeze(self, *args, **kwargs): + """Squeezes all tensors for a dimension in between `-self.batch_dims+1` and `self.batch_dims-1` and returns them in a new tensordict. + + Args: + dim (int | None): dimension along which to squeeze. If dim is + ``None``, all singleton dimensions will be squeezed. + Defaults to ``None``. + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 1, 4, 2), + ... }, batch_size=[3, 1, 4]) + >>> td = td.squeeze() + >>> td.shape + torch.Size([3, 4]) + >>> td.get("x").shape + torch.Size([3, 4, 2]) + + This operation can be used as a context manager too. Changes to the original + tensordict will occur out-place, i.e. the content of the original tensors + will not be altered. This also assumes that the tensordict is not locked + (otherwise, unlocking the tensordict is necessary). This functionality is + *not* compatible with implicit squeezing. + + >>> td = TensorDict({ + ... 'x': torch.arange(24).reshape(3, 1, 4, 2), + ... }, batch_size=[3, 1, 4]) + >>> with td.squeeze(1) as tds: + ... tds.set("y", torch.zeros(3, 4)) + >>> assert td.get("y").shape == [3, 1, 4] + + """ + _lazy_legacy = lazy_legacy() + + if _lazy_legacy: + return self._legacy_squeeze(*args, **kwargs) + else: + result = self._squeeze(*args, **kwargs) + if result._is_memmap or result._is_shared: + result.lock_() + return result + + @abc.abstractmethod + def _squeeze(self, dim=None): + raise NotImplementedError + + def _legacy_squeeze(self, dim: int | None = None) -> Self: + from tensordict._lazy import _SqueezedTensorDict + + if dim is None: + size = self.size() + if len(self.size()) == 1 or size.count(1) == 0: + return self + first_singleton_dim = size.index(1) + + squeezed_dict = _SqueezedTensorDict( + source=self, + custom_op="squeeze", + inv_op="unsqueeze", + custom_op_kwargs={"dim": first_singleton_dim}, + inv_op_kwargs={"dim": first_singleton_dim}, + ) + return squeezed_dict.squeeze(dim=None) + + if dim < 0: + dim = self.batch_dims + dim + + if self.batch_dims and (dim >= self.batch_dims or dim < 0): + raise RuntimeError( + f"squeezing is allowed for dims comprised between 0 and " + f"td.batch_dims only. Got dim={dim} and batch_size" + f"={self.batch_size}." + ) + + if dim >= self.batch_dims or self.batch_size[dim] != 1: + return self + + return _SqueezedTensorDict( + source=self, + custom_op="squeeze", + inv_op="unsqueeze", + custom_op_kwargs={"dim": dim}, + inv_op_kwargs={"dim": dim}, + ) + + @overload + def reshape(self, *shape: int): ... + + @overload + def reshape(self, shape: list | tuple): ... + + @abc.abstractmethod + def reshape( + self, + *args, + **kwargs, + ) -> Self: + """Returns a contiguous, reshaped tensor of the desired shape. + + Args: + *shape (int): new shape of the resulting tensordict. + + Returns: + A TensorDict with reshaped keys + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(12).reshape(3, 4), + ... }, batch_size=[3, 4]) + >>> td = td.reshape(12) + >>> print(td['x']) + torch.Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + + """ + raise NotImplementedError + + @abc.abstractmethod + def repeat_interleave( + self, + repeats: torch.Tensor | int, + dim: int | None = None, + *, + output_size: int | None = None, + ) -> Self: + """Repeat elements of a TensorDict. + + .. warning:: This is different from :meth:`~torch.Tensor.repeat` but similar to :func:`numpy.repeat`. + + Args: + repeats (torch.Tensor or int): The number of repetitions for each element. `repeats` is broadcast to fit + the shape of the given axis. + dim (int, optional): The dimension along which to repeat values. By default, use the flattened input + array, and return a flat output array. + + Keyword Args: + output_size (int, optional): Total output size for the given axis (e.g. sum of repeats). If given, it + will avoid stream synchronization needed to calculate output shape of the tensordict. + + Returns: + Repeated TensorDict which has the same shape as input, except along the given axis. + + Examples: + >>> import torch + >>> + >>> from tensordict import TensorDict + >>> + >>> td = TensorDict( + ... { + ... "a": torch.randn(3, 4, 5), + ... "b": TensorDict({ + ... "c": torch.randn(3, 4, 10, 1), + ... "a string": "a string!", + ... }, batch_size=[3, 4, 10]) + ... }, batch_size=[3, 4], + ... ) + >>> print(td.repeat_interleave(2, dim=0)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([6, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + a string: NonTensorData(data=a string!, batch_size=torch.Size([6, 4, 10]), device=None), + c: Tensor(shape=torch.Size([6, 4, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([6, 4, 10]), + device=None, + is_shared=False)}, + batch_size=torch.Size([6, 4]), + device=None, + is_shared=False) + + """ + raise NotImplementedError + + @overload + def repeat(self, repeats: torch.Size): ... + + def repeat(self, *repeats: int) -> Self: + """Repeats this tensor along the specified dimensions. + + Unlike :meth:`~.expand()`, this function copies the tensor's data. + + .. warning:: :meth:`~.repeat` behaves differently from :func:`~numpy.repeat`, but is more similar to + :func:`numpy.tile`. For the operator similar to :func:`numpy.repeat`, see :meth:`~tensordict.TensorDictBase.repeat_interleave`. + + Args: + repeat (torch.Size, int..., tuple of int or list of int): The number of times to repeat this tensor along + each dimension. + + Examples: + >>> import torch + >>> + >>> from tensordict import TensorDict + >>> + >>> td = TensorDict( + ... { + ... "a": torch.randn(3, 4, 5), + ... "b": TensorDict({ + ... "c": torch.randn(3, 4, 10, 1), + ... "a string": "a string!", + ... }, batch_size=[3, 4, 10]) + ... }, batch_size=[3, 4], + ... ) + >>> print(td.repeat(1, 2)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 8, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + a string: NonTensorData(data=a string!, batch_size=torch.Size([3, 8, 10]), device=None), + c: Tensor(shape=torch.Size([3, 8, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 8, 10]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 8]), + device=None, + is_shared=False) + + """ + if len(repeats) == 1 and not isinstance(repeats[0], int): + repeats = repeats[0] + if isinstance(repeats, torch.Size): + return self.repeat(*repeats[0]) + if isinstance(repeats, torch.Tensor): + # This will cause cuda to sync, which may not be desirable + return self.repeat(*repeats.tolist()) + raise ValueError( + f"repeats must be a sequence of integers, a tensor or a torch.Size object. Got {type(repeats)} instead." + ) + if len(repeats) != self.ndimension(): + raise ValueError( + f"The number of repeat elements must match the number of dimensions of the tensordict. Got {len(repeats)} but ndim={self.ndimension()}." + ) + return self._repeat(*repeats) + + @abc.abstractmethod + def _repeat(self, *repeats: int) -> Self: + raise NotImplementedError + + def copy(self) -> Self: + """Return a shallow copy of the tensordict (ie, copies the structure but not the data). + + Equivalent to `TensorDictBase.clone(recurse=False)` + """ + return self.clone(recurse=False) + + def cat_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: + """Concatenates entries into a new entry and possibly remove the original values. + + Args: + keys (sequence of NestedKey): entries to concatenate. + + Keyword Arguments: + out_key (NestedKey): new key name for the concatenated inputs. + keep_entries (bool, optional): if ``False``, entries in ``keys`` will be deleted. + Defaults to ``False``. + dim (int, optional): the dimension along which the concatenation must occur. + Defaults to ``0``. + + Returns: self + + Examples: + >>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1)) + >>> td.cat_tensors("a", "b", out_key="c") + >>> assert "a" not in td + >>> assert (td["c"] == torch.tensor([0, 1])).all() + + """ + if keep_entries: + entries = [self.get(key) for key in keys] + else: + entries = [self.pop(key) for key in keys] + return self.set(out_key, torch.cat(entries, dim=dim)) + + def stack_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: + """Stacks entries into a new entry and possibly remove the original values. + + Args: + keys (sequence of NestedKey): entries to stack. + + Keyword Arguments: + out_key (NestedKey): new key name for the stacked inputs. + keep_entries (bool, optional): if ``False``, entries in ``keys`` will be deleted. + Defaults to ``False``. + dim (int, optional): the dimension along which the stack must occur. + Defaults to ``0``. + + Returns: self + + Examples: + >>> td = TensorDict(a=torch.zeros(()), b=torch.ones(())) + >>> td.stack_tensors("a", "b", out_key="c") + >>> assert "a" not in td + >>> assert (td["c"] == torch.tensor([0, 1])).all() + + """ + if keep_entries: + entries = [self.get(key) for key in keys] + else: + entries = [self.pop(key) for key in keys] + return self.set(out_key, torch.stack(entries, dim=dim)) + + def cat_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | List[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: # noqa: D417 + """Concatenates all entries of a tensordict in a single tensor. + + Args: + dim (int, optional): the dimension along which the entries should be concatenated. + + Keyword Args: + sorted (bool or list of NestedKeys): if ``True``, the entries will be concatenated in alphabetical order. + If ``False`` (default), the dict order will be used. Alternatively, a list of key names can be provided + and the tensors will be concatenated accordingly. This incurs some overhead as the list of keys will + be checked against the list of leaf names in the tensordict. + out (torch.Tensor, optional): an optional destination tensor for the cat operation. + + """ + if sorted in (None, False): + tensors = list(self.values(True, True)) + elif sorted in (True,): + tensors = list(self.values(True, True, sort=True)) + else: + keys = unravel_key_list(sorted) + if set(keys) != set(self.keys(True, True)): + raise RuntimeError( + "The provided set of keys differs from the tensordict list of keys." + ) + tensors = [self.get(key) for key in keys] + return torch.cat(tensors, dim, out=out) + + def stack_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | List[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: # noqa: D417 + """Stacks all entries of a tensordict in a single tensor. + + Args: + dim (int, optional): the dimension along which the entries should be stacked. + + Keyword Args: + sorted (bool or list of NestedKeys): if ``True``, the entries will be stacked in alphabetical order. + If ``False`` (default), the dict order will be used. Alternatively, a list of key names can be provided + and the tensors will be stacked accordingly. This incurs some overhead as the list of keys will + be checked against the list of leaf names in the tensordict. + out (torch.Tensor, optional): an optional destination tensor for the stack operation. + + """ + if sorted in (None, False): + tensors = list(self.values(True, True)) + elif sorted in (True,): + tensors = list(self.values(True, True, sort=True)) + else: + keys = unravel_key_list(sorted) + if set(keys) != set(self.keys(True, True)): + raise RuntimeError( + "The provided set of keys differs from the tensordict list of keys." + ) + tensors = [self.get(key) for key in keys] + return torch.stack(tensors, dim, out=out) + + @classmethod + def stack(cls, input, dim: int = 0, *, out=None): + """Stacks tensordicts into a single tensordict along the given dimension. + + This call is equivalent to calling :func:`torch.stack` but is compatible with torch.compile. + + """ + from tensordict._torch_func import _stack + + if not _is_tensor_collection(type(input[0])): + return torch.stack(input, dim, out=out) + return _stack(input, dim, out=out) + + @classmethod + def cat(cls, input, dim: int = 0, *, out=None): + """Concatenates tensordicts into a single tensordict along the given dimension. + + This call is equivalent to calling :func:`torch.cat` but is compatible with torch.compile. + + """ + from tensordict._torch_func import _cat + + if not _is_tensor_collection(type(input[0])): + return torch.cat(input, dim, out=out) + return _cat(input, dim, out=out) + + @classmethod + def lazy_stack(cls, input, dim: int = 0, *, out=None, **kwargs): + """Creates a lazy stack of tensordicts. + + See :meth:`~tensordict.LazyStackTensorDict.lazy_stack` for details. + """ + from tensordict._lazy import LazyStackedTensorDict + + return LazyStackedTensorDict.lazy_stack(input, dim=dim, out=out, **kwargs) + + @classmethod + def maybe_dense_stack(cls, input, dim: int = 0, *, out=None, **kwargs): + """Attempts to make a dense stack of tensordicts, and falls back on lazy stack when required.. + + See :meth:`~tensordict.LazyStackTensorDict.maybe_dense_stack` for details. + """ + from tensordict._lazy import LazyStackedTensorDict + + return LazyStackedTensorDict.maybe_dense_stack( + input, dim=dim, out=out, **kwargs + ) + + @abc.abstractmethod + def split( + self, split_size: int | list[int], dim: int = 0 + ) -> list[TensorCollection]: + """Splits each tensor in the TensorDict with the specified size in the given dimension, like `torch.split`. + + Returns a list of ``TensorDict`` instances with the view of split chunks of items. + + Args: + split_size (int or List(int)): size of a single chunk or list of sizes for each chunk. + dim (int): dimension along which to split the tensor. + + Returns: + A list of TensorDict with specified size in given dimension. + + Examples: + >>> td = TensorDict({ + ... 'x': torch.arange(12).reshape(3, 4), + ... }, batch_size=[3, 4]) + >>> td0, td1 = td.split([1, 2], dim=0) + >>> print(td0['x']) + torch.Tensor([[0, 1, 2, 3]]) + """ + raise NotImplementedError + + def gather(self, dim: int, index: Tensor, out: T | None = None) -> Self: + """Gathers values along an axis specified by `dim`. + + Args: + dim (int): the dimension along which collect the elements + index (torch.Tensor): a long tensor which number of dimension matches + the one of the tensordict with only one dimension differring between + the two (the gathering dimension). Its elements refer to the + index to be gathered along the required dimension. + out (TensorDictBase, optional): a destination tensordict. It must + have the same shape as the index. + + Examples: + >>> td = TensorDict( + ... {"a": torch.randn(3, 4, 5), + ... "b": TensorDict({"c": torch.zeros(3, 4, 5)}, [3, 4, 5])}, + ... [3, 4]) + >>> index = torch.randint(4, (3, 2)) + >>> td_gather = td.gather(dim=1, index=index) + >>> print(td_gather) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 2, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 2]), + device=None, + is_shared=False) + + Gather keeps the dimension names. + + Examples: + >>> td.names = ["a", "b"] + >>> td_gather = td.gather(dim=1, index=index) + >>> td_gather.names + ["a", "b"] + """ + return torch.gather(self, dim, index, out=out) + + @overload + def view(self, *shape: int): ... + + @overload + def view(self, dtype): ... + + @overload + def view(self, shape: torch.Size): ... + + @abc.abstractmethod + def _view( + self, + *args, + **kwargs, + ) -> Self: + raise NotImplementedError + + @_as_context_manager() + def view( + self, + *shape: int, + size: list | tuple | torch.Size | None = None, + batch_size: torch.Size | None = None, + ): + """Returns a tensordict with views of the tensors according to a new shape, compatible with the tensordict batch_size. + + Alternatively, a dtype can be provided as a first unnamed argument. In that case, all tensors will be viewed + with the according dtype. Note that this assume that the new shapes will be compatible with the provided dtype. + See :meth:`~torch.view` for more information on dtype views. + + Args: + *shape (int): new shape of the resulting tensordict. + dtype (torch.dtype): alternatively, a dtype to use to represent the tensor content. + size: iterable + + Keyword Args: + batch_size (torch.Size, optional): if a dtype is provided, the batch-size can be reset using this + keyword argument. If the ``view`` is called with a shape, this is without effect. + + Returns: + a new tensordict with the desired batch_size. + + Examples: + >>> td = TensorDict(source={'a': torch.zeros(3,4,5), + ... 'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4])) + >>> td_view = td.view(12) + >>> print(td_view.get("a").shape) # torch.Size([12, 5]) + >>> print(td_view.get("b").shape) # torch.Size([12, 10, 1]) + >>> td_view = td.view(-1, 4, 3) + >>> print(td_view.get("a").shape) # torch.Size([1, 4, 3, 5]) + >>> print(td_view.get("b").shape) # torch.Size([1, 4, 3, 10, 1]) + + """ + if len(shape) == 1 and isinstance(shape[0], torch.dtype): + dtype = shape[0] + return self._view_dtype(dtype=dtype, batch_size=batch_size) + _lazy_legacy = lazy_legacy() + + if _lazy_legacy: + return self._legacy_view(*shape, size=size) + else: + result = self._view(size=size) if size is not None else self._view(*shape) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + def _view_dtype(self, *, dtype, batch_size): + # We use apply because we want to check the shapes + def view(x): + return x.view(dtype) + + return self.apply(view, batch_size=batch_size) + + def _legacy_view( + self, + *shape: int, + size: list | tuple | torch.Size | None = None, + ) -> Self: + if len(shape) == 0 and size is not None: + return self.view(*size) + elif len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)): + return self.view(*shape[0]) + elif not isinstance(shape, torch.Size): + shape = infer_size_impl(shape, self.numel()) + shape = torch.Size(shape) + if shape == self.shape: + return self + from tensordict._lazy import _ViewedTensorDict + + return _ViewedTensorDict( + source=self, + custom_op="view", + inv_op="view", + custom_op_kwargs={"size": shape}, + inv_op_kwargs={"size": self.batch_size}, + ) + + @_as_context_manager() + def transpose(self, dim0, dim1): + """Returns a tensordict that is a transposed version of input. The given dimensions ``dim0`` and ``dim1`` are swapped. + + In-place or out-place modifications of the transposed tensordict will + impact the original tensordict too as the memory is shared and the operations + are mapped back on the original tensordict. + + Examples: + >>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4]) + >>> tensordict_transpose = tensordict.transpose(0, 1) + >>> print(tensordict_transpose.shape) + torch.Size([4, 3]) + >>> tensordict_transpose.set("b",, torch.randn(4, 3)) + >>> print(tensordict.get("b").shape) + torch.Size([3, 4]) + """ + _lazy_legacy = lazy_legacy() + + if _lazy_legacy: + return self._legacy_transpose(dim0, dim1) + else: + ndim = self.ndim + if dim0 < 0: + dim0 = ndim + dim0 + if dim1 < 0: + dim1 = ndim + dim1 + if dim0 < 0 or dim1 < 0 or dim0 >= ndim or dim1 >= ndim: + raise ValueError( + "dim0 and dim1 must be within the range of the number of dimensions." + ) + dim0, dim1 = min(dim0, dim1), max(dim0, dim1) + if dim0 == dim1: + return self + result = self._transpose(dim0, dim1) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + @abc.abstractmethod + def _transpose(self, dim0, dim1): + raise NotImplementedError + + def _legacy_transpose(self, dim0, dim1): + if dim0 < 0: + dim0 = self.ndim + dim0 + if dim1 < 0: + dim1 = self.ndim + dim1 + if any((dim0 < 0, dim1 < 0)): + raise ValueError( + "The provided dimensions are incompatible with the tensordict batch-size." + ) + if dim0 == dim1: + return self + from tensordict._lazy import _TransposedTensorDict + + return _TransposedTensorDict( + source=self, + custom_op="transpose", + inv_op="transpose", + custom_op_kwargs={"dim0": dim0, "dim1": dim1}, + inv_op_kwargs={"dim0": dim0, "dim1": dim1}, + ) + + @_as_context_manager() + def swapaxes(self, axis0: int, axis1: int): + """Interchange two axes of the tensordict. + + This is an alias for :meth:`~.transpose`. + + Args: + axis0 (int): First axis. + axis1 (int): Second axis. + + Returns: + a new tensordict with the axes swapped. + + Examples: + >>> td = TensorDict({"a": torch.randn(3, 4, 5)}, batch_size=[3, 4]) + >>> print(td.swapaxes(0, 1).shape) + torch.Size([4, 3]) + """ + return self.transpose(axis0, axis1) + + # Alias for swapaxes (matching torch.swapdims) + swapdims = swapaxes + + @overload + def permute(self, *dims: int): ... + + @overload + def permute(self, dims: list | tuple): ... + + @_as_context_manager() + def permute(self, *args, **kwargs): + """Returns a view of a tensordict with the batch dimensions permuted according to dims. + + Args: + *dims_list (int): the new ordering of the batch dims of the tensordict. Alternatively, + a single iterable of integers can be provided. + dims (list of int): alternative way of calling permute(...). + + Returns: + a new tensordict with the batch dimensions in the desired order. + + Examples: + >>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4]) + >>> print(tensordict.permute([1, 0])) + PermutedTensorDict( + source=TensorDict( + fields={ + a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, + batch_size=torch.Size([3, 4]), + device=cpu, + is_shared=False), + op=permute(dims=[1, 0])) + >>> print(tensordict.permute(1, 0)) + PermutedTensorDict( + source=TensorDict( + fields={ + a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, + batch_size=torch.Size([3, 4]), + device=cpu, + is_shared=False), + op=permute(dims=[1, 0])) + >>> print(tensordict.permute(dims=[1, 0])) + PermutedTensorDict( + source=TensorDict( + fields={ + a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, + batch_size=torch.Size([3, 4]), + device=cpu, + is_shared=False), + op=permute(dims=[1, 0])) + """ + _lazy_legacy = lazy_legacy() + + if _lazy_legacy: + return self._legacy_permute(*args, **kwargs) + else: + result = self._permute(*args, **kwargs) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + @abc.abstractmethod + def _permute( + self, + *args, + **kwargs, + ): + raise NotImplementedError + + def _legacy_permute( + self, + *dims_list: int, + dims: list[int] | None = None, + ) -> Self: + if len(dims_list) == 0: + dims_list = dims + elif len(dims_list) == 1 and not isinstance(dims_list[0], int): + dims_list = dims_list[0] + if len(dims_list) != len(self.shape): + raise RuntimeError( + f"number of dims don't match in permute (got {len(dims_list)}, expected {len(self.shape)}" + ) + + if not len(dims_list) and not self.batch_dims: + return self + if np.array_equal(dims_list, range(self.batch_dims)): + return self + min_dim, max_dim = -self.batch_dims, self.batch_dims - 1 + seen = [False for dim in range(max_dim + 1)] + for idx in dims_list: + if idx < min_dim or idx > max_dim: + raise IndexError( + f"dimension out of range (expected to be in range of [{min_dim}, {max_dim}], but got {idx})" + ) + if seen[idx]: + raise RuntimeError("repeated dim in permute") + seen[idx] = True + + from tensordict._lazy import _PermutedTensorDict + + return _PermutedTensorDict( + source=self, + custom_op="permute", + inv_op="permute", + custom_op_kwargs={"dims": list(map(int, dims_list))}, + inv_op_kwargs={"dims": list(map(int, dims_list))}, + ) + + @_as_context_manager() + def movedim( + self, source: int | tuple[int, ...], destination: int | tuple[int, ...] + ): + """Moves the dimension(s) of input at the position(s) in source to the position(s) in destination. + + Other dimensions of input that are not explicitly moved remain in their + original order and appear at the positions not specified in destination. + + Args: + source (int or tuple of ints): Original positions of the dims to move. + These must be unique. + destination (int or tuple of ints): Destination positions for each of + the original dims. These must also be unique. + + Returns: + a new tensordict with the batch dimensions moved to the desired positions. + + Examples: + >>> td = TensorDict({"a": torch.randn(3, 4, 5)}, batch_size=[3, 4]) + >>> print(td.movedim(0, 1).shape) + torch.Size([4, 3]) + >>> print(td.movedim((0, 1), (1, 0)).shape) + torch.Size([4, 3]) + """ + ndim = self.ndim + + # Normalize source and destination to tuples + if isinstance(source, int): + source = (source,) + if isinstance(destination, int): + destination = (destination,) + + if len(source) != len(destination): + raise ValueError( + f"movedim: source and destination must have the same number of elements, " + f"got {len(source)} and {len(destination)}" + ) + + # Normalize negative indices + source = tuple(s if s >= 0 else ndim + s for s in source) + destination = tuple(d if d >= 0 else ndim + d for d in destination) + + # Validate indices + for s in source: + if s < 0 or s >= ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of [-{ndim}, {ndim - 1}], but got {s})" + ) + for d in destination: + if d < 0 or d >= ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of [-{ndim}, {ndim - 1}], but got {d})" + ) + + # Check for duplicates + if len(set(source)) != len(source): + raise RuntimeError("movedim: repeated dim in source") + if len(set(destination)) != len(destination): + raise RuntimeError("movedim: repeated dim in destination") + + # Fast path: if source == destination, return self + if source == destination: + return self + + # Convert movedim to permute dims + # Build the permutation by: + # 1. Create list of dims not in source + # 2. Insert source dims at destination positions + remaining_dims = [i for i in range(ndim) if i not in source] + + # Sort source by destination to insert in correct order + sorted_pairs = sorted(zip(destination, source)) + perm = list(remaining_dims) + for dest, src in sorted_pairs: + perm.insert(dest, src) + + result = self._permute(perm) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + # Alias for movedim (matching torch.moveaxis) + moveaxis = movedim + + @_as_context_manager() + def flip(self, dims: int | tuple[int, ...]): + """Reverse the order of elements in the tensordict along the given dimensions. + + The shape of the tensordict is preserved, but the elements are reordered. + + Args: + dims (int or tuple of ints): Dimensions to flip. + + Returns: + a new tensordict with the dimensions flipped. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.flip(0)["a"]) + tensor([[3, 4, 5], + [0, 1, 2]]) + """ + if isinstance(dims, int): + dims = (dims,) + + ndim = self.ndim + dims = tuple(d if d >= 0 else ndim + d for d in dims) + + # Validate dimensions + for d in dims: + if d < 0 or d >= ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of [-{ndim}, {ndim - 1}], but got {d})" + ) + + def _flip(tensor): + return tensor.flip(dims) + + result = self._fast_apply( + _flip, + batch_size=self.batch_size, + call_on_nested=True, + names=self._maybe_names(), + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + @_as_context_manager() + def fliplr(self): + """Flip the tensordict in the left/right direction. + + Flip the entries in each row in the left/right direction. + Columns are preserved, but appear in a different order than before. + + Requires the tensordict to have at least 2 batch dimensions. + + Returns: + a new tensordict with the second dimension flipped. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.fliplr()["a"]) + tensor([[2, 1, 0], + [5, 4, 3]]) + """ + if self.ndim < 2: + raise RuntimeError("fliplr requires at least 2 batch dimensions") + return self.flip(1) + + @_as_context_manager() + def flipud(self): + """Flip the tensordict in the up/down direction. + + Flip the entries in each column in the up/down direction. + Rows are preserved, but appear in a different order than before. + + Requires the tensordict to have at least 1 batch dimension. + + Returns: + a new tensordict with the first dimension flipped. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.flipud()["a"]) + tensor([[3, 4, 5], + [0, 1, 2]]) + """ + if self.ndim < 1: + raise RuntimeError("flipud requires at least 1 batch dimension") + return self.flip(0) + + @_as_context_manager() + def roll(self, shifts: int | tuple[int, ...], dims: int | tuple[int, ...] = None): + """Roll the tensordict along the given dimensions. + + Elements that are shifted beyond the last position are re-introduced at + the first position. + + Args: + shifts (int or tuple of ints): The number of places by which the elements + of the tensordict are shifted. If shifts is a tuple, dims must be a + tuple of the same size, and each dimension will be rolled by the + corresponding value. + dims (int or tuple of ints, optional): Axis along which to roll. + By default, the tensordict is flattened before rolling. + + Returns: + a new tensordict with elements rolled. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.roll(1, 0)["a"]) + tensor([[3, 4, 5], + [0, 1, 2]]) + """ + + def _roll(tensor): + return tensor.roll(shifts, dims) + + result = self._fast_apply( + _roll, + batch_size=self.batch_size, + call_on_nested=True, + names=self._maybe_names(), + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + @_as_context_manager() + def rot90(self, k: int = 1, dims: tuple[int, int] = (0, 1)): + """Rotate the tensordict by 90 degrees in the plane specified by dims. + + Rotation direction is from the first towards the second axis. + + Args: + k (int): Number of times to rotate. Default: 1. + dims (tuple of two ints): The plane to rotate in. Default: (0, 1). + + Returns: + a new tensordict rotated by 90 degrees. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.rot90()["a"]) + tensor([[2, 5], + [1, 4], + [0, 3]]) + """ + if self.ndim < 2: + raise RuntimeError("rot90 requires at least 2 batch dimensions") + if len(dims) != 2: + raise RuntimeError("rot90 requires exactly 2 dims") + + # Normalize dims + ndim = self.ndim + dims = tuple(d if d >= 0 else ndim + d for d in dims) + + # Calculate new batch size + k = k % 4 # Normalize k to [0, 3] + if k == 0: + return self + + batch_size = list(self.batch_size) + if k == 1 or k == 3: + batch_size[dims[0]], batch_size[dims[1]] = ( + batch_size[dims[1]], + batch_size[dims[0]], + ) + + if self._has_names(): + names = list(self.names) + if k == 1 or k == 3: + names[dims[0]], names[dims[1]] = names[dims[1]], names[dims[0]] + else: + names = None + + def _rot90(tensor): + return tensor.rot90(k, dims) + + result = self._fast_apply( + _rot90, + batch_size=torch.Size(batch_size), + call_on_nested=True, + names=names, + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + def narrow(self, dim: int, start: int, length: int): + """Returns a new tensordict that is a narrowed version of the input. + + The dimension dim is input from start to start + length. + + Args: + dim (int): The dimension along which to narrow. + start (int): Starting index. + length (int): Length of the narrowed dimension. + + Returns: + a new tensordict narrowed along the specified dimension. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.narrow(1, 1, 2)["a"]) + tensor([[1, 2], + [4, 5]]) + """ + ndim = self.ndim + if dim < 0: + dim = ndim + dim + if dim < 0 or dim >= ndim: + raise IndexError( + f"Dimension out of range (expected to be in range of [-{ndim}, {ndim - 1}], but got {dim})" + ) + + batch_size = list(self.batch_size) + batch_size[dim] = length + + def _narrow(tensor): + return tensor.narrow(dim, start, length) + + result = self._fast_apply( + _narrow, + batch_size=torch.Size(batch_size), + call_on_nested=True, + names=self._maybe_names(), + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + def tile(self, dims: tuple[int, ...]): + """Construct a tensordict by repeating the elements. + + The dims argument specifies the number of repetitions in each dimension. + + Args: + dims (tuple of ints): The number of repetitions per dimension. + + Returns: + a new tensordict with elements repeated. + + Examples: + >>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3]) + >>> print(td["a"]) + tensor([[0, 1, 2], + [3, 4, 5]]) + >>> print(td.tile((2, 1))["a"]) + tensor([[0, 1, 2], + [3, 4, 5], + [0, 1, 2], + [3, 4, 5]]) + """ + if isinstance(dims, int): + dims = (dims,) + + # Calculate new batch size + ndim = self.ndim + if len(dims) > ndim: + # If more dims than batch dims, prepend 1s to batch_size + new_batch_size = [1] * (len(dims) - ndim) + list(self.batch_size) + for i, d in enumerate(dims): + new_batch_size[i] *= d + else: + # Pad dims with leading 1s + new_batch_size = list(self.batch_size) + offset = ndim - len(dims) + for i, d in enumerate(dims): + new_batch_size[offset + i] *= d + + def _tile(tensor): + return tensor.tile(dims) + + result = self._fast_apply( + _tile, + batch_size=torch.Size(new_batch_size), + call_on_nested=True, + names=None, # tile invalidates names + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + def broadcast_to(self, shape: tuple[int, ...]): + """Broadcasts the tensordict to a new shape. + + The new shape must be compatible with the original shape. + + Args: + shape (tuple of ints): The desired shape. + + Returns: + a new tensordict with the shape broadcast. + + Examples: + >>> td = TensorDict({"a": torch.arange(3)}, batch_size=[3]) + >>> print(td.broadcast_to((2, 3)).shape) + torch.Size([2, 3]) + """ + shape = torch.Size(shape) + + def _broadcast_to(tensor): + return tensor.broadcast_to(shape + tensor.shape[self.ndim :]) + + result = self._fast_apply( + _broadcast_to, + batch_size=shape, + call_on_nested=True, + names=None, # broadcast invalidates names + propagate_lock=True, + ) + self._maybe_set_shared_attributes(result) + if result._is_shared or result._is_memmap: + result.lock_() + return result + + @_as_context_manager() + def atleast_1d(self): + """Returns the tensordict with at least 1 batch dimension. + + If the tensordict already has 1 or more batch dimensions, it is returned unchanged. + Otherwise, a dimension of size 1 is prepended. + + Returns: + a tensordict with at least 1 batch dimension. + + Examples: + >>> td = TensorDict({"a": torch.randn(3)}, batch_size=[]) + >>> print(td.atleast_1d().shape) + torch.Size([1]) + """ + if self.ndim >= 1: + return self + return self.unsqueeze(0) + + @_as_context_manager() + def atleast_2d(self): + """Returns the tensordict with at least 2 batch dimensions. + + If the tensordict already has 2 or more batch dimensions, it is returned unchanged. + Otherwise, dimensions of size 1 are prepended to reach 2 dimensions. + + Returns: + a tensordict with at least 2 batch dimensions. + + Examples: + >>> td = TensorDict({"a": torch.randn(3)}, batch_size=[3]) + >>> print(td.atleast_2d().shape) + torch.Size([1, 3]) + """ + if self.ndim >= 2: + return self + elif self.ndim == 1: + return self.unsqueeze(0) + else: + return self.unsqueeze(0).unsqueeze(0) + + @_as_context_manager() + def atleast_3d(self): + """Returns the tensordict with at least 3 batch dimensions. + + If the tensordict already has 3 or more batch dimensions, it is returned unchanged. + Otherwise, dimensions of size 1 are prepended to reach 3 dimensions. + + Returns: + a tensordict with at least 3 batch dimensions. + + Examples: + >>> td = TensorDict({"a": torch.randn(3)}, batch_size=[3]) + >>> print(td.atleast_3d().shape) + torch.Size([1, 1, 3]) + """ + if self.ndim >= 3: + return self + elif self.ndim == 2: + return self.unsqueeze(0) + elif self.ndim == 1: + return self.unsqueeze(0).unsqueeze(0) + else: + return self.unsqueeze(0).unsqueeze(0).unsqueeze(0) + + # Cache functionality + def _erase_cache(self): + self._cache = None + + # Dim names functionality + @property + @abc.abstractmethod + def names(self): + """The dimension names of the tensordict. + + The names can be set at construction time using the ``names`` argument. + + See also :meth:`~.refine_names` for details on how to set the names after + construction. + """ + raise NotImplementedError + + @names.setter + def names(self, value): + self._set_names(value) + + def _get_names_idx(self, idx): + if not self._has_names(): + return None + + def is_boolean(idx): + try: + from functorch import dim as ftdim + + except ImportError: + from tensordict.utils import _ftdim_mock as ftdim + + if isinstance(idx, ftdim.Dim): + return None + if isinstance(idx, tuple) and len(idx) == 1: + return is_boolean(idx[0]) + if hasattr(idx, "dtype") and idx.dtype is torch.bool: + return idx.ndim + return None + + num_boolean_dim = is_boolean(idx) + names = self.names + if num_boolean_dim: + names = [None] + names[num_boolean_dim:] + else: + if not isinstance(idx, tuple): + idx = (idx,) + if len([_idx for _idx in idx if _idx is not None]) < self.ndim: + idx = (*idx, Ellipsis) + idx_names = convert_ellipsis_to_idx(idx, self.batch_size) + # this will convert a [None, :, :, 0, None, 0] in [None, 0, 1, None, 3] + count = 0 + idx_to_take = [] + no_more_tensors = False + for _idx in idx_names: + if _idx is None: + idx_to_take.append(None) + elif _is_number(_idx): + count += 1 + elif isinstance(_idx, (torch.Tensor, np.ndarray)): + if not no_more_tensors: + idx_to_take.extend([count] * _idx.ndim) + count += 1 + no_more_tensors = True + else: + # skip this one + count += 1 + else: + idx_to_take.append(count) + count += 1 + names = [names[i] if i is not None else None for i in idx_to_take] + if all(name is None for name in names): + return None + return names + + @abc.abstractmethod + def _erase_names(self): + """Erases the dimension names from a tensordict.""" + raise NotImplementedError + + @abc.abstractmethod + def _rename_subtds(self, value): + """Renames all the sub-tensordicts dimension according to value. + + If value has less dimensions than the TD, the rest is just assumed to be None. + """ + raise NotImplementedError + + def _check_dim_name(self, name): + if name is None: + return False + if self._has_names() and name in self.names: + return True + for key in self.keys(): + if _is_tensor_collection(self.entry_class(key)): + if self._get_str(key, NO_DEFAULT)._check_dim_name(name): + return True + else: + return False + + def refine_names(self, *names) -> Self: + """Refines the dimension names of self according to names. + + Refining is a special case of renaming that "lifts" unnamed dimensions. + A None dim can be refined to have any name; a named dim can only be + refined to have the same name. + + Because named tensors can coexist with unnamed tensors, refining names + gives a nice way to write named-tensor-aware code that works with both + named and unnamed tensors. + + names may contain up to one Ellipsis (...). The Ellipsis is expanded + greedily; it is expanded in-place to fill names to the same length as + self.dim() using names from the corresponding indices of self.names. + + Returns: the same tensordict with dimensions named according to the input. + + Examples: + >>> td = TensorDict({}, batch_size=[3, 4, 5, 6]) + >>> tdr = td.refine_names(None, None, None, "d") + >>> assert tdr.names == [None, None, None, "d"] + >>> tdr = td.refine_names("a", None, None, "d") + >>> assert tdr.names == ["a", None, None, "d"] + + """ + # replace ellipsis if any + names_copy = list(names) + if any(name is Ellipsis for name in names): + ellipsis_name = [NO_DEFAULT for _ in range(self.ndim - len(names) + 1)] + names = [] + for name in names_copy: + if name is Ellipsis: + names += ellipsis_name + else: + names.append(name) + # check that the names that are set are either None or identical + curr_names = self.names + for i, name in enumerate(names): + if name is NO_DEFAULT: + # whatever value is ok + names[i] = curr_names[i] + continue + else: + if curr_names[i] is None: + continue + if self.names[i] == name: + continue + else: + raise RuntimeError( + f"refine_names: cannot coerce TensorDict names {self.names} with {names_copy}." + ) + self._set_names(names) + # we also need to rename the sub-tensordicts + # self._rename_subtds(self.names) + return self + + def rename(self, *names, **rename_map): + """Returns a clone of the tensordict with dimensions renamed. + + Examples: + >>> td = TensorDict({}, batch_size=[1, 2, 3 ,4]) + >>> td.names = list("abcd") + >>> td_rename = td.rename(c="g") + >>> assert td_rename.names == list("abgd") + + """ + clone = self.clone(recurse=False) + if len(names) == 1 and names[0] is None: + clone.names = None + if rename_map and names: + raise ValueError( + "Passed both a name map and a name list. Only one is accepted." + ) + elif not rename_map and not names: + raise ValueError( + "Neither a name map nor a name list was passed. " + "Only one is accepted." + ) + elif rename_map: + cnames = list(clone.names) + for i, name in enumerate(cnames): + new_name = rename_map.pop(name, NO_DEFAULT) + if new_name is not NO_DEFAULT: + cnames[i] = new_name + clone.names = cnames + if rename_map: + raise ValueError( + f"Some names to be renamed were not part of the tensordict names: {rename_map.keys()} vs {self.names}." + ) + else: + clone.names = names + return clone + + def rename_(self, *names, **rename_map): + """Same as :meth:`~.rename`, but executes the renaming in-place. + + Examples: + >>> td = TensorDict({}, batch_size=[1, 2, 3 ,4]) + >>> td.names = list("abcd") + >>> assert td.rename_(c="g") + >>> assert td.names == list("abgd") + """ + if len(names) == 1 and names[0] is None: + self._set_names(None) + if rename_map and names: + raise ValueError( + "Passed both a name map and a name list. " "Only one is accepted." + ) + elif not rename_map and not names and self.batch_dims: + raise ValueError( + "Neither a name map nor a name list was passed. " + "Only one is accepted." + ) + elif rename_map: + cnames = list(self.names) + for i, name in enumerate(cnames): + new_name = rename_map.pop(name, NO_DEFAULT) + if new_name is not NO_DEFAULT: + cnames[i] = new_name + if rename_map: + raise ValueError( + f"Some names to be renamed were not part of the tensordict names: {rename_map.keys()} vs {self.names}." + ) + self._set_names(cnames) + else: + self._set_names(names) + return self + + @abc.abstractmethod + def _has_names(self) -> bool: + raise NotImplementedError + + def _maybe_names(self) -> Sequence[str] | None: + if self._has_names(): + return self.names + return None + + @property + def _has_non_tensor(self): + """Checks if the tensordict has non-tensor data.""" + for value in self.values(True, True, is_leaf=_is_leaf_nontensor): + if _is_non_tensor(type(value)): + return True + return False + + # Device functionality: device is optional. If provided, it will enforce + # all data is on the same device + @property + @abc.abstractmethod + def device(self) -> torch.device | None: + """Device of a TensorDict. + + If the TensorDict has a specified device, all + its tensors (incl. nested ones) must live on the same device. + If the TensorDict device is ``None``, different values can be located + on different devices. + + Returns: + torch.device object indicating the device where the tensors + are placed, or None if TensorDict does not have a device. + + Examples: + >>> td = TensorDict({ + ... "cpu": torch.randn(3, device='cpu'), + ... "cuda": torch.randn(3, device='cuda'), + ... }, batch_size=[], device=None) + >>> td['cpu'].device + device(type='cpu') + >>> td['cuda'].device + device(type='cuda') + >>> td = TensorDict({ + ... "x": torch.randn(3, device='cpu'), + ... "y": torch.randn(3, device='cuda'), + ... }, batch_size=[], device='cuda') + >>> td['x'].device + device(type='cuda') + >>> td['y'].device + device(type='cuda') + >>> td = TensorDict({ + ... "x": torch.randn(3, device='cpu'), + ... "y": TensorDict({'z': torch.randn(3, device='cpu')}, batch_size=[], device=None), + ... }, batch_size=[], device='cuda') + >>> td['x'].device + device(type='cuda') + >>> td['y'].device # nested tensordicts are also mapped onto the appropriate device. + device(type='cuda') + >>> td['y', 'x'].device + device(type='cuda') + + """ + raise NotImplementedError + + @device.setter + @abc.abstractmethod + def device(self, value: DeviceType) -> None: + raise NotImplementedError + + @lock_blocked + def clear(self) -> Self: + """Erases the content of the tensordict.""" + for key in list(self.keys()): + del self[key] + return self + + @classmethod + def fromkeys(cls, keys: List[NestedKey], value: Any = 0): + """Creates a tensordict from a list of keys and a single value. + + Args: + keys (list of NestedKey): An iterable specifying the keys of the new dictionary. + value (compatible type, optional): The value for all keys. Defaults to ``0``. + """ + from tensordict._td import TensorDict + + return TensorDict(dict.fromkeys(keys, value), batch_size=[]) + + @abc.abstractmethod + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + """Removes the item that was last inserted into the TensorDict. + + ``popitem`` will only return non-nested values. + """ + raise NotImplementedError + + def clear_device_(self) -> Self: + """Clears the device of the tensordict. + + Returns: self + + """ + self._device = None + for value in self.values(): + if _is_tensor_collection(type(value)): + value.clear_device_() + return self + + def _set_device(self, device: torch.device) -> Self: + self._device = device + for value in self.values(): + if _is_tensor_collection(type(value)): + value._set_device(device=device) + return self + + @cache # noqa: B019 + def param_count(self, *, count_duplicates: bool = True) -> int: + """Counts the number of parameters (total number of indexable items), accounting for tensors only. + + Keyword Args: + count_duplicates (bool): Whether to count duplicated tensor as independent or not. + If ``False``, only strictly identical tensors will be discarded (same views but different + ids from a common base tensor will be counted twice). Defaults to `True` (each tensor is assumed + to be a single copy). + + """ + vals = self._values_list(True, True) + total = 0 + if not count_duplicates: + vals = set(vals) + for v in vals: + total += v.numel() + return total + + @cache # noqa: B019 + def bytes(self, *, count_duplicates: bool = True) -> int: + """Counts the number of bytes of the contained tensors. + + Keyword Args: + count_duplicates (bool): Whether to count duplicated tensor as independent or not. + If ``False``, only strictly identical tensors will be discarded (same views but different + ids from a common base tensor will be counted twice). Defaults to `True` (each tensor is assumed + to be a single copy). + + """ + set_of_tensors = set() if not count_duplicates else [] + + def add(tensor): + if count_duplicates: + set_of_tensors.append(tensor) + else: + set_of_tensors.add(tensor) + + def count_bytes(tensor): + if tensor.is_nested: + if not tensor.layout == torch.jagged: + raise RuntimeError( + "NTs that are not jagged are not supported by the bytes method. Please use the jagged layout instead " + "or raise and issue on https://github.com/pytorch/tensordict/issues instead." + ) + attrs, ctx = tensor.__tensor_flatten__() + for attr in attrs: + t = getattr(tensor, attr) + count_bytes(t) + return + if isinstance(tensor, torch.Tensor): + if isinstance(tensor, MemoryMappedTensor): + add(tensor) + return + if type(tensor) in (Tensor, Parameter, Buffer): + pass + elif hasattr(tensor, "__tensor_flatten__"): + attrs, ctx = tensor.__tensor_flatten__() + for attr in attrs: + t = getattr(tensor, attr) + count_bytes(t) + return + else: + warnings.warn( + "The sub-tensor doesn't ot have a __tensor_flatten__ attribute, making it " + "impossible to count the bytes it contains. Falling back on regular count.", + category=UserWarning, + ) + count_bytes(torch.as_tensor(tensor)) + return + + grad = getattr(tensor, "grad", None) + if grad is not None: + count_bytes(grad) + count_bytes(tensor.data) + else: + add(tensor) + return + + vals = self._values_list(True, True) + for v in vals: + count_bytes(v) + total = 0 + for tensor in set_of_tensors: + total += tensor.numel() * tensor.dtype.itemsize + return total + + def pin_memory(self, num_threads: int | None = None, inplace: bool = False) -> Self: + """Calls :meth:`~torch.Tensor.pin_memory` on the stored tensors. + + Args: + num_threads (int or str): if provided, the number of threads to use + to call ``pin_memory`` on the leaves. Defaults to ``None``, which sets a high + number of threads in :class:`~concurrent.futures.ThreadPoolExecutor(max_workers=None)`. + To execute all the calls to :meth:`~torch.Tensor.pin_memory` on the main thread, pass + ``num_threads=0``. + inplace (bool, optional): if ``True``, the tensordict is modified in-place. + Defaults to ``False``. + + """ + + def pin_memory(x): + return x.pin_memory() + + return self._fast_apply( + pin_memory, + num_threads=num_threads, + inplace=inplace, + propagate_lock=True, + ) + + def pin_memory_(self, num_threads: int | str = 0) -> Self: + """Calls :meth:`~torch.Tensor.pin_memory` on the stored tensors and returns the TensorDict modifies in-place. + + Args: + num_threads (int or str): if provided, the number of threads to use + to call ``pin_memory`` on the leaves. If ``"auto"`` is passed, the + number of threads is automatically determined. + + """ + return self.pin_memory(num_threads=num_threads, inplace=True) + + def cpu(self, **kwargs) -> Self: + """Casts a tensordict to CPU. + + This function also supports all the keyword arguments of :meth:`~.to`. + """ + return self.to("cpu", **kwargs) + + def cuda(self, device: int | None = None, **kwargs) -> Self: + """Casts a tensordict to a cuda device (if not already on it). + + Args: + device (int, optional): if provided, the cuda device on which the + tensor should be cast. + + This function also supports all the keyword arguments of :meth:`~.to`. + + """ + if device is None: + return self.to(torch.device("cuda")) + return self.to(f"cuda:{device}", **kwargs) + + @property + def is_cuda(self): + return self.device is not None and self.device.type == "cuda" + + @property + def is_cpu(self): + return self.device is not None and self.device.type == "cpu" + + # Serialization functionality + def state_dict( + self, + destination=None, + prefix="", + keep_vars=False, + flatten=False, + ) -> OrderedDict[str, Any]: + """Produces a state_dict from the tensordict. + + The structure of the state-dict will still be nested, unless ``flatten`` is set to ``True``. + + A tensordict state-dict contains all the tensors and meta-data needed + to rebuild the tensordict (names are currently not supported). + + Args: + destination (dict, optional): If provided, the state of tensordict will + be updated into the dict and the same object is returned. + Otherwise, an ``OrderedDict`` will be created and returned. + Default: ``None``. + prefix (str, optional): a prefix added to tensor + names to compose the keys in state_dict. Default: ``''``. + keep_vars (bool, optional): by default the :class:`torch.Tensor` items + returned in the state dict are detached from autograd. If it's + set to ``True``, detaching will not be performed. + Default: ``False``. + flatten (bool, optional): whether the structure should be flattened + with the ``"."`` character or not. + Defaults to ``False``. + + Examples: + >>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) + >>> sd = data.state_dict() + >>> print(sd) + OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)]) + >>> sd = data.state_dict(flatten=True) + OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)]) + + """ + out = collections.OrderedDict() + source = self + if flatten: + source = source.flatten_keys(".") + for key, item in source.items(): + if not _is_tensor_collection(type(item)): + if not keep_vars: + out[prefix + key] = item.detach().clone() + else: + out[prefix + key] = item + else: + out[prefix + key] = item.state_dict(keep_vars=keep_vars) + if "__batch_size" in out: + raise KeyError( + "Cannot retrieve the state_dict of a TensorDict with `'__batch_size'` key" + ) + if "__device" in out: + raise KeyError( + "Cannot retrieve the state_dict of a TensorDict with `'__batch_size'` key" + ) + out[prefix + "__batch_size"] = source.batch_size + out[prefix + "__device"] = source.device + if destination is not None: + destination.update(out) + return destination + return out + + def load_state_dict( + self, + state_dict: OrderedDict[str, Any], + strict=True, + assign=False, + from_flatten=False, + ) -> Self: + """Loads a state-dict, formatted as in :meth:`~.state_dict`, into the tensordict. + + Args: + state_dict (OrderedDict): the state_dict of to be copied. + strict (bool, optional): whether to strictly enforce that the keys + in :attr:`state_dict` match the keys returned by this tensordict's + :meth:`torch.nn.Module.state_dict` function. Default: ``True`` + assign (bool, optional): whether to assign items in the state + dictionary to their corresponding keys in the tensordict instead + of copying them inplace into the tensordict's current tensors. + When ``False``, the properties of the tensors in the current + module are preserved while when ``True``, the properties of the + Tensors in the state dict are preserved. + Default: ``False`` + from_flatten (bool, optional): if ``True``, the input state_dict is + assumed to be flattened. + Defaults to ``False``. + + Examples: + >>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) + >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) + >>> sd = data.state_dict() + >>> data_zeroed.load_state_dict(sd) + >>> print(data_zeroed["3", "3"]) + tensor(3) + >>> # with flattening + >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) + >>> data_zeroed.load_state_dict(data.state_dict(flatten=True), from_flatten=True) + >>> print(data_zeroed["3", "3"]) + tensor(3) + + + """ + if from_flatten: + self_flatten = self.flatten_keys(".") + self_flatten.load_state_dict(state_dict, strict=strict, assign=assign) + if not assign: + # modifications are done in-place so we should be fine returning self + return self + else: + # run a check over keys, if we any key with a '.' in name we're doomed + DOT_ERROR = "Cannot use load_state_dict(..., from_flatten=True, assign=True) when some keys contain a dot character." + for key in self.keys(True, True): + if isinstance(key, tuple): + for subkey in key: + if "." in subkey: + raise RuntimeError(DOT_ERROR) + elif "." in key: + raise RuntimeError(DOT_ERROR) + return self.update(self_flatten.unflatten_keys(".")) + + # copy since we'll be using pop + if is_compiling(): + state_dict = type(state_dict)(state_dict) + else: + state_dict = copy(state_dict) + batch_size = state_dict.pop("__batch_size") + device = state_dict.pop("__device", None) + + if strict and set(state_dict.keys()) != set(self.keys()): + set_sd = set(state_dict.keys()) + set_td = set(self.keys()) + + # if there are keys in state-dict that point to an empty tensordict + # or if the local tensordicts are empty, we can skip + def _is_empty_dict(sd, key=None): + if key is not None: + if not isinstance(sd[key], dict): + return False + return _is_empty_dict(sd[key]) + for key, item in sd.items(): + if key in ("__batch_size", "__device"): + continue + if isinstance(item, dict): + if not _is_empty_dict(item): + return False + continue + return False + else: + return True + + def check_is_empty(target, key): + item = target.get(key) + if not is_tensor_collection(item) or not item.is_empty(): + return False + return True + + if not all(check_is_empty(self, key) for key in set_td - set_sd) or not all( + _is_empty_dict(state_dict, key) for key in set_sd - set_td + ): + raise RuntimeError( + "Cannot load state-dict because the key sets don't match: got " + f"state_dict extra keys \n{set_sd - set_td}\n and tensordict extra keys\n{set_td - set_sd}\n" + ) + + self.batch_size = batch_size + if device is not None and self.device is not None and device != self.device: + raise RuntimeError("Loading data from another device is not yet supported.") + + for key, item in state_dict.items(): + if isinstance(item, dict): + dest = self.get(key, None) + if dest is None: + dest = self.empty() + dest.load_state_dict(item, assign=assign, strict=strict) + self.set( + key, + dest, + inplace=not assign, + ) + else: + self.set(key, item, inplace=not assign) + return self + + def is_shared(self) -> bool: + """Checks if tensordict is in shared memory. + + If a TensorDict instance is in shared memory, it is locked (entries cannot + be renamed, removed or added). If a ``TensorDict`` is created with + tensors that are all in shared memory, this does __not__ mean that ``is_shared`` + will return ``True`` (as a new tensor may or may not be in shared memory). + Only if one calls `tensordict.share_memory_()` or places the tensordict + on a device where the content is shared by default (eg, ``"cuda"``) + will the tensordict be considered in shared memory. + + This is always ``True`` for tensordicts on a CUDA device. + + """ + if self.device and not self._is_memmap: + return self.device.type == "cuda" or self._is_shared + return self._is_shared + + def is_memmap(self) -> bool: + """Checks if tensordict is memory-mapped. + + If a TensorDict instance is memory-mapped, it is locked (entries cannot + be renamed, removed or added). If a ``TensorDict`` is created with + tensors that are all memory-mapped, this does __not__ mean that ``is_memmap`` + will return ``True`` (as a new tensor may or may not be memory-mapped). + Only if one calls `tensordict.memmap_()` will the tensordict be + considered as memory-mapped. + + This is always ``True`` for tensordicts on a CUDA device. + + """ + return self._is_memmap + + @abc.abstractmethod + def share_memory_(self) -> Self: + """Places all the tensors in shared memory. + + The TensorDict is then locked, meaning that any writing operations that + isn't in-place will throw an exception (eg, rename, set or remove an + entry). + Conversely, once the tensordict is unlocked, the share_memory attribute + is turned to ``False``, because cross-process identity is not + guaranteed anymore. + + Returns: + self + + """ + raise NotImplementedError + + @abc.abstractmethod + def _memmap_( + self, + *, + prefix: str | None, + copy_existing: bool, + executor, + futures, + inplace, + like, + share_non_tensor, + existsok, + robust_key, + ) -> Self: + raise NotImplementedError + + def densify(self, layout: torch.layout = torch.strided): + """Attempts to represent the lazy stack with contiguous tensors (plain tensors or nested). + + Keyword Args: + layout (torch.layout): the layout of the nested tensors, if any. Defaults to + :class:`~torch.strided`. + + """ + any_set = False + out_dict = {} + for key, val in self.items(): + if is_tensor_collection(val): + val_dense = val.densify(layout=layout) + any_set = any_set | (val_dense is not val) + val = val_dense + out_dict[key] = val + if any_set: + result = self.empty() + for key, val in out_dict.items(): + result._set_str(key, val, validated=True, inplace=False) + return result + return self + + @property + def saved_path(self): + """Returns the path where a memmap saved TensorDict is being stored. + + This argument valishes as soon as is_memmap() returns ``False`` (e.g., when the tensordict is unlocked). + """ + if self.is_memmap(): + path = self._memmap_prefix + return path + raise AttributeError( + f"The tensordict has no saved path (memmap={self.is_memmap()}, path={self._memmap_prefix})." + ) + + # Generic method to get a class metadata + def _reduce_get_metadata(self): + return { + "device": str(self.device) if self.device is not None else None, + "names": self.names, + "batch_size": list(self.batch_size), + "is_locked": self._is_locked, + } + + # @cache # noqa: B019 + def _reduce_vals_and_metadata(self, *, dtype=NO_DEFAULT, requires_metadata): + """Returns a nested dictionary of metadata, a flat Dict[NestedKey, Tensor] containing tensor data and a list of tensor sizes.""" + if dtype is NO_DEFAULT: + dtype = self.dtype + need_padding = dtype is None + # If the dtype is not unique (self.dtype is None) then we need the metadata + # because we need a custom unpickler + requires_metadata = requires_metadata | need_padding + + if requires_metadata: + # metadata is nested + cls = type(self) + from tensordict._reductions import CLS_MAP + + if cls.__name__ in CLS_MAP: + cls = cls.__name__ + else: + pass + metadata_dict = { + "cls": cls, + "non_tensors": {}, + "leaves": {}, + "cls_metadata": self._reduce_get_metadata(), + } + else: + metadata_dict = None + + # flat_key_values is flat + flat_key_values = {} + + flat_size = [] + start = 0 + + def add_single_value(value, key, metadata_dict, dtype, shape, flat_size): + nonlocal start + n = value.element_size() * value.numel() + if need_padding: + pad = n % 8 + if pad != 0: + pad = 8 - pad + else: + pad = 0 + flat_size.append(sum([n, pad])) + # Using sum to tell dynamo to use sym_sum + stop = sum([start, flat_size[-1]]) + if requires_metadata: + metadata_dict["leaves"][key] = ( + _DTYPE_TO_STR_DTYPE[dtype], + list(shape), + # _DEVICE2STRDEVICE[device], + start, + stop, + pad, + ) + start = stop + + def assign( + key, + value, + track_key=(), + metadata_dict=metadata_dict, + flat_size=flat_size, + ): + total_key = key if isinstance(key, tuple) else (key,) + total_key = track_key + total_key + cls = type(value) + if issubclass(cls, torch.Tensor): + pass + # We want to skip NonTensorStacks + elif _is_non_tensor(cls) and not issubclass(cls, TensorDictBase): + if requires_metadata: + metadata_dict["non_tensors"][key] = ( + value.data, + list(value.batch_size), + str(value.device) if value.device is not None else None, + ) + return + elif _is_tensor_collection(cls): + metadata_dict_key = None + if requires_metadata: + from tensordict._reductions import CLS_MAP + + if cls.__name__ in CLS_MAP: + cls = cls.__name__ + else: + pass + metadata_dict_key = metadata_dict[key] = { + "cls": cls, + "non_tensors": {}, + "leaves": {}, + "cls_metadata": value._reduce_get_metadata(), + } + + def local_assign(*t): + return assign( + *t, + track_key=total_key, + metadata_dict=metadata_dict_key, + flat_size=flat_size, + ) + + value._fast_apply( + local_assign, + named=True, + nested_keys=True, + call_on_nested=True, + is_leaf=_NESTED_TENSORS_AS_LISTS_NONTENSOR, + ) + return + # Tensors: DTensor, nested and then regular + if hasattr(value, "full_tensor"): + raise NotImplementedError("DTensor is not supported yet") + if getattr(value, "is_nested", False): + if value.layout is torch.jagged: + # Get the values + values = value._values + shape = [v if isinstance(v, int) else -1 for v in values.shape] + # Get the offsets + offsets = value._offsets + # Get the lengths + lengths = value._lengths + + # Now we're saving the two tensors + # We will rely on the fact that the writing order is preserved in python dict + # (since python 3.7). Later, we will read the NJT then the NJT offset in that order + # to do the allocation. + flat_key_values[_prefix_last_key(total_key, "")] = value + flat_size.append(0) + flat_key_values[_prefix_last_key(total_key, "")] = ( + values + ) + add_single_value( + values, + _prefix_last_key(key, ""), + metadata_dict, + values.dtype, + shape, + flat_size, + ) + # Lengths + if lengths is not None: + flat_key_values[ + _prefix_last_key(total_key, "") + ] = lengths + add_single_value( + lengths, + _prefix_last_key(key, ""), + metadata_dict, + lengths.dtype, + lengths.shape, + flat_size, + ) + # Offsets + flat_key_values[_prefix_last_key(total_key, "")] = ( + offsets + ) + add_single_value( + offsets, + _prefix_last_key(key, ""), + metadata_dict, + offsets.dtype, + offsets.shape, + flat_size, + ) + + else: + raise NotImplementedError( + "NST is not supported, please use layout=torch.jagged when building the nested tensor." + ) + return + flat_key_values[total_key] = value + add_single_value( + value, + key, + metadata_dict, + value.dtype, + value.shape, + # value.device, + flat_size, + ) + + self._fast_apply( + assign, + named=True, + call_on_nested=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS_NONTENSOR, + filter_empty=True, + ) + return metadata_dict, flat_key_values, flat_size, need_padding + + def consolidate( + self, + filename: Path | str | None = None, + *, + num_threads=0, + device: torch.device | None = None, + non_blocking: bool = False, + inplace: bool = False, + return_early: bool = False, + use_buffer: bool = False, + share_memory: bool = False, + pin_memory: bool = False, + metadata: bool = False, + ) -> None: + """Consolidates the tensordict content in a single storage for fast serialization. + + Args: + filename (Path, optional): an optional file path for a memory-mapped tensor + to use as a storage for the tensordict. + + Keyword Args: + num_threads (integer, optional): the number of threads to use for populating + the storage. + device (torch.device, optional): an optional device where the storage must be + instantiated. + non_blocking (bool, optional): ``non_blocking`` argument passed to :meth:`~torch.Tensor.copy_`. + inplace (bool, optional): if ``True``, the resulting tensordict is the same + as ``self`` with updated values. Defaults to ``False``. + return_early (bool, optional): if ``True`` and ``num_threads>0``, + the method will return a future of the tensordict. The resulting + tensordict can be queried using `future.result()`. + use_buffer (bool, optional): if ``True`` and a filename is passed, an intermediate + local buffer will be created in shared memory, and the data will be copied at + the storage location as a last step. This may be faster than writing directly + to a distant physical memory (e.g., NFS). + Defaults to ``False``. + share_memory (bool, optional): if ``True``, the storage will be placed in shared memory. + Defaults to ``False``. + pin_memory (bool, optional): whether the consolidated data should be placed in pinned + memory. Defaults to ``False``. + metadata (bool, optional): if ``True``, the metadata will be stored alongisde the + common storage. If a filename is provided, this is without effect. + Storing the metadata can be useful when one wants to control how serialization + is achieved, as TensorDict handles the pickling/unpickling of consolidated TDs + differently if the metadata is or isn't available. + + .. note:: + If the tensordict is already consolidated, all arguments are ignored and ``self`` + is returned. Call :meth:`~.contiguous` to re-consolidate. + + Examples: + >>> import pickle + >>> import tempfile + >>> import torch + >>> import tqdm + >>> from torch.utils.benchmark import Timer + >>> from tensordict import TensorDict + >>> data = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) + >>> data_consolidated = data.consolidate() + >>> # check that the data has a single data_ptr() + >>> assert torch.tensor([ + ... v.untyped_storage().data_ptr() for v in data_c.values(True, True) + ... ]).unique().numel() == 1 + >>> # Serializing the tensordict will be faster with data_consolidated + >>> with open("data.pickle", "wb") as f: + ... print("regular", Timer("pickle.dump(data, f)", globals=globals()).adaptive_autorange()) + >>> with open("data_c.pickle", "wb") as f: + ... print("consolidated", Timer("pickle.dump(data_consolidated, f)", globals=globals()).adaptive_autorange()) + + + """ + if self.is_consolidated(): + return self + + ( + metadata_dict, + flat_dict, + flat_size, + need_padding, + ) = self._reduce_vals_and_metadata( + requires_metadata=filename is not None or metadata, dtype=None + ) + filesize = sum(flat_size) + device = torch.device(device) if device is not None else None + if filename is None: + storage = torch.empty( + filesize, + dtype=torch.uint8, + device=device if device else self.device, + pin_memory=pin_memory, + ) + if share_memory and not ( + device is not None and device.type == "cuda" + ): # cuda device is always shared + storage.share_memory_() + else: + # Convert the dict to json + try: + from tensordict.utils import json_dumps + + metadata_dict_json = json_dumps(metadata_dict) + except TypeError as e: + raise RuntimeError( + "Failed to convert the metatdata to json. " + "This is usually due to a nested class that is unaccounted for by the serializer, " + "such as custom TensorClass. " + "If you encounter this error, please file an issue on github." + ) from e + # Represent as a tensor + if isinstance(metadata_dict_json, str): + metadata_dict_json = metadata_dict_json.encode("utf-8") + metadata_dict_json = torch.as_tensor( + bytearray(metadata_dict_json), dtype=torch.uint8 + ) + len_metadata = torch.tensor( + [metadata_dict_json.numel()], dtype=torch.int64 + ).view(torch.uint8) + + if device not in (torch.device("cpu"), None): + raise RuntimeError( + "device and filename are mutually exclusive arguments." + ) + suffix = len_metadata.numel() + metadata_dict_json.numel() + if not use_buffer: + total_storage = torch.from_file( + str(filename), + size=filesize + suffix, + dtype=torch.uint8, + shared=True, + # needed when device ctx differs + device=torch.device("cpu"), + ) + else: + total_storage = MemoryMappedTensor.empty( + shape=(filesize + suffix,), + dtype=torch.uint8, + ) + + total_storage[-8:] = len_metadata + total_storage[-8 - metadata_dict_json.numel() : -8] = metadata_dict_json + storage = total_storage[:-suffix] + # assert len(storage.untyped_storage()) == filesize + + offsets = torch.tensor([0] + flat_size).cumsum(0).tolist() + + def view_old_as_new(v, oldv): + v = v.view(oldv.dtype) + if v.numel() > oldv.numel(): + return v[: oldv.numel()].view(oldv.shape) + return v.view(oldv.shape) + + if num_threads is None: + num_threads = 0 + if num_threads > 0: + + def assign( + *, + k, + v, + start, + stop, + njts, + storage=storage, + non_blocking=non_blocking, + ): + """Reads a slice of the storage and assigns the resulting tensor in flat_dict.""" + # v may need padding + if k[-1].startswith(""): + njts[k] = v + return + v_pad = v.view(-1).view(torch.uint8) + exp_length = stop - start + pad = exp_length - v_pad.numel() + if pad: + v_pad = torch.cat([v_pad, v_pad.new_zeros(pad)]) + storage[start:stop].copy_(v_pad, non_blocking=non_blocking) + + storage_slice = storage[start:stop] + shape, dtype = v.shape, v.dtype + new_v = storage_slice.view(dtype) + if pad: + new_v = new_v[: v.numel()] + new_v = new_v.view(shape) + flat_dict[k] = new_v + + njts = {} + if num_threads > 1: + executor = ThreadPoolExecutor(num_threads) + r = [] + for i, (k, v) in enumerate(flat_dict.items()): + r.append( + executor.submit( + assign, + k=k, + v=v, + start=offsets[i], + stop=offsets[i + 1], + njts=njts, + ) + ) + if not return_early: + wait(r) + else: + # TODO: We'd need to merge the second half of this function to make this a thing + raise NotImplementedError( + "return_early is not implemented yet for `consolidate`." + ) + else: + for i, (k, v) in enumerate(flat_dict.items()): + assign( + k=k, + v=v, + start=offsets[i], + stop=offsets[i + 1], + njts=njts, + ) + for njt_key, njt in njts.items(): + newkey = njt_key[:-1] + (njt_key[-1].replace("", ""),) + njt_key_values = njt_key[:-1] + ( + njt_key[-1].replace("", ""), + ) + njt_key_offset = njt_key[:-1] + ( + njt_key[-1].replace("", ""), + ) + njt_key_lengths = njt_key[:-1] + ( + njt_key[-1].replace("", ""), + ) + val = _rebuild_njt_from_njt( + njt, + values=flat_dict.pop(njt_key_values), + offsets=flat_dict.pop(njt_key_offset), + lengths=flat_dict.pop(njt_key_lengths, None), + ) + del flat_dict[njt_key] + flat_dict[newkey] = val + + if non_blocking and device.type != "cuda": + # sync if needed + self._sync_all() + else: + + def _view_and_pad(tensor): + result = tensor.view(-1).view(torch.uint8) + # result must always have a multiple of 8 elements + pad = 0 + if need_padding: + pad = result.numel() % 8 + if pad != 0: + result = torch.cat([result, result.new_zeros(8 - pad)]) + return result, pad + + items = [] + for v in flat_dict.values(): + if v.is_nested: + continue + if v.device != storage.device: + v = v.to(storage.device, non_blocking=non_blocking) + stride = v.stride() + if is_compiling(): + if not v.is_contiguous(): + v = v.clone(memory_format=torch.contiguous_format) + elif (stride and stride[-1] != 1) or v.storage_offset(): + v = v.clone(memory_format=torch.contiguous_format) + v, pad = _view_and_pad(v) + items.append(v) + if non_blocking and device.type != "cuda": + # sync if needed + self._sync_all() + if items: + torch.cat(items, out=storage) + for v, (k, oldv) in _zip_strict( + storage.split(flat_size), list(flat_dict.items()) + ): + if not k[-1].startswith("<"): + flat_dict[k] = view_old_as_new(v, oldv) + elif k[-1].startswith(""): + # NJT/NT always comes before offsets/shapes + nt = oldv + nt_lengths = None + del flat_dict[k] + elif k[-1].startswith(""): + nt_vaues = view_old_as_new(v, oldv) + del flat_dict[k] + elif k[-1].startswith(""): + nt_lengths = view_old_as_new(v, oldv) + del flat_dict[k] + elif k[-1].startswith(""): + newk = k[:-1] + (k[-1].replace("", ""),) + nt_offsets = view_old_as_new(v, oldv) + del flat_dict[k] + + val = _rebuild_njt_from_njt( + nt, values=nt_vaues, offsets=nt_offsets, lengths=nt_lengths + ) + + flat_dict[newk] = val + + # delete the nested value to make sure that if there was an + # ordering mismatch we wouldn't be looking at the value key of + # another nested tensor. + del nt, nt_vaues, nt_offsets, nt_lengths + else: + flat_dict[k] = view_old_as_new(v, oldv) + + def assign_val(key, val): + if isinstance(key, str): + key = (key,) + return flat_dict.get(key, val) + + if filename is None: + device = self.device + elif not inplace: + device = torch.device("cpu") + elif self.device is not None and self.device != torch.device("cpu"): + self.clear_device_() + device = None + else: + device = None + result = self._fast_apply( + assign_val, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS_NONTENSOR, + out=self if inplace else None, + device=device, + ) + result._consolidated = {"storage": storage, "metadata": metadata_dict} + # Lock the consolidated TensorDict to prevent modifications that could break consolidation + result.lock_() + if filename is not None: + if use_buffer: + with open(filename, "w+b") as f: + f.write(total_storage._handler.buffer) + # with open(Path(filename).with_suffix(".json"), "wb") as f: + # metadata_dict["size"] = filesize + # f.write(json.dumps(metadata_dict)) + return result + + @classmethod + def from_consolidated(cls, filename): + # with open(Path(filename).with_suffix(".json"), "rb") as f: + # metadata = json.loads(f.read()) + file = torch.from_file( + str(filename), + dtype=torch.uint8, + size=os.path.getsize(filename), + # needed when device ctx differs + device=torch.device("cpu"), + ) + metadata_size = file[-8:].clone().view(torch.int64) + metadata = file[-metadata_size - 8 : -8] + metadata = json.loads(bytes(metadata.tolist())) + + from ._reductions import _rebuild_tensordict_files_consolidated + + return _rebuild_tensordict_files_consolidated( + metadata, file[: -metadata_size - 8] + ) + + def is_consolidated(self): + """Checks if a TensorDict has a consolidated storage.""" + return hasattr(self, "_consolidated") + + to_mds = to_mds + + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: + """Writes all tensors onto a corresponding memory-mapped Tensor, in-place. + + Args: + prefix (str): directory prefix where the memory-mapped tensors will + be stored. The directory tree structure will mimic the tensordict's. + copy_existing (bool): If False (default), an exception will be raised if an + entry in the tensordict is already a tensor stored on disk + with an associated file, but is not saved in the correct + location according to prefix. + If ``True``, any existing Tensor will be copied to the new location. + + Keyword Args: + num_threads (int, optional): the number of threads used to write the memmap + tensors. Defaults to `0`. + return_early (bool, optional): if ``True`` and ``num_threads>0``, + the method will return a future of the tensordict. The resulting + tensordict can be queried using `future.result()`. + share_non_tensor (bool, optional): if ``True``, the non-tensor data will be + shared between the processes and writing operation (such as inplace update + or set) on any of the workers within a single node will update the value + on all other workers. If the number of non-tensor leaves is high (e.g., + sharing large stacks of non-tensor data) this may result in OOM or similar + errors. Defaults to ``False``. + existsok (bool, optional): if ``False``, an exception will be raised if a tensor already + exists in the same path. Defaults to ``True``. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + The TensorDict is then locked, meaning that any writing operations that + isn't in-place will throw an exception (eg, rename, set or remove an + entry). + Once the tensordict is unlocked, the memory-mapped attribute is turned to ``False``, + because cross-process identity is not guaranteed anymore. + + Returns: + self if ``return_early=False``, otherwise a :class:`~tensordict.utils.TensorDictFuture` instance. + + Note: + Serialising in this fashion might be slow with deeply nested tensordicts, so + it is not recommended to call this method inside a training loop. + """ + prefix = Path(prefix) if prefix is not None else self._memmap_prefix + if num_threads > 1: + with ( + ThreadPoolExecutor(max_workers=num_threads) + if not return_early + else contextlib.nullcontext() + ) as executor: + if return_early: + executor = ThreadPoolExecutor(max_workers=num_threads) + futures = [] + result = self._memmap_( + prefix=prefix, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=True, + like=False, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if not return_early: + concurrent.futures.wait(futures) + return result + else: + return TensorDictFuture(futures, result) + return self._memmap_( + prefix=prefix, + copy_existing=copy_existing, + inplace=True, + futures=None, + executor=None, + like=False, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ).lock_() + + @abc.abstractmethod + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + """Creates an empty memory-mapped tensor given a shape and possibly a dtype. + + .. warning:: + This method is not lock-safe by design. A memory-mapped TensorDict instance present on multiple nodes + will need to be updated using the method :meth:`~tensordict.TensorDictBase.memmap_refresh_`. + + Writing an existing entry will result in an error. + + Args: + key (NestedKey): the key of the new entry to write. If the key is already present in the tensordict, an + exception is raised. + shape (torch.Size or equivalent, torch.Tensor for nested tensors): the shape of the tensor to write. + + Keyword arguments: + dtype (torch.dtype, optional): the dtype of the new tensor. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + Returns: + A new memory mapped tensor. + + """ + raise NotImplementedError + + @abc.abstractmethod + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + """Creates an empty memory-mapped tensor given a storage, a shape and possibly a dtype. + + .. warning:: + This method is not lock-safe by design. A memory-mapped TensorDict instance present on multiple nodes + will need to be updated using the method :meth:`~tensordict.TensorDictBase.memmap_refresh_`. + + .. note:: + If the storage has a filename associated, it must match the new filename for the file. + If it has not a filename associated but the tensordict has an associated path, this will result in an + exception. + + Args: + key (NestedKey): the key of the new entry to write. If the key is already present in the tensordict, an + exception is raised. + storage (torch.UntypedStorage): the storage to use for the new MemoryMappedTensor. Must be a physical memory + storage. + shape (torch.Size or equivalent, torch.Tensor for nested tensors): the shape of the tensor to write. + + Keyword arguments: + dtype (torch.dtype, optional): the dtype of the new tensor. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + Returns: + A new memory mapped tensor with the given storage. + + """ + raise NotImplementedError + + @abc.abstractmethod + def make_memmap_from_tensor( + self, + key: NestedKey, + tensor: torch.Tensor, + *, + copy_data: bool = True, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + """Creates an empty memory-mapped tensor given a tensor. + + .. warning:: + This method is not lock-safe by design. A memory-mapped TensorDict instance present on multiple nodes + will need to be updated using the method :meth:`~tensordict.TensorDictBase.memmap_refresh_`. + + This method always copies the storage content if ``copy_data`` is ``True`` (i.e., the storage is not shared). + + Args: + key (NestedKey): the key of the new entry to write. If the key is already present in the tensordict, an + exception is raised. + tensor (torch.Tensor): the tensor to replicate on physical memory. + + Keyword arguments: + copy_data (bool, optionaL): if ``False``, the new tensor will share the metadata of the input such as + shape and dtype, but the content will be empty. Defaults to ``True``. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + Returns: + A new memory mapped tensor with the given storage. + + """ + raise NotImplementedError + + def save( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: + """Saves the tensordict to disk. + + This function is a proxy to :meth:`~.memmap`. + """ + return self.memmap( + prefix=prefix, + copy_existing=copy_existing, + num_threads=num_threads, + return_early=return_early, + share_non_tensor=share_non_tensor, + robust_key=robust_key, + ) + + dumps = save + + def memmap( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: + """Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict. + + Args: + prefix (str): directory prefix where the memory-mapped tensors will + be stored. The directory tree structure will mimic the tensordict's. + copy_existing (bool): If False (default), an exception will be raised if an + entry in the tensordict is already a tensor stored on disk + with an associated file, but is not saved in the correct + location according to prefix. + If ``True``, any existing Tensor will be copied to the new location. + + Keyword Args: + num_threads (int, optional): the number of threads used to write the memmap + tensors. Defaults to `0`. + return_early (bool, optional): if ``True`` and ``num_threads>0``, + the method will return a future of the tensordict. + share_non_tensor (bool, optional): if ``True``, the non-tensor data will be + shared between the processes and writing operation (such as inplace update + or set) on any of the workers within a single node will update the value + on all other workers. If the number of non_tensor leaves is high (e.g., + sharing large stacks of non-tensor data) this may result in OOM or similar + errors. Defaults to ``False``. + existsok (bool, optional): if ``False``, an exception will be raised if a tensor already + exists in the same path. Defaults to ``True``. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + The TensorDict is then locked, meaning that any writing operations that + isn't in-place will throw an exception (eg, rename, set or remove an + entry). + Once the tensordict is unlocked, the memory-mapped attribute is turned to ``False``, + because cross-process identity is not guaranteed anymore. + + Returns: + A new tensordict with the tensors stored on disk if ``return_early=False``, + otherwise a :class:`~tensordict.utils.TensorDictFuture` instance. + + Note: + Serialising in this fashion might be slow with deeply nested tensordicts, so + it is not recommended to call this method inside a training loop. + """ + prefix = Path(prefix) if prefix is not None else self._memmap_prefix + + if num_threads > 1: + with ( + ThreadPoolExecutor(max_workers=num_threads) + if not return_early + else contextlib.nullcontext() + ) as executor: + if return_early: + executor = ThreadPoolExecutor(max_workers=num_threads) + futures = [] + result = self._memmap_( + prefix=prefix, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=False, + like=False, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if not return_early: + concurrent.futures.wait(futures) + return result + else: + return TensorDictFuture(futures, result) + + return self._memmap_( + prefix=prefix, + copy_existing=copy_existing, + inplace=False, + executor=None, + like=False, + futures=None, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ).lock_() + + def memmap_like( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + existsok: bool = True, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: + """Creates a contentless Memory-mapped tensordict with the same shapes as the original one. + + Args: + prefix (str): directory prefix where the memory-mapped tensors will + be stored. The directory tree structure will mimic the tensordict's. + copy_existing (bool): If False (default), an exception will be raised if an + entry in the tensordict is already a tensor stored on disk + with an associated file, but is not saved in the correct + location according to prefix. + If ``True``, any existing Tensor will be copied to the new location. + + Keyword Args: + num_threads (int, optional): the number of threads used to write the memmap + tensors. Defaults to `0`. + return_early (bool, optional): if ``True`` and ``num_threads>0``, + the method will return a future of the tensordict. + share_non_tensor (bool, optional): if ``True``, the non-tensor data will be + shared between the processes and writing operation (such as inplace update + or set) on any of the workers within a single node will update the value + on all other workers. If the number of non-tensor leaves is high (e.g., + sharing large stacks of non-tensor data) this may result in OOM or similar + errors. Defaults to ``False``. + existsok (bool, optional): if ``False``, an exception will be raised if a tensor already + exists in the same path. Defaults to ``True``. + robust_key (bool, optional): if ``True``, uses robust key encoding that safely + handles keys with path separators and special characters. If ``False``, + uses legacy behavior (keys used as-is). If ``None`` (default), emits a + deprecation warning and falls back to legacy behavior. Will default to + ``True`` in v0.12. + + The TensorDict is then locked, meaning that any writing operations that + isn't in-place will throw an exception (eg, rename, set or remove an + entry). + Once the tensordict is unlocked, the memory-mapped attribute is turned to ``False``, + because cross-process identity is not guaranteed anymore. + + Returns: + A new ``TensorDict`` instance with data stored as memory-mapped tensors if ``return_early=False``, + otherwise a :class:`~tensordict.utils.TensorDictFuture` instance. + + .. note:: + This is the recommended method to write a set of large buffers + on disk, as :meth:`~.memmap_()` will copy the information, which can + be slow for large content. + + Examples: + >>> td = TensorDict({ + ... "a": torch.zeros((3, 64, 64), dtype=torch.uint8), + ... "b": torch.zeros(1, dtype=torch.int64), + ... }, batch_size=[]).expand(1_000_000) # expand does not allocate new memory + >>> buffer = td.memmap_like("/path/to/dataset") + + """ + prefix = Path(prefix) if prefix is not None else self._memmap_prefix + if num_threads > 1: + with ( + ThreadPoolExecutor(max_workers=num_threads) + if not return_early + else contextlib.nullcontext() + ) as executor: + if return_early: + executor = ThreadPoolExecutor(max_workers=num_threads) + futures = [] + + # we create an empty copy of self + # This is because calling MMapTensor.from_tensor(mmap_tensor) does nothing + # if both are in filesystem + def empty(x): + return torch.empty((), device=x.device, dtype=x.dtype).expand( + x.shape + ) + + input = self.apply(empty) + result = input._memmap_( + prefix=prefix, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=False, + like=True, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if not return_early: + concurrent.futures.wait(futures) + return result + else: + return TensorDictFuture(futures, result) + + def empty_expand(x): + return torch.empty((), device=x.device, dtype=x.dtype).expand(x.shape) + + input = self.apply(empty_expand) + return input._memmap_( + prefix=prefix, + copy_existing=copy_existing, + inplace=False, + like=True, + executor=None, + futures=None, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ).lock_() + + @classmethod + def load(cls, prefix: str | Path, *args, **kwargs) -> Self: + """Loads a tensordict from disk. + + This class method is a proxy to :meth:`~.load_memmap`. + """ + return cls.load_memmap(prefix, *args, **kwargs) + + def load_(self, prefix: str | Path, *args, **kwargs): + """Loads a tensordict from disk within the current tensordict. + + This class method is a proxy to :meth:`~.load_memmap_`. + """ + return self.load_memmap_(prefix, *args, **kwargs) + + @classmethod + def load_memmap( + cls, + prefix: str | Path, + device: torch.device | None = None, + non_blocking: bool = False, + *, + out: TensorDictBase | None = None, + robust_key: bool | None = None, + ) -> Self: + """Loads a memory-mapped tensordict from disk. + + Args: + prefix (str or Path to folder): the path to the folder where the + saved tensordict should be fetched. + device (torch.device or equivalent, optional): if provided, the + data will be asynchronously cast to that device. + Supports `"meta"` device, in which case the data isn't loaded + but a set of empty "meta" tensors are created. This is + useful to get a sense of the total model size and structure + without actually opening any file. + non_blocking (bool, optional): if ``True``, synchronize won't be + called after loading tensors on device. Defaults to ``False``. + out (TensorDictBase, optional): optional tensordict where the data + should be written. + robust_key (bool, optional): if ``True``, expects robust key encoding was used + when saving and decodes filenames accordingly. If ``False``, uses legacy + behavior. If ``None`` (default), emits a deprecation warning and falls + back to legacy behavior. Will default to ``True`` in v0.12. + + Examples: + >>> from tensordict import TensorDict + >>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0) + >>> td.memmap("./saved_td") + >>> td_load = TensorDict.load_memmap("./saved_td") + >>> assert (td == td_load).all() + + This method also allows loading nested tensordicts. + + Examples: + >>> nested = TensorDict.load_memmap("./saved_td/nested") + >>> assert nested["e"] == 0 + + A tensordict can also be loaded on "meta" device or, alternatively, + as a fake tensor. + + Examples: + >>> import tempfile + >>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) + >>> with tempfile.TemporaryDirectory() as path: + ... td.save(path) + ... td_load = TensorDict.load_memmap(path, device="meta") + ... print("meta:", td_load) + ... from torch._subclasses import FakeTensorMode + ... with FakeTensorMode(): + ... td_load = TensorDict.load_memmap(path) + ... print("fake:", td_load) + meta: TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=meta, + is_shared=False)}, + batch_size=torch.Size([]), + device=meta, + is_shared=False) + fake: TensorDict( + fields={ + a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False)}, + batch_size=torch.Size([]), + device=cpu, + is_shared=False) + + """ + prefix = Path(prefix) + + metadata = _load_metadata(prefix) + type_name = metadata["_type"] + if type_name != str(cls): + import tensordict + + for other_cls in tensordict.base._ACCEPTED_CLASSES: + if str(other_cls) == type_name: + return other_cls._load_memmap( + prefix, metadata, robust_key=robust_key + ) + else: + raise RuntimeError( + f"Could not find name {type_name} in {tensordict.base._ACCEPTED_CLASSES}. " + f"Did you call _register_tensor_class(cls) on {type_name}?" + ) + if device is not None: + device = torch.device(device) + out = cls._load_memmap( + prefix, metadata, device=device, out=out, robust_key=robust_key + ) + if ( + not non_blocking + and device is not None + and device.type not in ("meta", "cuda") + ): + out._sync_all() + return out + + def load_memmap_( + self, + prefix: str | Path, + robust_key: bool | None = None, + ): + """Loads the content of a memory-mapped tensordict within the tensordict where ``load_memmap_`` is called. + + See :meth:`~tensordict.TensorDictBase.load_memmap` for more info. + """ + is_memmap = self.is_memmap() + with self.unlock_() if is_memmap else contextlib.nullcontext(): + self.load_memmap( + prefix=prefix, device=self.device, out=self, robust_key=robust_key + ) + if is_memmap: + self.memmap_() + return self + + def memmap_refresh_(self): + """Refreshes the content of the memory-mapped tensordict if it has a :attr:`~tensordict.TensorDict.saved_path`. + + This method will raise an exception if no path is associated with it. + + """ + if not self.is_memmap() or self._memmap_prefix is None: + raise RuntimeError( + "Cannot refresh a TensorDict that is not memory mapped or has no path associated." + ) + return self.load_memmap_(prefix=self.saved_path) + + @classmethod + @abc.abstractmethod + def _load_memmap( + cls, + prefix: Path, + metadata: dict, + device: torch.device | None = None, + *, + robust_key, + out=None, + ): + raise NotImplementedError + + # Key functionality: set, get, set_, set_at_, update, update_ + @abc.abstractmethod + def entry_class(self, key: NestedKey) -> type: + """Returns the class of an entry, possibly avoiding a call to `isinstance(td.get(key), type)`. + + This method should be preferred to ``tensordict.get(key).shape`` whenever + :meth:`.get` can be expensive to execute. + + """ + raise NotImplementedError + + def set( + self, + key: NestedKey, + item: CompatibleType, + inplace: bool = False, + *, + non_blocking: bool = False, + **kwargs: Any, + ) -> Self: + """Sets a new key-value pair. + + Args: + key (str, tuple of str): name of the key to be set. + item (torch.Tensor or equivalent, TensorDictBase instance): value + to be stored in the tensordict. + inplace (bool, optional): if ``True`` and if a key matches an existing + key in the tensordict, then the update will occur in-place + for that key-value pair. If inplace is ``True`` and + the entry cannot be found, it will be added. For a more restrictive + in-place operation, use :meth:`~.set_` instead. + Defaults to ``False``. + + Keyword Args: + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Returns: + self + + Examples: + >>> td = TensorDict({}, batch_size[3, 4]) + >>> td.set("x", torch.randn(3, 4)) + >>> y = torch.randn(3, 4, 5) + >>> td.set("y", y, inplace=True) # works, even if 'y' is not present yet + >>> td.set("y", torch.zeros_like(y), inplace=True) + >>> assert (y==0).all() # y values are overwritten + >>> td.set("y", torch.ones(5), inplace=True) # raises an exception as shapes mismatch + + """ + key = _unravel_key_to_tuple(key) + # inplace is loose here, but for set_ it is constraining. We translate it + # to None to tell _set_str and others to drop it if the key isn't found + inplace = BEST_ATTEMPT_INPLACE if inplace else False + return self._set_tuple( + key, item, inplace=inplace, validated=False, non_blocking=non_blocking + ) + + @abc.abstractmethod + def _set_str( + self, + key: str, + value: Any, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ): + raise NotImplementedError + + @abc.abstractmethod + def _set_tuple(self, key, value, *, inplace, validated, non_blocking: bool): + raise NotImplementedError + + @lock_blocked + def set_non_tensor(self, key: NestedKey, value: Any): + """Registers a non-tensor value in the tensordict using :class:`tensordict.tensorclass.NonTensorData`. + + The value can be retrieved using :meth:`TensorDictBase.get_non_tensor` + or directly using `get`, which will return the :class:`tensordict.tensorclass.NonTensorData` + object. + + return: self + + Examples: + >>> data = TensorDict({}, batch_size=[]) + >>> data.set_non_tensor(("nested", "the string"), "a string!") + >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" + >>> # regular `get` works but returns a NonTensorData object + >>> data.get(("nested", "the string")) + NonTensorData( + data='a string!', + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + key = unravel_key(key) + return self._set_non_tensor(key, value) + + def _set_non_tensor(self, key: NestedKey, value: Any): + if isinstance(key, tuple): + if len(key) == 1: + return self._set_non_tensor(key[0], value) + sub_td = self._get_str(key[0], None) + if sub_td is None: + sub_td = self._create_nested_str(key[0]) + sub_td._set_non_tensor(key[1:], value) + return self + from tensordict.tensorclass import NonTensorData + + self._set_str( + key, + NonTensorData( + data=value, + batch_size=self.batch_size, + device=self.device, + names=self._maybe_names(), + ), + validated=True, + inplace=False, + non_blocking=False, + ) + return self + + def get_non_tensor(self, key: NestedKey, default=NO_DEFAULT): + """Gets a non-tensor value, if it exists, or `default` if the non-tensor value is not found. + + This method is robust to tensor/TensorDict values, meaning that if the + value gathered is a regular tensor it will be returned too (although + this method comes with some overhead and should not be used out of its + natural scope). + + See :meth:`~tensordict.TensorDictBase.set_non_tensor` for more information + on how to set non-tensor values in a tensordict. + + Args: + key (NestedKey): the location of the NonTensorData object. + default (Any, optional): the value to be returned if the key cannot + be found. + + Returns: the content of the :class:`tensordict.tensorclass.NonTensorData`, + or the entry corresponding to the ``key`` if it isn't a + :class:`tensordict.tensorclass.NonTensorData` (or ``default`` if the + entry cannot be found). + + Examples: + >>> data = TensorDict({}, batch_size=[]) + >>> data.set_non_tensor(("nested", "the string"), "a string!") + >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" + >>> # regular `get` works but returns a NonTensorData object + >>> data.get(("nested", "the string")) + NonTensorData( + data='a string!', + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + key = unravel_key(key) + return self._get_non_tensor(key, default=default) + + def _get_non_tensor(self, key: NestedKey, default=NO_DEFAULT): + if isinstance(key, tuple): + if len(key) == 1: + return self._get_non_tensor(key[0], default=default) + subtd = self._get_str(key[0], default=default) + if subtd is default: + return subtd + return subtd._get_non_tensor(key[1:], default=default) + value = self._get_str(key, default=default) + + if is_non_tensor(value): + from tensordict import NonTensorStack + + if isinstance(value, NonTensorStack) and not capture_non_tensor_stack(): + return value.tolist(as_linked_list=True) + data = getattr(value, "data", None) + if data is None: + return value.tolist(as_linked_list=True) + return data + return value + + def filter_non_tensor_data(self) -> Self: + """Filters out all non-tensor-data.""" + + def _filter(x): + if not is_non_tensor(x): + if is_tensor_collection(x): + return x.filter_non_tensor_data() + return x + + return self._apply_nest(_filter, call_on_nested=True, filter_empty=False) + + def filter_empty_(self): + """Filters out all empty tensordicts in-place.""" + for key, val in reversed( + list(self.items(True, is_leaf=_NESTED_TENSORS_AS_LISTS, sort=True)) + ): + if _is_tensor_collection(type(val)) and val.is_empty(): + del self[key] + return self + + def _convert_inplace(self, inplace, key): + if inplace is not False: + has_key = key in self.keys() + if inplace is True and not has_key: # inplace could be None + raise KeyError( + _KEY_ERROR.format(key, type(self).__name__, sorted(self.keys())) + ) + inplace = has_key + return inplace + + def set_at_( + self, + key: NestedKey, + value: CompatibleType, + index: IndexType, + *, + non_blocking: bool = False, + ) -> Self: + """Sets the values in-place at the index indicated by ``index``. + + Args: + key (str, tuple of str): key to be modified. + value (torch.Tensor): value to be set at the index `index` + index (int, tensor or tuple): index where to write the values. + + Keyword Args: + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Returns: + self + + Examples: + >>> td = TensorDict({}, batch_size[3, 4]) + >>> x = torch.randn(3, 4) + >>> td.set("x", x) + >>> td.set_at_("x", value=torch.ones(1, 4), index=slice(1)) + >>> assert (x[0] == 1).all() + """ + key = _unravel_key_to_tuple(key) + return self._set_at_tuple( + key, value, index, validated=False, non_blocking=non_blocking + ) + + @abc.abstractmethod + def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool): + raise NotImplementedError + + @abc.abstractmethod + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool): + raise NotImplementedError + + def set_( + self, + key: NestedKey, + item: CompatibleType, + *, + non_blocking: bool = False, + ) -> Self: + """Sets a value to an existing key while keeping the original storage. + + Args: + key (str): name of the value + item (torch.Tensor or compatible type, TensorDictBase): value to + be stored in the tensordict + + Keyword Args: + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Returns: + self + + Examples: + >>> td = TensorDict({}, batch_size[3, 4]) + >>> x = torch.randn(3, 4) + >>> td.set("x", x) + >>> td.set_("x", torch.zeros_like(x)) + >>> assert (x == 0).all() + + """ + key = _unravel_key_to_tuple(key) + return self._set_tuple( + key, item, inplace=True, validated=False, non_blocking=non_blocking + ) + + # Stack functionality + @abc.abstractmethod + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: + """Stacks a list of values onto an existing key while keeping the original storage. + + Args: + key (str): name of the value + list_item (list of torch.Tensor): value to be stacked and stored in the tensordict. + dim (int): dimension along which the tensors should be stacked. + + Returns: + self + + """ + raise NotImplementedError + + def _stack_onto_at_( + self, + key: NestedKey, + list_item: list[CompatibleType], + dim: int, + idx: IndexType, + ) -> Self: + """Similar to _stack_onto_ but on a specific index. Only works with regular TensorDicts.""" + raise RuntimeError( + f"Cannot call _stack_onto_at_ with {type(self).__name__}. " + "Make sure your sub-classed tensordicts are turned into regular tensordicts by calling to_tensordict() " + "before calling __getindex__ and stack." + ) + + def _default_get(self, key: NestedKey, default: Any = NO_DEFAULT) -> CompatibleType: + if default is not NO_DEFAULT: + return default + else: + # raise KeyError + raise KeyError( + _KEY_ERROR.format(key, type(self).__name__, sorted(self.keys())) + ) + + @overload + def get(self, key): ... + @overload + def get(self, key, default): ... + + def get(self, key: NestedKey, *args, **kwargs) -> CompatibleType: + """Gets the value stored with the input key. + + Args: + key (str, tuple of str): key to be queried. If tuple of str it is + equivalent to chained calls of getattr. + default: default value if the key is not found in the tensordict. Defaults to ``None``. + + .. warning:: + Previously, if a key was not present in the tensordict and no default + was passed, a `KeyError` was raised. From v0.7, this behaviour has been changed + and a `None` value is returned instead (in accordance with the what dict.get behavior). + To adopt the old behavior, set the environment variable `export TD_GET_DEFAULTS_TO_NONE='0'` or call + :func`~tensordict.set_get_defaults_to_none(False)`. + + .. note:: Keyword arguments can be passed to :meth:`~.get` when dealing with ragged tensors. + See :meth:`~tensordict.LazyStackedTensorDict.get` for a complete overview. + + Examples: + >>> td = TensorDict({"x": 1}, batch_size=[]) + >>> td.get("x") + tensor(1) + >>> td.get("y") + None + """ + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + # Find what the default is + if args: + default = args[0] + if len(args) > 1: + raise TypeError("Only one arg is allowed in TD.get.") + elif "default" in kwargs: + raise TypeError("'default' arg was passed twice.") + elif "default" in kwargs: + default = kwargs.pop("default") + if args: + raise TypeError("'default' arg was passed twice.") + elif _GET_DEFAULTS_TO_NONE: + default = None + else: + default = NO_DEFAULT + return self._get_tuple(key, default=default, **kwargs) + + @abc.abstractmethod + def _get_str(self, key, default, **kwargs): + raise NotImplementedError + + @abc.abstractmethod + def _get_tuple(self, key, default, **kwargs): + raise NotImplementedError + + def _get_tuple_maybe_non_tensor(self, key, default, **kwargs): + result = self._get_tuple(key, default, **kwargs) + if _pass_through(result): + # Only lazy stacks of non tensors are actually tensordict instances + if isinstance(result, TensorDictBase): + return result.tolist(as_linked_list=True) + return result.data + return result + + @overload + def get_at(self, key, index): ... + + @overload + def get_at(self, key, index, default): ... + + def get_at( + self, + key: NestedKey, + *args, + **kwargs, + ) -> CompatibleType: + """Get the value of a tensordict from the key `key` at the index `idx`. + + Args: + key (str, tuple of str): key to be retrieved. + index (int, slice, torch.Tensor, iterable): index of the tensor. + default (torch.Tensor): default value to return if the key is + not present in the tensordict. + + Returns: + indexed tensor. + + Examples: + >>> td = TensorDict({"x": torch.arange(3)}, batch_size=[]) + >>> td.get_at("x", index=1) + tensor(1) + + """ + # TODO: check that this works with masks, and add to docstring + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + + try: + if len(args): + index = args[0] + args = args[1:] + else: + index = kwargs.pop("index") + except KeyError: + raise TypeError("index argument missing from get_at") + + # Find what the default is + if args: + default = args[0] + if len(args) > 1: + raise TypeError("only one (keyword) argument is allowed.") + elif "default" in kwargs: + default = kwargs.pop("default") + elif _GET_DEFAULTS_TO_NONE: + default = None + else: + default = NO_DEFAULT + + return self._get_at_tuple(key, index, default, **kwargs) + + def _get_at_str(self, key, idx, default, **kwargs): + out = self._get_str(key, default, **kwargs) + if out is default: + return out + return out[idx] + + def _get_at_tuple(self, key, idx, default, **kwargs): + out = self._get_tuple(key, default, **kwargs) + if out is default: + return out + return out[idx] + + def get_item_shape(self, key: NestedKey): + """Returns the shape of the entry, possibly avoiding recurring to :meth:`~.get`.""" + return _shape(self.get(key)) + + @lock_blocked + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> Self: + """Updates the TensorDict with values from either a dictionary or another TensorDict. + + .. warning:: `update` will corrupt the data if called within a try/except block. Do not user this method within + such blocks hoping to catch and patch errors that occur during the execution. + + Args: + input_dict_or_td (TensorDictBase or dict): input data to be written + in self. + clone (bool, optional): whether the tensors in the input ( + tensor) dict should be cloned before being set. + Defaults to ``False``. + inplace (bool, optional): if ``True`` and if a key matches an existing + key in the tensordict, then the update will occur in-place + for that key-value pair. If the entry cannot be found, it will be + added. Defaults to ``False``. + + Keyword Args: + keys_to_update (sequence of NestedKeys, optional): if provided, only + the list of keys in ``key_to_update`` will be updated. + This is aimed at avoiding calls to + ``data_dest.update(data_src.select(*keys_to_update))``. + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + is_leaf (Callable[[Type], bool], optional): a callable that indicates + whether an object type is to be considered a leaf and swapped + or a tensor collection. + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`. + + update_batch_size (bool, optional): if ``True``, ``update`` will attempt to update the batch-size + of the destination (`self`) if it mismatches the source's batch size. Defaults to ``False``. + + .. note:: In cases where the batch size does not match, :class:`~tensordict.LazyStackTensorDict` + instances will be emptied of their content and copies of the tensordicts from the source + will be used to repopulate the container. + + .. note:: This argument assumes that `keys_to_update` is left empty, and that `inplace=False`. + If the keys of the destination (`self`) is not a subset of the keys of the source, + an exception will be raised, as TensorDict will be unable to infer what to do with the extra + destination entries. + + ignore_lock (bool, optional): if ``True``, any tensordict can be updated regardless of its locked status. + Defaults to `False`. + + .. note:: When updating a :class:`~tensordict.LazyStackedTensorDict` with N elements with another + :class:`~tensordict.LazyStackedTensorDict` with M elements, with M > N, along the stack dimension, + the ``update`` method will append copies of the extra tensordicts to the dest (self) lazy stack. + This allows users to rely on ``update`` to increment lazy stacks progressively. + + Returns: + self + + Examples: + >>> td = TensorDict({}, batch_size=[3]) + >>> a = torch.randn(3) + >>> b = torch.randn(3, 4) + >>> other_td = TensorDict({"a": a, "b": b}, batch_size=[]) + >>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found + >>> assert td['a'] is other_td['a'] + >>> other_td = other_td.clone().zero_() + >>> td.update(other_td) + >>> assert td['a'] is not other_td['a'] + + """ + batch_size_changed = False + if input_dict_or_td is self: + # no op + return self + if is_leaf is None: + is_leaf = _is_leaf_nontensor + if keys_to_update is not None: + if len(keys_to_update) == 0: + return self + keys_to_update = unravel_key_list(keys_to_update) + + if ( + _is_tensor_collection(type(input_dict_or_td)) + # here, we do a loose check on the batch sizes: it could be that the source has batch_size (1,) and self (1, 2) + # and that all the values have an appropriate shape for the new batch size. + # What we want to catch is any batch size item that is obviously mismatching, like (1, 2, 3) and (1, 4, 3). + and self.batch_size[: input_dict_or_td.batch_dims] + != input_dict_or_td.batch_size[: self.batch_dims] + ): + if not update_batch_size: + raise RuntimeError( + "update_batch_size must be set to True to be able to update " + "tensordicts of different batch size. Got sizes {}" + ) + if inplace: + raise RuntimeError( + "Source and destination tensor collection shapes mismatch, but " + "update was called with inplace=True, which cannot be achieved." + ) + if keys_to_update is not None: + raise RuntimeError( + "Updating tensordicts of different batch-size with keys_to_update " + "is currently not supported." + ) + # This could be expensive but we must run it + keys_source = set(input_dict_or_td.keys(True)) + keys_dest = set(self.keys(True)) + if not keys_dest.issubset(keys_source): + raise RuntimeError( + "Some keys of the dest tensordict are not present in the source " + "during update with mismatching batch-size. " + f"batch_size of source={input_dict_or_td.batch_size}, batch_size of dest={self.batch_size}, " + f"keys in dest but not in source: {{{keys_dest - keys_source}}}." + ) + # We can swap target with value if the batch sizes are incongruent. We must make sure the id of target + # stays the same though + self.batch_size = () + # Remove all leaves, and update + ks = self.keys(True, True, is_leaf=is_leaf) + self.exclude(*ks, inplace=True) + self.batch_size = input_dict_or_td.batch_size + self.update(input_dict_or_td, update_batch_size=True) + return self + + for key, value in input_dict_or_td.items(): + key = _unravel_key_to_tuple(key) + firstkey, subkey = key[0], key[1:] + if keys_to_update and not any( + firstkey == ktu if isinstance(ktu, str) else firstkey == ktu[0] + for ktu in keys_to_update + ): + continue + target = self._get_str(firstkey, None) + if clone and hasattr(value, "clone"): + value = value.clone() + elif clone: + value = tree_map(torch.clone, value) + # the key must be a string by now. Let's check if it is present + if target is not None: + if not is_leaf(type(target)) and not is_leaf(type(value)): + if subkey: + sub_keys_to_update = _prune_selected_keys( + keys_to_update, firstkey + ) + target.update( + {subkey: value}, + inplace=inplace, + clone=clone, + keys_to_update=sub_keys_to_update, + non_blocking=non_blocking, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + continue + elif isinstance(value, (dict,)) or _is_tensor_collection( + type(value) + ): + from tensordict._lazy import LazyStackedTensorDict + + value_is_lazy_stack = isinstance(value, LazyStackedTensorDict) + target_is_lazy_stack = isinstance(target, LazyStackedTensorDict) + if value_is_lazy_stack and not target_is_lazy_stack: + sub_keys_to_update = _prune_selected_keys( + keys_to_update, firstkey + ) + self._set_tuple( + key, + LazyStackedTensorDict( + *target.unbind(value.stack_dim), + stack_dim=value.stack_dim, + ).update( + value, + inplace=inplace, + clone=clone, + keys_to_update=sub_keys_to_update, + non_blocking=non_blocking, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ), + validated=True, + inplace=False, + non_blocking=non_blocking, + ) + + else: + sub_keys_to_update = _prune_selected_keys( + keys_to_update, firstkey + ) + target.update( + value, + inplace=inplace, + clone=clone, + non_blocking=non_blocking, + keys_to_update=sub_keys_to_update, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + continue + # A tensor collection may still be a leaf so we need to duplicate the logic here + if ( + update_batch_size + and _is_tensor_collection(type(target)) + and type(target) is type(value) + and target.shape != value.shape + ): + batch_size_changed = True + from tensordict._lazy import LazyStackedTensorDict + + # We can swap target with value if the batch sizes are incongruent. We must make sure the id of target + # stays the same though + if isinstance(target, LazyStackedTensorDict): + target.__init__( + *value.unbind(target.stack_dim), + stack_dim=target.stack_dim, + hook_out=target.hook_out, + hook_in=target.hook_in, + stack_dim_name=target._td_dim_name, + ) + else: + target = target.exclude( + *target.keys(True, True, is_leaf=is_leaf), inplace=True + ) + target.update(value, update_batch_size=update_batch_size) + target.batch_size = value.batch_size + continue + + self._set_tuple( + key, + value, + inplace=BEST_ATTEMPT_INPLACE if inplace else False, + validated=False, + non_blocking=non_blocking, + ) + if batch_size_changed: + bd = self.batch_dims + self.batch_size = () + # self.batch_size = () + self.auto_batch_size_(bd, keep_compliant_size=True) + return self + + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: + """Updates the TensorDict in-place with values from either a dictionary or another TensorDict. + + Unlike :meth:`~.update`, this function will throw an error if the key is unknown to ``self``. + + Args: + input_dict_or_td (TensorDictBase or dict): input data to be written + in self. + clone (bool, optional): whether the tensors in the input ( + tensor) dict should be cloned before being set. Defaults to ``False``. + + Keyword Args: + keys_to_update (sequence of NestedKeys, optional): if provided, only + the list of keys in ``key_to_update`` will be updated. + This is aimed at avoiding calls to + ``data_dest.update_(data_src.select(*keys_to_update))``. + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Returns: + self + + Examples: + >>> a = torch.randn(3) + >>> b = torch.randn(3, 4) + >>> td = TensorDict({"a": a, "b": b}, batch_size=[3]) + >>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[]) + >>> td.update_(other_td) + >>> assert td['a'] is not other_td['a'] + >>> assert (td['a'] == other_td['a']).all() + >>> assert (td['a'] == 0).all() + + """ + if input_dict_or_td is self: + # no op + return self + + if not _is_tensor_collection(type(input_dict_or_td)): + from tensordict import TensorDict + + input_dict_or_td = TensorDict.from_dict( + input_dict_or_td, batch_dims=self.batch_dims + ) + + if keys_to_update is not None: + if len(keys_to_update) == 0: + return self + keys_to_update = [_unravel_key_to_tuple(key) for key in keys_to_update] + + named = True + + def inplace_update(name, source, dest): + if source is None: + return None + name = _unravel_key_to_tuple(name) + for key in keys_to_update: + if key == name[: len(key)]: + if dest is None: + raise KeyError( + f"The key {name} was not found in the dest tensordict." + ) + dest.copy_(source, non_blocking=non_blocking) + + else: + # Fastest route using _foreach_copy_ + keys, vals = self._items_list(True, True) + new_keys, other_val = input_dict_or_td._items_list( + True, True, sorting_keys=keys, default="intersection" + ) + if len(new_keys) and _foreach_copy_ is not None: + if len(other_val) != len(vals): + vals = dict(zip(keys, vals)) + vals = [vals[k] for k in new_keys] + _foreach_copy_(vals, other_val, non_blocking=non_blocking) + return self + named = True + + def inplace_update(name, source, dest): + if source is None: + return None + if dest is None: + raise KeyError( + f"The key {name} was not found in the dest tensordict." + ) + dest.copy_(source, non_blocking=non_blocking) + + input_dict_or_td._apply_nest( + inplace_update, + self, + nested_keys=True, + default=None, + filter_empty=True, + named=named, + is_leaf=_is_leaf_nontensor, + ) + return self + + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + idx: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: + """Updates the TensorDict in-place at the specified index with values from either a dictionary or another TensorDict. + + Unlike TensorDict.update, this function will throw an error if the key is unknown to the TensorDict. + + Args: + input_dict_or_td (TensorDictBase or dict): input data to be written + in self. + idx (int, torch.Tensor, iterable, slice): index of the tensordict + where the update should occur. + clone (bool, optional): whether the tensors in the input ( + tensor) dict should be cloned before being set. Default is + `False`. + + Keyword Args: + keys_to_update (sequence of NestedKeys, optional): if provided, only + the list of keys in ``key_to_update`` will be updated. + non_blocking (bool, optional): if ``True`` and this copy is between + different devices, the copy may occur asynchronously with respect + to the host. + + Returns: + self + + Examples: + >>> td = TensorDict({ + ... 'a': torch.zeros(3, 4, 5), + ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) + >>> td.update_at_( + ... TensorDict({ + ... 'a': torch.ones(1, 4, 5), + ... 'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]), + ... slice(1, 2)) + TensorDict( + fields={ + a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32), + b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> assert (td[1] == 1).all() + + """ + if idx == (): + return self.update_( + input_dict_or_td=input_dict_or_td, + keys_to_update=keys_to_update, + clone=clone, + non_blocking=non_blocking, + ) + if keys_to_update is not None: + if len(keys_to_update) == 0: + return self + keys_to_update = unravel_key_list(keys_to_update) + for key, value in input_dict_or_td.items(): + firstkey, *nextkeys = _unravel_key_to_tuple(key) + if keys_to_update and not any( + firstkey == ktu if isinstance(ktu, str) else firstkey == ktu[0] + for ktu in keys_to_update + ): + continue + if not isinstance(value, _ACCEPTED_CLASSES): + raise TypeError( + f"Expected value to be one of types {_ACCEPTED_CLASSES} " + f"but got {type(value)}" + ) + if clone: + value = value.clone() + self.set_at_((firstkey, *nextkeys), value, idx, non_blocking=non_blocking) + return self + + def replace(self, *args, **kwargs): + """Creates a shallow copy of the tensordict where entries have been replaced. + + Accepts one unnamed argument which must be a dictionary of a :class:`~tensordict.TensorDictBase` subclass. + Additionally, first-level entries can be updated with the named keyword arguments. + + Returns: + a copy of ``self`` with updated entries if the input is non-empty. If an empty dict or no dict is provided + and the kwargs are empty, ``self`` is returned. + + """ + if args: + if len(args) > 1: + raise RuntimeError( + "Only a single argument containing a dictionary-like " + f"structure of entries to replace can be passed to replace. Received {len(args)} " + f"arguments instead." + ) + dict_to_replace = args[0] + else: + dict_to_replace = {} + if kwargs: + dict_to_replace.update(kwargs) + is_dict = isinstance(dict_to_replace, dict) + if is_dict: + if not dict_to_replace: + return self + else: + if not is_tensor_collection(dict_to_replace): + raise RuntimeError( + f"Cannot use object type {type(dict_to_replace)} to update values in tensordict." + ) + if dict_to_replace.is_empty(): + return self + result = self.copy() + # using update makes sure that any optimization (e.g. for lazy stacks) is done properly + result.update(dict_to_replace) + return result + + @lock_blocked + def create_nested(self, key): + """Creates a nested tensordict of the same shape, device and dim names as the current tensordict. + + If the value already exists, it will be overwritten by this operation. + This operation is blocked in locked tensordicts. + + Examples: + >>> data = TensorDict({}, [3, 4, 5]) + >>> data.create_nested("root") + >>> data.create_nested(("some", "nested", "value")) + >>> print(data) + TensorDict( + fields={ + root: TensorDict( + fields={ + }, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False), + some: TensorDict( + fields={ + nested: TensorDict( + fields={ + value: TensorDict( + fields={ + }, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3, 4, 5]), + device=None, + is_shared=False) + """ + key = _unravel_key_to_tuple(key) + self._create_nested_tuple(key) + return self + + def _create_nested_str(self, key): + out = self.empty() + self._set_str(key, out, inplace=False, validated=True, non_blocking=False) + return out + + def _create_nested_tuple(self, key): + td = self._create_nested_str(key[0]) + if len(key) > 1: + td._create_nested_tuple(key[1:]) + + def copy_(self, tensordict: T, non_blocking: bool = False) -> Self: + """See :obj:`TensorDictBase.update_`. + + The non-blocking argument will be ignored and is just present for + compatibility with :func:`torch.Tensor.copy_`. + """ + return self.update_(tensordict, non_blocking=non_blocking) + + def copy_at_( + self, tensordict: T, idx: IndexType, non_blocking: bool = False + ) -> Self: + """See :obj:`TensorDictBase.update_at_`.""" + return self.update_at_(tensordict, idx, non_blocking=non_blocking) + + def is_empty(self) -> bool: + """Checks if the tensordict contains any leaf.""" + for _ in self.keys(True, True): + return False + return True + + # Dict features: setdefault, items, values, keys, ... + def setdefault( + self, key: NestedKey, default: CompatibleType | Any, inplace: bool = False + ) -> CompatibleType: + """Insert the ``key`` entry with a value of ``default`` if ``key`` is not in the tensordict. + + Return the value for ``key`` if ``key`` is in the tensordict, else ``default``. + + Args: + key (str or nested key): the name of the value. + default (torch.Tensor or compatible type, TensorDictBase): value + to be stored in the tensordict if the key is not already present. + + Returns: + The value of key in the tensordict. Will be default if the key was not + previously set. + + Examples: + >>> td = TensorDict({}, batch_size=[3, 4]) + >>> val = td.setdefault("a", torch.zeros(3, 4)) + >>> assert (val == 0).all() + >>> val = td.setdefault("a", torch.ones(3, 4)) + >>> assert (val == 0).all() # output is still 0 + + """ + if key not in self.keys(include_nested=isinstance(key, tuple)): + self.set(key, default, inplace=inplace) + return self.get(key) + + def items( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf=None, + *, + sort: bool = False, + ) -> Iterator[tuple[str, CompatibleType]]: # noqa: D417 + """Returns a generator of key-value pairs for the tensordict. + + Args: + include_nested (bool, optional): if ``True``, nested values will be returned. + Defaults to ``False``. + leaves_only (bool, optional): if ``False``, only leaves will be + returned. Defaults to ``False``. + is_leaf (callable, optional): a callable over a class type returning + a bool indicating if this class has to be considered as a leaf. + + .. note:: The purpose of `is_leaf` is not to prevent recursive calls into nested tensordicts, but + rather to mark certain types as "leaves" for the purpose of filtering when `leaves_only=True`. + Even if `is_leaf(cls)` returns `True`, the nested structure of the tensordict will still be + traversed if `include_nested=True`. + In other words, `is_leaf` does not control the recursion depth, but rather provides a way to filter + out certain types from the result when `leaves_only=True`. This means that a node in the tree can + be both a leaf and a node with children. + In practice, the default value of ``is_leaf`` does exclude tensordict and tensorclass instances + from the leaf set. + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`. + + Keyword Args: + sort (bool, optional): whether the keys should be sorted. For nested keys, + the keys are sorted according to their joined name (ie, ``("a", "key")`` will + be counted as ``"a.key"`` for sorting). Be mindful that sorting may incur + significant overhead when dealing with large tensordicts. + Defaults to ``False``. + + """ + if sort: + + def keyfunc(item): + return item[0] if isinstance(item[0], str) else ".".join(item[0]) + + yield from sorted( + self.items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + ), + key=keyfunc, + ) + else: + + if is_leaf is None: + is_leaf = _default_is_leaf + + if include_nested: + # check the conditions once only + for k in self.keys(): + val = self._get_str(k, NO_DEFAULT) + cls = type(val) + if not leaves_only or is_leaf(cls): + yield k, val + if _is_tensor_collection(cls): + if not is_non_tensor(cls): + yield from ( + (_unravel_key_to_tuple((k, _key)), _val) + for _key, _val in val.items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + ) + ) + elif leaves_only: + for k in self.keys(): + val = self._get_str(k, NO_DEFAULT) + if is_leaf(type(val)): + yield k, val + else: + for k in self.keys(): + yield k, self._get_str(k, NO_DEFAULT) + + def non_tensor_items(self, include_nested: bool = False): + """Returns all non-tensor leaves, maybe recursively.""" + return tuple( + self.items( + include_nested, + leaves_only=True, + is_leaf=_is_non_tensor, + ) + ) + + def values( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf=None, + *, + sort: bool = False, + ) -> Iterator[CompatibleType]: # noqa: D417 + """Returns a generator representing the values for the tensordict. + + Args: + include_nested (bool, optional): if ``True``, nested values will be returned. + Defaults to ``False``. + leaves_only (bool, optional): if ``False``, only leaves will be + returned. Defaults to ``False``. + is_leaf (callable, optional): a callable over a class type returning + a bool indicating if this class has to be considered as a leaf. + + .. note:: The purpose of `is_leaf` is not to prevent recursive calls into nested tensordicts, but + rather to mark certain types as "leaves" for the purpose of filtering when `leaves_only=True`. + Even if `is_leaf(cls)` returns `True`, the nested structure of the tensordict will still be + traversed if `include_nested=True`. + In other words, `is_leaf` does not control the recursion depth, but rather provides a way to filter + out certain types from the result when `leaves_only=True`. This means that a node in the tree can + be both a leaf and a node with children. + In practice, the default value of ``is_leaf`` does exclude tensordict and tensorclass instances + from the leaf set. + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`. + + Keyword Args: + sort (bool, optional): whether the keys should be sorted. For nested keys, + the keys are sorted according to their joined name (ie, ``("a", "key")`` will + be counted as ``"a.key"`` for sorting). Be mindful that sorting may incur + significant overhead when dealing with large tensordicts. + Defaults to ``False``. + + """ + if sort: + for _, value in self.items(include_nested, leaves_only, is_leaf, sort=sort): + yield value + else: + + if is_leaf is None: + is_leaf = _default_is_leaf + + # check the conditions once only + if include_nested: + for k in self.keys(): + val = self._get_str(k, NO_DEFAULT) + cls = type(val) + if not leaves_only or is_leaf(cls): + yield val + if include_nested and _is_tensor_collection(cls): + if not is_non_tensor(cls): + yield from val.values( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + ) + elif leaves_only: + for k in self.keys(sort=sort): + val = self._get_str(k, NO_DEFAULT) + if is_leaf(type(val)): + yield val + else: + for k in self.keys(sort=sort): + yield self._get_str(k, NO_DEFAULT) + + @cache # noqa: B019 + def _values_list( + self, + include_nested: bool = False, + leaves_only: bool = False, + *, + collapse: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + sorting_keys: List[NestedKey] | None = None, + ) -> List: + if sorting_keys is None: + return list( + self.values( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=_NESTED_TENSORS_AS_LISTS if not collapse else is_leaf, + ) + ) + else: + keys, vals = self._items_list( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + collapse=collapse, + ) + if is_compiling(): + key_to_index = {key: i for i, key in enumerate(keys)} + return [vals[key_to_index[key]] for key in sorting_keys] + else: + source = dict(zip(keys, vals)) + return [source[key] for key in sorting_keys] + + @cache # noqa: B019 + def _items_list( + self, + include_nested: bool = False, + leaves_only: bool = False, + *, + collapse: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + sorting_keys: List[NestedKey] | None = None, + default: str | CompatibleType | None = None, + ) -> Tuple[List, List]: + items = self.items( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=_NESTED_TENSORS_AS_LISTS if not collapse else None, + ) + keys_vals = tuple(zip(*items)) + if not keys_vals: + return (), () + keys, vals = keys_vals + if sorting_keys is None: + return list(keys), list(vals) + if default is None: + # TODO: check that lists are identical + if is_compiling(): + key_to_index = {key: i for i, key in enumerate(keys)} + new_vals = [vals[key_to_index[key]] for key in sorting_keys] + if len(new_vals) < len(vals): + raise KeyError( + f"Some keys were not found: {set(sorting_keys).symmetric_difference(keys)}." + ) + else: + source = dict(zip(keys, vals)) + new_vals = [source[key] for key in sorting_keys] + if len(new_vals) < len(vals): + raise KeyError( + f"Some keys were not found: {set(sorting_keys).symmetric_difference(keys)}." + ) + return sorting_keys, new_vals + if isinstance(default, str) and default == "intersection": + new_keys = [ + key for key in sorting_keys if key in set(keys) + ] # intersection does not keep the sorting + else: + new_keys = list(set(sorting_keys).union(keys)) + source = dict(zip(keys, vals)) + vals = [source.get(key, default) for key in new_keys] + return new_keys, vals + + def _grad(self): + # We can't cache this because zero_grad can be called outside (eg from optimizer) and we want the tensors + # to clear out when that is done. + keys, vals = self._items_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + grads = [val.grad for val in vals] + items = dict(zip(keys, grads)) + + def get(name, val): + return items[name] + + return self._fast_apply( + get, + named=True, + nested_keys=True, + propagate_lock=True, + filter_empty=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + ) + + def _data(self): + keys, vals = self._items_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + data = [val.data for val in vals] + items = dict(zip(keys, data)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + propagate_lock=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + ) + + def _data_setter(self, value: Self): + keys, vals = value._items_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS) + keys_self, vals_self = value._items_list( + True, True, is_leaf=_NESTED_TENSORS_AS_LISTS + ) + self_dict = dict(zip(keys_self, vals_self)) + for key, val in zip(keys, vals): + val_self = self_dict[key] + if hasattr(val_self, "data"): + val_self.data = val + + @abc.abstractmethod + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ): + """Returns a generator of tensordict keys. + + .. warning:: + TensorDict ``keys()`` method returns a lazy view of the keys. If the ``keys`` + are queried but not iterated over and then the tensordict is modified, iterating over + the keys later will return the new configuration of the keys. + + Args: + include_nested (bool, optional): if ``True``, nested values will be returned. + Defaults to ``False``. + leaves_only (bool, optional): if ``False``, only leaves will be + returned. Defaults to ``False``. + is_leaf (callable, optional): a callable over a class type returning + a bool indicating if this class has to be considered as a leaf. + + .. note:: The purpose of `is_leaf` is not to prevent recursive calls into nested tensordicts, but + rather to mark certain types as "leaves" for the purpose of filtering when `leaves_only=True`. + Even if `is_leaf(cls)` returns `True`, the nested structure of the tensordict will still be + traversed if `include_nested=True`. + In other words, `is_leaf` does not control the recursion depth, but rather provides a way to filter + out certain types from the result when `leaves_only=True`. This means that a node in the tree can + be both a leaf and a node with children. + In practice, the default value of ``is_leaf`` does exclude tensordict and tensorclass instances + from the leaf set. + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`. + + Keyword Args: + sort (bool, optional): whether the keys shoulbe sorted. For nested keys, + the keys are sorted according to their joined name (ie, ``("a", "key")`` will + be counted as ``"a.key"`` for sorting). Be mindful that sorting may incur + significant overhead when dealing with large tensordicts. + Defaults to ``False``. + + Examples: + >>> from tensordict import TensorDict + >>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[]) + >>> data.keys() + ['0', '1'] + >>> list(data.keys(leaves_only=True)) + ['0'] + >>> list(data.keys(include_nested=True, leaves_only=True)) + ['0', '1', ('1', '2')] + """ + raise NotImplementedError + + def pop(self, key: NestedKey, default: Any = NO_DEFAULT) -> CompatibleType: + """Removes and returns a value from a tensordict. + + If the value is not present and no default value is provided, a KeyError + is thrown. + + Args: + key (str or nested key): the entry to look for. + default (Any, optional): the value to return if the key cannot be found. + + Examples: + >>> td = TensorDict({"1": 1}, []) + >>> one = td.pop("1") + >>> assert one == 1 + >>> none = td.pop("1", default=None) + >>> assert none is None + """ + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + # Use _UNSET sentinel to detect if key exists without try/except (compile-friendly) + out = self.get(key, _UNSET) + if out is _UNSET: + # Key not found + if default is NO_DEFAULT: + raise KeyError( + f"You are trying to pop key `{key}` which is not in dict " + f"without providing default value. " + f"Keys={self.keys(include_nested=isinstance(key, tuple))}." + ) + return default + self.del_(key) + return out + + @property + @cache # noqa: B019 + def sorted_keys(self) -> list[NestedKey]: + """Returns the keys sorted in alphabetical order. + + Does not support extra arguments. + + If the TensorDict is locked, the keys are cached until the tensordict + is unlocked for faster execution. + + """ + return sorted(self.keys()) + + @_as_context_manager() + def flatten(self, start_dim: int | None = None, end_dim: int | None = None): + """Flattens all the tensors of a tensordict. + + Args: + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + + Examples: + >>> td = TensorDict({ + ... "a": torch.arange(60).view(3, 4, 5), + ... "b": torch.arange(12).view(3, 4)}, batch_size=[3, 4]) + >>> td_flat = td.flatten(0, 1) + >>> td_flat.batch_size + torch.Size([12]) + >>> td_flat["a"] + tensor([[ 0, 1, 2, 3, 4], + [ 5, 6, 7, 8, 9], + [10, 11, 12, 13, 14], + [15, 16, 17, 18, 19], + [20, 21, 22, 23, 24], + [25, 26, 27, 28, 29], + [30, 31, 32, 33, 34], + [35, 36, 37, 38, 39], + [40, 41, 42, 43, 44], + [45, 46, 47, 48, 49], + [50, 51, 52, 53, 54], + [55, 56, 57, 58, 59]]) + >>> td_flat["b"] + tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + + """ + if start_dim in (None, 0) and end_dim in (None, -1, 0) and not self.ndim: + return self.unsqueeze(0) + if start_dim is None: + start_dim = 0 + if end_dim is None: + end_dim = -1 + if start_dim < 0: + start_dim = self.ndim + start_dim + if end_dim < 0: + end_dim = self.ndim + end_dim + if end_dim < 0: + raise ValueError( + f"Incompatible end_dim {end_dim} for tensordict with shape {self.shape}." + ) + if end_dim < start_dim: + raise ValueError( + "The end dimension must be greater or equal to the start dim." + ) + + def flatten(tensor): + return torch.flatten(tensor, start_dim, end_dim) + + nelt = prod(self.batch_size[start_dim : end_dim + 1]) + if start_dim > 0: + batch_size = ( + list(self.batch_size)[:start_dim] + + [nelt] + + list(self.batch_size[end_dim + 1 :]) + ) + else: + batch_size = [nelt] + list(self.batch_size[end_dim + 1 :]) + # TODO: check that this works with nested tds of different batch size + if self._has_names(): + names = [ + name + for i, name in enumerate(self.names) + if (i < start_dim or i > end_dim) + ] + names.insert(start_dim, None) + else: + names = None + out = self._fast_apply( + flatten, + batch_size=batch_size, + propagate_lock=True, + names=names, + call_on_nested=True, + ) + return out + + @_as_context_manager() + def unflatten(self, dim, unflattened_size): + """Unflattens a tensordict dim expanding it to a desired shape. + + Args: + dim (int): specifies the dimension of the input tensor to be + unflattened. + unflattened_size (shape): is the new shape of the unflattened + dimension of the tensordict. + + Examples: + >>> td = TensorDict({ + ... "a": torch.arange(60).view(3, 4, 5), + ... "b": torch.arange(12).view(3, 4)}, + ... batch_size=[3, 4]) + >>> td_flat = td.flatten(0, 1) + >>> td_unflat = td_flat.unflatten(0, [3, 4]) + >>> assert (td == td_unflat).all() + """ + dim = _maybe_correct_neg_dim(dim, self.batch_size) + + def unflatten(tensor): + return torch.unflatten( + tensor, + dim, + unflattened_size, + ) + + if dim > 0: + batch_size = ( + list(self.batch_size)[:dim] + + list(unflattened_size) + + list(self.batch_size[dim + 1 :]) + ) + else: + batch_size = list(unflattened_size) + list(self.batch_size[1:]) + # TODO: check that this works with nested tds of different batch size + out = self._fast_apply( + unflatten, batch_size=batch_size, propagate_lock=True, call_on_nested=True + ) + if self._has_names(): + names = list(self.names) + for _ in range(len(unflattened_size) - 1): + names.insert(dim, None) + out.names = names + return out + + def _transform_keys(self, key_transform: Callable[[NestedKey], NestedKey]) -> Self: + """Transform all keys using the provided function. + + Args: + key_transform (Callable): A function that takes a NestedKey and returns a new NestedKey. + For string keys, it receives a string. For tuple keys, it receives a tuple. + + Returns: + A new TensorDict with transformed keys. + + Examples: + >>> td = TensorDict({"a": torch.randn(3), "b": torch.randn(3)}, [3]) + >>> td_transformed = td._transform_keys(lambda key: f"avg_{key}") + >>> print(td_transformed.keys()) + ["avg_a", "avg_b"] + + >>> td_nested = TensorDict({("a", "b"): torch.randn(3)}, [3]) + >>> td_transformed = td_nested._transform_keys(lambda key: tuple(f"avg_{k}" for k in key)) + >>> print(td_transformed.keys()) + [("avg_a", "avg_b")] + + """ + new_td = self.empty() + for key, value in self.items(): + new_key = key_transform(key) + # If the value is a TensorDict, recursively transform its keys + if hasattr(value, "_transform_keys"): + value = value._transform_keys(key_transform) + new_td[new_key] = value + return new_td + + @abc.abstractmethod + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> Self: + """Renames a key with a new string and returns the same tensordict with the updated key name. + + Args: + old_key (str or nested key): key to be renamed. + new_key (str or nested key): new name of the entry. + safe (bool, optional): if ``True``, an error is thrown when the new + key is already present in the TensorDict. + + Returns: + self + + """ + raise NotImplementedError + + @abc.abstractmethod + def del_(self, key: NestedKey) -> Self: + """Deletes a key of the tensordict. + + Args: + key (NestedKey): key to be deleted + + Returns: + self + + """ + raise NotImplementedError + + # Distributed functionality + def gather_and_stack( + self, dst: int, group: "torch.distributed.ProcessGroup" | None = None + ) -> Self | None: + """Gathers tensordicts from various workers and stacks them onto self in the destination worker. + + Args: + dst (int): the rank of the destination worker where :func:`gather_and_stack` will be called. + group (torch.distributed.ProcessGroup, optional): if set, the specified process group + will be used for communication. Otherwise, the default process group + will be used. + Defaults to ``None``. + + Example: + >>> from torch import multiprocessing as mp + >>> from tensordict import TensorDict + >>> import torch + >>> + >>> def client(): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=1, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... # Create a single tensordict to be sent to server + ... td = TensorDict( + ... {("a", "b"): torch.randn(2), + ... "c": torch.randn(2)}, [2] + ... ) + ... td.gather_and_stack(0) + ... + >>> def server(): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=0, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... # Creates the destination tensordict on server. + ... # The first dim must be equal to world_size-1 + ... td = TensorDict( + ... {("a", "b"): torch.zeros(2), + ... "c": torch.zeros(2)}, [2] + ... ).expand(1, 2).contiguous() + ... td.gather_and_stack(0) + ... assert td["a", "b"] != 0 + ... print("yuppie") + ... + >>> if __name__ == "__main__": + ... mp.set_start_method("spawn") + ... + ... main_worker = mp.Process(target=server) + ... secondary_worker = mp.Process(target=client) + ... + ... main_worker.start() + ... secondary_worker.start() + ... + ... main_worker.join() + ... secondary_worker.join() + """ + from torch import distributed as dist + + output = ( + [None for _ in range(dist.get_world_size(group=group))] + if dst == dist.get_rank(group=group) + else None + ) + dist.gather_object(self, output, dst=dst, group=group) + if dst == dist.get_rank(group=group): + # remove self from output + output = [item for i, item in enumerate(output) if i != dst] + self.update(torch.stack(output, 0), inplace=True) + return self + return None + + def send( + self, + dst: int, + *, + group: "torch.distributed.ProcessGroup" | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> None: # noqa: D417 + """Sends the content of a tensordict to a distant worker. + + Args: + dst (int): the rank of the destination worker where the content + should be sent. + + Keyword Args: + group (torch.distributed.ProcessGroup, optional): if set, the specified process group + will be used for communication. Otherwise, the default process group + will be used. + Defaults to ``None``. + init_tag (int): the initial tag to be used to mark the tensors. + Note that this will be incremented by as much as the number of + tensors contained in the TensorDict. + pseudo_rand (bool): if True, the sequence of tags will be pseudo- + random, allowing to send multiple data from different nodes + without overlap. Notice that the generation of these pseudo-random + numbers is expensive (1e-5 sec/number), meaning that it could + slow down the runtime of your algorithm. + Defaults to ``False``. + + Example: + >>> from torch import multiprocessing as mp + >>> from tensordict import TensorDict + >>> import torch + >>> + >>> + >>> def client(): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=1, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... + ... td = TensorDict( + ... { + ... ("a", "b"): torch.randn(2), + ... "c": torch.randn(2, 3), + ... "_": torch.ones(2, 1, 5), + ... }, + ... [2], + ... ) + ... td.send(0) + ... + >>> + >>> def server(queue): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=0, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... td = TensorDict( + ... { + ... ("a", "b"): torch.zeros(2), + ... "c": torch.zeros(2, 3), + ... "_": torch.zeros(2, 1, 5), + ... }, + ... [2], + ... ) + ... td.recv(1) + ... assert (td != 0).all() + ... queue.put("yuppie") + ... + >>> + >>> if __name__=="__main__": + ... queue = mp.Queue(1) + ... main_worker = mp.Process(target=server, args=(queue,)) + ... secondary_worker = mp.Process(target=client) + ... + ... main_worker.start() + ... secondary_worker.start() + ... out = queue.get(timeout=10) + ... assert out == "yuppie" + ... main_worker.join() + ... secondary_worker.join() + + """ + self._send(dst, _tag=init_tag - 1, pseudo_rand=pseudo_rand, group=group) + + def _send( + self, + dst: int, + _tag: int = -1, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + ) -> int: + from torch import distributed as dist + + for key in self.sorted_keys: + value = self._get_str(key, NO_DEFAULT) + if isinstance(value, Tensor): + pass + elif _is_tensor_collection(type(value)): + _tag = value._send(dst, _tag=_tag, pseudo_rand=pseudo_rand, group=group) + continue + else: + raise NotImplementedError(f"Type {type(value)} is not supported.") + if not pseudo_rand: + _tag += 1 + else: + _tag = int_generator(_tag + 1) + dist.send(value, dst=dst, tag=_tag, group=group) + + return _tag + + def recv( + self, + src: int, + *, + group: "torch.distributed.ProcessGroup" | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> int: # noqa: D417 + """Receives the content of a tensordict and updates content with it. + + Check the example in the `send` method for context. + + Args: + src (int): the rank of the source worker. + + Keyword Args: + group (torch.distributed.ProcessGroup, optional): if set, the specified process group + will be used for communication. Otherwise, the default process group + will be used. + Defaults to ``None``. + init_tag (int): the ``init_tag`` used by the source worker. + pseudo_rand (bool): if True, the sequence of tags will be pseudo- + random, allowing to send multiple data from different nodes + without overlap. Notice that the generation of these pseudo-random + numbers is expensive (1e-5 sec/number), meaning that it could + slow down the runtime of your algorithm. + This value must match the one passed to :func:`send`. + Defaults to ``False``. + """ + return self._recv(src, _tag=init_tag - 1, pseudo_rand=pseudo_rand, group=group) + + def _recv( + self, + src: int, + _tag: int = -1, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + non_blocking: bool = False, + ) -> int: + from torch import distributed as dist + + for key in self.sorted_keys: + value = self._get_str(key, NO_DEFAULT) + if isinstance(value, Tensor): + pass + elif _is_tensor_collection(type(value)): + _tag = value._recv(src, _tag=_tag, pseudo_rand=pseudo_rand, group=group) + continue + else: + raise NotImplementedError(f"Type {type(value)} is not supported.") + if not pseudo_rand: + _tag += 1 + else: + _tag = int_generator(_tag + 1) + dist.recv(value, src=src, tag=_tag, group=group) + self._set_str( + key, value, inplace=True, validated=True, non_blocking=non_blocking + ) + + return _tag + + def init_remote( + self, + dst: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ) -> None: + """Initializes a remote tensordict by sending its metadata and content. + + This method sends the metadata (shape, dtype, etc.) of the current tensordict to the specified destination rank (`dst`). + + It then asynchronously sends the actual tensordict content. + + Args: + dst (int): The rank of the destination process. + group ("ProcessGroup", optional): The process group to use for communication. Defaults to None. + device (torch.device, optional): The device to use for tensor operations. Defaults to None. + + .. seealso:: + The receiving process should call `~.from_remote_init` or an equivalent method to receive and initialize a new tensordict based on the sent metadata. + + Examples: + >>> import os + >>> import torch + >>> import torch.distributed as dist + >>> from tensordict import TensorDict, MemoryMappedTensor + >>> import multiprocessing as mp + >>> + >>> def server(queue): + ... # Set environment variables for distributed communication + ... os.environ["MASTER_ADDR"] = "localhost" + ... os.environ["MASTER_PORT"] = "29505" + ... + ... # Initialize the distributed backend + ... dist.init_process_group("gloo", rank=0, world_size=2) + ... + ... # Create a sample tensordict + ... td = ( + ... TensorDict( + ... { + ... ("a", "b"): torch.ones(2), + ... "c": torch.ones(2), + ... ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)), + ... }, + ... [2], + ... ) + ... .expand(1, 2) + ... .contiguous() + ... ) + ... + ... # Send the tensordict metadata and content to the client + ... td.init_remote(dst=1) + ... + >>> def client(queue): + ... # Set environment variables for distributed communication + ... os.environ["MASTER_ADDR"] = "localhost" + ... os.environ["MASTER_PORT"] = "29505" + ... + ... # Initialize the distributed backend + ... dist.init_process_group("gloo", rank=1, world_size=2) + ... + ... # Receive the tensordict metadata and content from the server + ... received_td = TensorDict.from_remote_init(src=0) + ... + ... # Verify that the received tensordict matches the expected structure and values + ... assert set(received_td.keys()) == {"a", "c", "d"} + ... assert (received_td == 1).all() + ... + ... # Signal that the test has completed successfully + ... queue.put("yuppie") + >>> + >>> if __name__ == "__main__": + ... queue = mp.Queue(1) + ... + ... # Create and start the server and client processes + ... main_worker = mp.Process(target=server, args=(queue,)) + ... secondary_worker = mp.Process(target=client, args=(queue,)) + ... + ... main_worker.start() + ... secondary_worker.start() + ... + ... try: + ... out = queue.get(timeout=10) # Wait for the signal with a timeout + ... print(out) # Should print "yuppie" + ... finally: + ... queue.close() + ... main_worker.join(timeout=10) + ... secondary_worker.join(timeout=10) + """ + # Get a list of key - specs + data = [ + { + k: (tuple(val.shape), str(val.dtype), str(val.device)) + for k, val in self.items(True, True) + }, + self.batch_size, + self.device, + self.is_locked, + ] + torch.distributed.send_object_list( + data, + dst=dst, + group=group, + device=device, + ) + self.isend(dst, group=group) + + @classmethod + def from_remote_init( + cls: T, + src: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ) -> Self: + """Creates a new tensordict instance initialized from remotely sent metadata. + + This class method receives the metadata sent by `init_remote`, creates a new tensordict with matching shape and dtype, + and then asynchronously receives the actual tensordict content. + + Args: + src (int): The rank of the source process that sent the metadata. + group ("ProcessGroup", optional): The process group to use for communication. Defaults to None. + device (torch.device, optional): The device to use for tensor operations. Defaults to None. + + Returns: + TensorDict: A new tensordict instance initialized with the received metadata and content. + + .. seealso:: + The sending process should have called `~.init_remote` to send the metadata and content. + """ + from tensordict import TensorDict + + if not issubclass(cls, TensorDict): + raise TypeError( + f"remote initialization is currently only supported for TensorDict objects, got {cls=}." + ) + + data = [None, None, None, None] + torch.distributed.recv_object_list( + data, + src=src, + group=group, + device=device, + ) + metadata = data[0] + td = cls( + { + k: torch.empty(v[0], dtype=_STR_DTYPE_TO_DTYPE[v[1]], device=v[2]) + for k, v in metadata.items() + }, + batch_size=data[1], + device=data[2], + ) + if data[3]: + td.lock_() + td.irecv(src=src, group=group) + return td + + def isend( + self, + dst: int, + *, + group: "torch.distributed.ProcessGroup" | None = None, # noqa: F821 + init_tag: int = 0, + pseudo_rand: bool = False, + return_early: bool = False, + ) -> int | List["Work"]: # noqa: D417, F821 + """Sends the content of the tensordict asynchronously. + + Args: + dst (int): the rank of the destination worker where the content + should be sent. + + Keyword Args: + group (torch.distributed.ProcessGroup, optional): if set, the specified process group + will be used for communication. Otherwise, the default process group + will be used. + Defaults to ``None``. + init_tag (int): the initial tag to be used to mark the tensors. + Note that this will be incremented by as much as the number of + tensors contained in the TensorDict. + pseudo_rand (bool): if True, the sequence of tags will be pseudo- + random, allowing to send multiple data from different nodes + without overlap. Notice that the generation of these pseudo-random + numbers is expensive (1e-5 sec/number), meaning that it could + slow down the runtime of your algorithm. + Defaults to ``False``. + return_early (bool, optional): if True, a list of futures + will be returned instead of the tag of the last tensor sent. + Defaults to ``False``. + + Example: + >>> import torch + >>> from tensordict import TensorDict + >>> from torch import multiprocessing as mp + >>> def client(): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=1, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... + ... td = TensorDict( + ... { + ... ("a", "b"): torch.randn(2), + ... "c": torch.randn(2, 3), + ... "_": torch.ones(2, 1, 5), + ... }, + ... [2], + ... ) + ... td.isend(0) + ... + >>> + >>> def server(queue, return_premature=True): + ... torch.distributed.init_process_group( + ... "gloo", + ... rank=0, + ... world_size=2, + ... init_method=f"tcp://localhost:10003", + ... ) + ... td = TensorDict( + ... { + ... ("a", "b"): torch.zeros(2), + ... "c": torch.zeros(2, 3), + ... "_": torch.zeros(2, 1, 5), + ... }, + ... [2], + ... ) + ... out = td.irecv(1, return_premature=return_premature) + ... if return_premature: + ... for fut in out: + ... fut.wait() + ... assert (td != 0).all() + ... queue.put("yuppie") + ... + >>> + >>> if __name__ == "__main__": + ... queue = mp.Queue(1) + ... main_worker = mp.Process( + ... target=server, + ... args=(queue, ) + ... ) + ... secondary_worker = mp.Process(target=client) + ... + ... main_worker.start() + ... secondary_worker.start() + ... out = queue.get(timeout=10) + ... assert out == "yuppie" + ... main_worker.join() + ... secondary_worker.join() + + """ + return self._isend( + dst, + _tag=init_tag - 1, + pseudo_rand=pseudo_rand, + group=group, + return_early=return_early, + ) + + def _isend( + self, + dst: int, + _tag: int = -1, + _futures: list[torch.Future] | None = None, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + return_early: bool = False, + ) -> int: + from torch import distributed as dist + + root = False + if _futures is None: + root = True + _futures = [] + for key in self.sorted_keys: + value = self._get_str(key, NO_DEFAULT) + if _is_tensor_collection(type(value)): + _tag = value._isend( + dst, + _tag=_tag, + pseudo_rand=pseudo_rand, + _futures=_futures, + group=group, + return_early=return_early, + ) + continue + elif isinstance(value, Tensor): + pass + else: + raise NotImplementedError(f"Type {type(value)} is not supported.") + if not pseudo_rand: + _tag += 1 + else: + _tag = int_generator(_tag + 1) + _future = dist.isend(value, dst=dst, tag=_tag, group=group) + _futures.append(_future) + if root and not return_early: + for _future in _futures: + _future.wait() + elif root and return_early: + return _futures + return _tag + + def irecv( + self, + src: int, + *, + group: "torch.distributed.ProcessGroup" | None = None, + return_premature: bool = False, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: + """Receives the content of a tensordict and updates content with it asynchronously. + + Check the example in the :meth:`~.isend` method for context. + + Args: + src (int): the rank of the source worker. + + Keyword Args: + group (torch.distributed.ProcessGroup, optional): if set, the specified process group + will be used for communication. Otherwise, the default process group + will be used. + Defaults to ``None``. + return_premature (bool): if ``True``, returns a list of futures to wait + upon until the tensordict is updated. Defaults to ``False``, + i.e. waits until update is completed withing the call. + init_tag (int): the ``init_tag`` used by the source worker. + pseudo_rand (bool): if True, the sequence of tags will be pseudo- + random, allowing to send multiple data from different nodes + without overlap. Notice that the generation of these pseudo-random + numbers is expensive (1e-5 sec/number), meaning that it could + slow down the runtime of your algorithm. + This value must match the one passed to :func:`isend`. + Defaults to ``False``. + + Returns: + if ``return_premature=True``, a list of futures to wait + upon until the tensordict is updated. + """ + return self._irecv( + src, + return_premature=return_premature, + _tag=init_tag - 1, + pseudo_rand=pseudo_rand, + group=group, + ) + + def _irecv( + self, + src: int, + return_premature: bool = False, + _tag: int = -1, + _future_list: list[torch.Future] = None, + pseudo_rand: bool = False, + group: "torch.distributed.ProcessGroup" | None = None, + ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: + from torch import distributed as dist + + root = False + if _future_list is None: + _future_list = [] + root = True + + for key in self.sorted_keys: + value = self._get_str(key, NO_DEFAULT) + if _is_tensor_collection(type(value)): + _tag, _future_list = value._irecv( + src, + _tag=_tag, + _future_list=_future_list, + pseudo_rand=pseudo_rand, + group=group, + ) + continue + elif isinstance(value, Tensor): + pass + else: + raise NotImplementedError(f"Type {type(value)} is not supported.") + if not pseudo_rand: + _tag += 1 + else: + _tag = int_generator(_tag + 1) + _future_list.append(dist.irecv(value, src=src, tag=_tag, group=group)) + if not root: + return _tag, _future_list + elif return_premature: + return _future_list + else: + for future in _future_list: + future.wait() + return + + def reduce( + self, + dst, + op=None, + async_op=False, + return_premature=False, + group=None, + ) -> None: + """Reduces the tensordict across all machines. + + Only the process with ``rank`` dst is going to receive the final result. + + """ + from torch import distributed as dist + + if op is None: + op = dist.ReduceOp.SUM + return self._reduce(dst, op, async_op, return_premature, group=group) + + def _reduce( + self, + dst, + op=None, + async_op=False, + return_premature=False, + _future_list=None, + group=None, + ): + from torch import distributed as dist + + if op is None: + op = dist.ReduceOp.SUM + root = False + if _future_list is None: + _future_list = [] + root = True + for key in self.sorted_keys: + value = self._get_str(key, NO_DEFAULT) + if _is_tensor_collection(type(value)): + _future_list = value._reduce( + dst=dst, + op=op, + async_op=async_op, + _future_list=_future_list, + ) + continue + elif isinstance(value, Tensor): + pass + else: + raise NotImplementedError(f"Type {type(value)} is not supported.") + _future_list.append( + dist.reduce(value, dst=dst, op=op, async_op=async_op, group=group) + ) + if not root: + return _future_list + elif async_op and return_premature: + return _future_list + elif async_op: + for future in _future_list: + future.wait() + return + + # Apply and map functionality + def apply_(self, fn: Callable, *others, **kwargs) -> Self: + """Applies a callable to all values stored in the tensordict and re-writes them in-place. + + Args: + fn (Callable): function to be applied to the tensors in the + tensordict. + *others (sequence of TensorDictBase, optional): the other + tensordicts to be used. + + Keyword Args: See :meth:`~.apply`. + + Returns: + self or a copy of self with the function applied + + """ + return self.apply(fn, *others, inplace=True, **kwargs) + + def apply( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + default: Any = NO_DEFAULT, + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorDictBase | None = None, + **constructor_kwargs, + ) -> Self | None: + """Applies a callable to all values stored in the tensordict and sets them in a new tensordict. + + The callable signature must be ``Callable[Tuple[Tensor, ...], Tensor | TensorDictBase | None]``. + + Args: + fn (Callable): function to be applied to the tensors in the + tensordict. + *others (TensorDictBase instances, optional): if provided, these + tensordict instances should have a structure matching the one + of self. The ``fn`` argument should receive as many + unnamed inputs as the number of tensordicts, including self. + If other tensordicts have missing entries, a default value + can be passed through the ``default`` keyword argument. + + Keyword Args: + batch_size (sequence of int, optional): if provided, + the resulting TensorDict will have the desired batch_size. + The :obj:`batch_size` argument should match the batch_size after + the transformation. This is a keyword only argument. + device (torch.device, optional): the resulting device, if any. + names (list of str, optional): the new dimension names, in case the + batch_size is modified. + inplace (bool, optional): if True, changes are made in-place. + Default is False. This is a keyword only argument. + default (Any, optional): default value for missing entries in the + other tensordicts. If not provided, missing entries will + raise a `KeyError`. + filter_empty (bool, optional): if ``True``, empty tensordicts will be + filtered out. This also comes with a lower computational cost as + empty data structures won't be created and destroyed. Non-tensor data + is considered as a leaf and thereby will be kept in the tensordict even + if left untouched by the function. + Defaults to ``False`` for backward compatibility. + propagate_lock (bool, optional): if ``True``, a locked tensordict will produce + another locked tensordict. Defaults to ``False``. + call_on_nested (bool, optional): if ``True``, the function will be called on first-level tensors + and containers (TensorDict or tensorclass). In this scenario, ``func`` is responsible of + propagating its calls to nested levels. This allows a fine-grained behaviour + when propagating the calls to nested tensordicts. + If ``False``, the function will only be called on leaves, and ``apply`` will take care of dispatching + the function to all leaves. + + >>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]}) + >>> def mean_tensor_only(val): + ... if is_tensor_collection(val): + ... raise RuntimeError("Unexpected!") + ... return val.mean() + >>> td_mean = td.apply(mean_tensor_only) + >>> def mean_any(val): + ... if is_tensor_collection(val): + ... # Recurse + ... return val.apply(mean_any, call_on_nested=True) + ... return val.mean() + >>> td_mean = td.apply(mean_any, call_on_nested=True) + out (TensorDictBase, optional): a tensordict where to write the results. This can be used to avoid + creating a new tensordict: + + >>> td = TensorDict({"a": 0}) + >>> td.apply(lambda x: x+1, out=td) + >>> assert (td==1).all() + + .. warning:: + If the operation executed on the tensordict requires multiple keys to be accessed for + a single computation, providing an ``out`` argument equal to ``self`` can cause the operation + to provide silently wrong results. + For instance: + + >>> td = TensorDict({"a": 1, "b": 1}) + >>> td.apply(lambda x: x+td["a"])["b"] # Right! + tensor(2) + >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong! + tensor(3) + + **constructor_kwargs: additional keyword arguments to be passed to the + TensorDict constructor. + + Returns: + a new tensordict with transformed_in tensors. + + Example: + >>> td = TensorDict({ + ... "a": -torch.ones(3), + ... "b": {"c": torch.ones(3)}}, + ... batch_size=[3]) + >>> td_1 = td.apply(lambda x: x+1) + >>> assert (td_1["a"] == 0).all() + >>> assert (td_1["b", "c"] == 2).all() + >>> td_2 = td.apply(lambda x, y: x+y, td) + >>> assert (td_2["a"] == -2).all() + >>> assert (td_2["b", "c"] == 2).all() + + .. note:: + If ``None`` is returned by the function, the entry is ignored. This + can be used to filter the data in the tensordict: + + >>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, []) + >>> def filter(tensor): + ... if tensor == 1: + ... return tensor + >>> td.apply(filter) + TensorDict( + fields={ + 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + .. note:: + The apply method will return an :class:`~tensordict.TensorDict` instance, + regardless of the input type. To keep the same type, one can execute + + >>> out = td.clone(False).update(td.apply(...)) + + + """ + result = self._apply_nest( + fn, + *others, + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=False, + default=default, + filter_empty=filter_empty, + call_on_nested=call_on_nested, + out=out, + **constructor_kwargs, + ) + if propagate_lock and not inplace and self.is_locked and result is not None: + result.lock_() + return result + + def named_apply( + self, + fn: Callable, + *others: T, + nested_keys: bool = False, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + default: Any = NO_DEFAULT, + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorDictBase | None = None, + **constructor_kwargs, + ) -> Self | None: + """Applies a key-conditioned callable to all values stored in the tensordict and sets them in a new atensordict. + + The callable signature must be ``Callable[Tuple[str, Tensor, ...], Tensor | TensorDictBase | None]``. + + Args: + fn (Callable): function to be applied to the (name, tensor) pairs in the + tensordict. For each leaf, only its leaf name will be used (not + the full `NestedKey`). + *others (TensorDictBase instances, optional): if provided, these + tensordict instances should have a structure matching the one + of self. The ``fn`` argument should receive as many + unnamed inputs as the number of tensordicts, including self. + If other tensordicts have missing entries, a default value + can be passed through the ``default`` keyword argument. + nested_keys (bool, optional): if ``True``, the complete path + to the leaf will be used. Defaults to ``False``, i.e. only the last + string is passed to the function. + batch_size (sequence of int, optional): if provided, + the resulting TensorDict will have the desired batch_size. + The :obj:`batch_size` argument should match the batch_size after + the transformation. This is a keyword only argument. + device (torch.device, optional): the resulting device, if any. + names (list of str, optional): the new dimension names, in case the + batch_size is modified. + inplace (bool, optional): if True, changes are made in-place. + Default is False. This is a keyword only argument. + default (Any, optional): default value for missing entries in the + other tensordicts. If not provided, missing entries will + raise a `KeyError`. + filter_empty (bool, optional): if ``True``, empty tensordicts will be + filtered out. This also comes with a lower computational cost as + empty data structures won't be created and destroyed. Defaults to + ``False`` for backward compatibility. + propagate_lock (bool, optional): if ``True``, a locked tensordict will produce + another locked tensordict. Defaults to ``False``. + call_on_nested (bool, optional): if ``True``, the function will be called on first-level tensors + and containers (TensorDict or tensorclass). In this scenario, ``func`` is responsible of + propagating its calls to nested levels. This allows a fine-grained behaviour + when propagating the calls to nested tensordicts. + If ``False``, the function will only be called on leaves, and ``apply`` will take care of dispatching + the function to all leaves. + + >>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]}) + >>> def mean_tensor_only(val): + ... if is_tensor_collection(val): + ... raise RuntimeError("Unexpected!") + ... return val.mean() + >>> td_mean = td.apply(mean_tensor_only) + >>> def mean_any(val): + ... if is_tensor_collection(val): + ... # Recurse + ... return val.apply(mean_any, call_on_nested=True) + ... return val.mean() + >>> td_mean = td.apply(mean_any, call_on_nested=True) + + out (TensorDictBase, optional): a tensordict where to write the results. This can be used to avoid + creating a new tensordict: + + >>> td = TensorDict({"a": 0}) + >>> td.apply(lambda x: x+1, out=td) + >>> assert (td==1).all() + + .. warning:: + If the operation executed on the tensordict requires multiple keys to be accessed for + a single computation, providing an ``out`` argument equal to ``self`` can cause the operation + to provide silently wrong results. + For instance: + + >>> td = TensorDict({"a": 1, "b": 1}) + >>> td.apply(lambda x: x+td["a"])["b"] # Right! + tensor(2) + >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong! + tensor(3) + + **constructor_kwargs: additional keyword arguments to be passed to the + TensorDict constructor. + + Returns: + a new tensordict with transformed_in tensors. + + Example: + >>> td = TensorDict({ + ... "a": -torch.ones(3), + ... "nested": {"a": torch.ones(3), "b": torch.zeros(3)}}, + ... batch_size=[3]) + >>> def name_filter(name, tensor): + ... if name == "a": + ... return tensor + >>> td.named_apply(name_filter) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + nested: TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> def name_filter(name, *tensors): + ... if name == "a": + ... r = 0 + ... for tensor in tensors: + ... r = r + tensor + ... return tensor + >>> out = td.named_apply(name_filter, td) + >>> print(out) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + nested: TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> print(out["a"]) + tensor([-1., -1., -1.]) + + .. note:: + If ``None`` is returned by the function, the entry is ignored. This + can be used to filter the data in the tensordict: + + >>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, []) + >>> def name_filter(name, tensor): + ... if name == "1": + ... return tensor + >>> td.named_apply(name_filter) + TensorDict( + fields={ + 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + result = self._apply_nest( + fn, + *others, + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=False, + default=default, + named=True, + nested_keys=nested_keys, + filter_empty=filter_empty, + call_on_nested=call_on_nested, + **constructor_kwargs, + ) + if propagate_lock and not inplace and self.is_locked and result is not None: + result.lock_() + return result + + @abc.abstractmethod + def _multithread_apply_flat( + self, + fn: Callable, + *others: T, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + is_leaf: Callable[[Type], bool] | None = None, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + ) -> None: + raise NotImplementedError + + @abc.abstractmethod + def _multithread_rebuild( + self, + *, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + out: TensorDictBase | None = None, + filter_empty: bool = False, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + subs_results: Dict[Future, Any] | None = None, + multithread_set: bool = False, # Experimental + **constructor_kwargs, + ) -> None: + raise NotImplementedError + + def _multithread_apply_nest( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + filter_empty: bool | None = None, + is_leaf: Callable[[Type], bool] | None = None, + out: TensorDictBase | None = None, + num_threads: int, + call_when_done: Callable | None = None, + **constructor_kwargs, + ) -> Self | None: + """A deadlock-safe multithread wrapper around TD.apply. + + First launches fn for all the leaves, then rebuilds the tensordicts out of them. + + An optional ``call_when_done`` function can be passed to execute a method on the main thread + after a future is completed. + + """ + if call_on_nested: + warnings.warn( + "Multithreaded apply with call_on_nested=True should not be used for deep TensorDicts. " + "In the best cases, it will be inefficient, in the worst an arbitrary large number of " + "threads will be spawn." + ) + # We create 2 structures that will have the same elements within: + # futures is a flat list of all the futures we need to wait for, + # local_futures is a nested representation of this flat structure. + # In local_futures, the order of the items can be used to link the items to their key. + futures = [] + local_futures = [] + executor = ThreadPoolExecutor(max_workers=num_threads) + self._multithread_apply_flat( + fn, + *others, + call_on_nested=call_on_nested, + default=default, + named=named, + nested_keys=nested_keys, + prefix=prefix, + is_leaf=is_leaf, + executor=executor, + futures=futures, + local_futures=local_futures, + ) + if call_when_done is not None: + subs_results = {} + + def cb(fut): + fut._result = call_when_done(fut.result()) + return fut + + for fut in futures: + fut.add_done_callback(cb) + # wait(futures) + # for fut in as_completed(futures): + # subs_results[fut] = call_when_done(fut.result()) + # fut._result = None + # # futures.remove(fut) + # # del fut + else: + subs_results = None + return self._multithread_rebuild( + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=checked, + out=out, + filter_empty=filter_empty, + executor=executor, + futures=futures, + local_futures=local_futures, + subs_results=subs_results, + **constructor_kwargs, + ) + + @abc.abstractmethod + def _apply_nest( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + filter_empty: bool | None = None, + is_leaf: Callable[[Type], bool] | None = None, + out: TensorDictBase | None = None, + **constructor_kwargs, + ) -> Self | None: + raise NotImplementedError + + def _fast_apply( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + # filter_empty must be False because we use _fast_apply for all sorts of ops like expand etc + # and non-tensor data will disappear if we use True by default. + filter_empty: bool | None = False, + is_leaf: Callable[[Type], bool] | None = None, + propagate_lock: bool = False, + out: TensorDictBase | None = None, + num_threads: int = 0, + checked: bool = True, + **constructor_kwargs, + ) -> Self | None: + """A faster apply method. + + This method does not run any check after performing the func. This + means that one to make sure that the metadata of the resulting tensors + (device, shape etc.) match the :meth:`~.apply` ones. + + """ + if num_threads: + + def func(*args, **kwargs): + return self._multithread_apply_nest( + *args, **kwargs, num_threads=num_threads + ) + + else: + func = self._apply_nest + result = func( + fn, + *others, + batch_size=batch_size, + device=device, + names=names, + inplace=inplace, + checked=checked, + call_on_nested=call_on_nested, + named=named, + default=default, + nested_keys=nested_keys, + filter_empty=filter_empty, + is_leaf=is_leaf, + out=out, + **constructor_kwargs, + ) + if propagate_lock and not inplace and self.is_locked and result is not None: + result.lock_() + return result + + def map( + self, + fn: Callable[[TensorCollection], TensorCollection | None], + dim: int = 0, + num_workers: int | None = None, + *, + out: TensorCollection | None = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = False, + pbar: bool = False, + mp_start_method: str | None = None, + ) -> Self: + """Maps a function to splits of the tensordict across one dimension. + + This method will apply a function to a tensordict instance by chunking + it in tensordicts of equal size and dispatching the operations over the + desired number of workers. + + The function signature should be ``Callabe[[TensorDict], TensorDict | Tensor]``. + The output must support the :func:`torch.cat` operation. The function + must be serializable. + + .. note:: + This method is particularly useful when working with large + datasets stored on disk (e.g. memory-mapped tensordicts) where + chunks will be zero-copied slices of the original data which can + be passed to the processes with virtually zero-cost. This allows + to tread very large datasets (eg. over a Tb big) to be processed + at little cost. + + Args: + fn (callable): function to apply to the tensordict. + Signatures similar to ``Callabe[[TensorDict], TensorDict | Tensor]`` + are supported. + dim (int, optional): the dim along which the tensordict will be chunked. + num_workers (int, optional): the number of workers. Exclusive with ``pool``. + If none is provided, the number of workers will be set to the + number of cpus available. + + Keyword Args: + out (TensorDictBase, optional): an optional container for the output. + Its batch-size along the ``dim`` provided must match ``self.ndim``. + If it is shared or memmap (:meth:`~.is_shared` or :meth:`~.is_memmap` + returns ``True``) it will be populated within the remote processes, + avoiding data inward transfers. Otherwise, the data from the ``self`` + slice will be sent to the process, collected on the current process + and written inplace into ``out``. + chunksize (int, optional): The size of each chunk of data. + A ``chunksize`` of 0 will unbind the tensordict along the + desired dimension and restack it after the function is applied, + whereas ``chunksize>0`` will split the tensordict and call + :func:`torch.cat` on the resulting list of tensordicts. + If none is provided, the number of chunks will equate the number + of workers. For very large tensordicts, such large chunks + may not fit in memory for the operation to be done and + more chunks may be needed to make the operation practically + doable. This argument is exclusive with ``num_chunks``. + num_chunks (int, optional): the number of chunks to split the tensordict + into. If none is provided, the number of chunks will equate the number + of workers. For very large tensordicts, such large chunks + may not fit in memory for the operation to be done and + more chunks may be needed to make the operation practically + doable. This argument is exclusive with ``chunksize``. + pool (mp.Pool, optional): a multiprocess Pool instance to use + to execute the job. If none is provided, a pool will be created + within the ``map`` method. + generator (torch.Generator, optional): a generator to use for seeding. + A base seed will be generated from it, and each worker + of the pool will be seeded with the provided seed incremented + by a unique integer from ``0`` to ``num_workers``. If no generator + is provided, a random integer will be used as seed. + To work with unseeded workers, a pool should be created separately + and passed to :meth:`map` directly. + + .. note:: + Caution should be taken when providing a low-valued seed as + this can cause autocorrelation between experiments, example: + if 8 workers are asked and the seed is 4, the workers seed will + range from 4 to 11. If the seed is 5, the workers seed will range + from 5 to 12. These two experiments will have an overlap of 7 + seeds, which can have unexpected effects on the results. + + .. note:: + The goal of seeding the workers is to have independent seed on + each worker, and NOT to have reproducible results across calls + of the `map` method. In other words, two experiments may and + probably will return different results as it is impossible to + know which worker will pick which job. However, we can make sure + that each worker has a different seed and that the pseudo-random + operations on each will be uncorrelated. + + max_tasks_per_child (int, optional): the maximum number of jobs picked + by every child process. Defaults to ``None``, i.e., no restriction + on the number of jobs. + worker_threads (int, optional): the number of threads for the workers. + Defaults to ``1``. + index_with_generator (bool, optional): if ``True``, the splitting / chunking + of the tensordict will be done during the query, sparing init time. + Note that :meth:`~.chunk` and :meth:`~.split` are much more + efficient than indexing (which is used within the generator) + so a gain of processing time at init time may have a negative + impact on the total runtime. Defaults to ``False``. + pbar (bool, optional): if ``True``, a progress bar will be displayed. + Requires tqdm to be available. Defaults to ``False``. + mp_start_method (str, optional): the start method for multiprocessing. + If not provided, the default start method will be used. + Accepted strings are ``"fork"`` and ``"spawn"``. Keep in mind that + ``"cuda"`` tensors cannot be shared between processes with the + ``"fork"`` start method. This is without effect if the ``pool`` + is passed to the ``map`` method. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> + >>> def process_data(data): + ... data.set("y", data.get("x") + 1) + ... return data + >>> if __name__ == "__main__": + ... data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_() + ... data = data.map(process_data, dim=1) + ... print(data["y"][:, :10]) + ... + tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]) + """ + from torch import multiprocessing as mp + + if pool is None: + if num_workers is None: + num_workers = mp.cpu_count() # Get the number of CPU cores + if generator is None: + generator = torch.Generator() + seed = ( + torch.empty((), dtype=torch.int64).random_(generator=generator).item() + ) + if mp_start_method is not None: + ctx = mp.get_context(mp_start_method) + else: + ctx = mp.get_context() + + queue = ctx.Queue(maxsize=num_workers) + for i in range(num_workers): + queue.put(i) + with ctx.Pool( + processes=num_workers, + initializer=_proc_init, + initargs=(seed, queue, worker_threads), + maxtasksperchild=max_tasks_per_child, + ) as pool: + return self._map( + fn=fn, + dim=dim, + chunksize=chunksize, + num_chunks=num_chunks, + pool=pool, + pbar=pbar, + out=out, + index_with_generator=index_with_generator, + iterable=False, + shuffle=False, + ) + else: + return self._map( + fn=fn, + dim=dim, + chunksize=chunksize, + num_chunks=num_chunks, + pool=pool, + pbar=pbar, + out=out, + index_with_generator=index_with_generator, + iterable=False, + shuffle=False, + ) + + def map_iter( + self, + fn: Callable[[TensorCollection], TensorCollection | None], + dim: int = 0, + num_workers: int | None = None, + *, + shuffle: bool = False, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = True, + pbar: bool = False, + mp_start_method: str | None = None, + ) -> Iterator[T]: + """Maps a function to splits of the tensordict across one dimension iteratively. + + This is the iterable version of :meth:`~TensorDictBase.map`. + + This method will apply a function to a tensordict instance by chunking + it in tensordicts of equal size and dispatching the operations over the + desired number of workers. It will yield the results one at a time. + + The function signature should be ``Callabe[[TensorDict], TensorDict | Tensor]``. + The function must be serializable. + + .. note:: + This method is particularly useful when working with large + datasets stored on disk (e.g. memory-mapped tensordicts) where + chunks will be zero-copied slices of the original data which can + be passed to the processes with virtually zero-cost. This allows + to tread very large datasets (eg. over a Tb big) to be processed + at little cost. + + .. note:: + This function be used to represent a dataset and load from it, + in a dataloader-like fashion. + + Args: + fn (callable): function to apply to the tensordict. + Signatures similar to ``Callabe[[TensorDict], TensorDict | Tensor]`` + are supported. + dim (int, optional): the dim along which the tensordict will be chunked. + num_workers (int, optional): the number of workers. Exclusive with ``pool``. + If none is provided, the number of workers will be set to the + number of cpus available. + + Keyword Args: + shuffle (bool, optional): whether the indices should be globally shuffled. + If ``True``, each batch will contain non-contiguous samples. + If ``index_with_generator=False`` and `shuffle=True``, an error will be raised. + Defaults to ``False``. + chunksize (int, optional): The size of each chunk of data. + A ``chunksize`` of 0 will unbind the tensordict along the + desired dimension and restack it after the function is applied, + whereas ``chunksize>0`` will split the tensordict and call + :func:`torch.cat` on the resulting list of tensordicts. + If none is provided, the number of chunks will equate the number + of workers. For very large tensordicts, such large chunks + may not fit in memory for the operation to be done and + more chunks may be needed to make the operation practically + doable. This argument is exclusive with ``num_chunks``. + num_chunks (int, optional): the number of chunks to split the tensordict + into. If none is provided, the number of chunks will equate the number + of workers. For very large tensordicts, such large chunks + may not fit in memory for the operation to be done and + more chunks may be needed to make the operation practically + doable. This argument is exclusive with ``chunksize``. + pool (mp.Pool, optional): a multiprocess Pool instance to use + to execute the job. If none is provided, a pool will be created + within the ``map`` method. + generator (torch.Generator, optional): a generator to use for seeding. + A base seed will be generated from it, and each worker + of the pool will be seeded with the provided seed incremented + by a unique integer from ``0`` to ``num_workers``. If no generator + is provided, a random integer will be used as seed. + To work with unseeded workers, a pool should be created separately + and passed to :meth:`map` directly. + + .. note:: + Caution should be taken when providing a low-valued seed as + this can cause autocorrelation between experiments, example: + if 8 workers are asked and the seed is 4, the workers seed will + range from 4 to 11. If the seed is 5, the workers seed will range + from 5 to 12. These two experiments will have an overlap of 7 + seeds, which can have unexpected effects on the results. + + .. note:: + The goal of seeding the workers is to have independent seed on + each worker, and NOT to have reproducible results across calls + of the `map` method. In other words, two experiments may and + probably will return different results as it is impossible to + know which worker will pick which job. However, we can make sure + that each worker has a different seed and that the pseudo-random + operations on each will be uncorrelated. + + max_tasks_per_child (int, optional): the maximum number of jobs picked + by every child process. Defaults to ``None``, i.e., no restriction + on the number of jobs. + worker_threads (int, optional): the number of threads for the workers. + Defaults to ``1``. + index_with_generator (bool, optional): if ``True``, the splitting / chunking + of the tensordict will be done during the query, sparing init time. + Note that :meth:`~.chunk` and :meth:`~.split` are much more + efficient than indexing (which is used within the generator) + so a gain of processing time at init time may have a negative + impact on the total runtime. Defaults to ``True``. + + .. note:: + The default value of ``index_with_generator`` differs for ``map_iter`` + and ``map`` and the former assumes that it is prohibitively expensive to + store a split version of the TensorDict in memory. + + pbar (bool, optional): if ``True``, a progress bar will be displayed. + Requires tqdm to be available. Defaults to ``False``. + mp_start_method (str, optional): the start method for multiprocessing. + If not provided, the default start method will be used. + Accepted strings are ``"fork"`` and ``"spawn"``. Keep in mind that + ``"cuda"`` tensors cannot be shared between processes with the + ``"fork"`` start method. This is without effect if the ``pool`` + is passed to the ``map`` method. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> + >>> def process_data(data): + ... data.unlock_() + ... data.set("y", data.get("x") + 1) + ... return data + >>> if __name__ == "__main__": + ... data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_() + ... for sample in data.map_iter(process_data, dim=1, chunksize=5): + ... print(sample["y"]) + ... break + ... + tensor([[1., 1., 1., 1., 1.]]) + + """ + from torch import multiprocessing as mp + + if pool is None: + if num_workers is None: + num_workers = mp.cpu_count() # Get the number of CPU cores + if generator is None: + generator = torch.Generator() + seed = ( + torch.empty((), dtype=torch.int64).random_(generator=generator).item() + ) + if mp_start_method is not None: + ctx = mp.get_context(mp_start_method) + else: + ctx = mp.get_context() + + queue = ctx.Queue(maxsize=num_workers) + for i in range(num_workers): + queue.put(i) + pool = ctx.Pool( + processes=num_workers, + initializer=_proc_init, + initargs=(seed, queue, worker_threads), + maxtasksperchild=max_tasks_per_child, + ) + try: + yield from self._map( + fn=fn, + dim=dim, + chunksize=chunksize, + num_chunks=num_chunks, + pool=pool, + pbar=pbar, + out=None, + index_with_generator=index_with_generator, + iterable=True, + shuffle=shuffle, + ) + finally: + try: + pool.close() + pool.join() + except Exception: + pool.terminate() + else: + yield from self._map( + fn=fn, + dim=dim, + chunksize=chunksize, + num_chunks=num_chunks, + pool=pool, + pbar=pbar, + out=None, + index_with_generator=index_with_generator, + iterable=True, + shuffle=shuffle, + ) + + def _map( + self, + fn: Callable[[TensorDictBase], TensorDictBase | None], + dim: int = 0, + *, + shuffle: bool = False, + out: TensorDictBase | None = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + index_with_generator: bool = False, + pbar: bool = False, + iterable: bool, + ): + num_workers = pool._processes + dim = _maybe_correct_neg_dim(dim, self.batch_size) + + self_split = _split_tensordict( + self, + chunksize, + num_chunks, + num_workers, + dim, + shuffle=shuffle, + use_generator=index_with_generator, + ) + if not index_with_generator: + length = len(self_split) + else: + length = None + call_chunksize = 1 + + if out is not None and (out.is_shared() or out.is_memmap()): + + def wrap_fn_with_out(fn, out): + @wraps(fn) + def newfn(item_and_out): + item, out = item_and_out + result = fn(item) + out.update_(result) + return + + out_split = _split_tensordict( + out, + chunksize, + num_chunks, + num_workers, + dim, + shuffle=shuffle, + use_generator=index_with_generator, + ) + return _CloudpickleWrapper(newfn), _zip_strict(self_split, out_split) + + fn, self_split = wrap_fn_with_out(fn, out) + out = None + + imap_fn = pool.imap if not shuffle else pool.imap_unordered + imap = imap_fn(fn, self_split, call_chunksize) + + if pbar and importlib.util.find_spec("tqdm", None) is not None: + import tqdm + + imap = tqdm.tqdm(imap, total=length) + + if iterable: + return imap + else: + imaplist = [] + start = 0 + base_index = (slice(None),) * dim + for item in imap: + if item is not None: + if out is not None: + if chunksize == 0: + out[base_index + (start,)].update_(item) + start += 1 + else: + end = start + item.shape[dim] + chunk = base_index + (slice(start, end),) + out[chunk].update_(item) + start = end + else: + imaplist.append(item) + del imap + + # support inplace modif + if imaplist: + if chunksize == 0: + from tensordict._lazy import LazyStackedTensorDict + + # We want to be able to return whichever data structure + with set_capture_non_tensor_stack(False): + out = LazyStackedTensorDict.maybe_dense_stack(imaplist, dim) + else: + out = torch.cat(imaplist, dim) + return out + + # Stream + def record_stream(self, stream: torch.cuda.Stream) -> Self: + """Marks the tensordict as having been used by this stream. + + When the tensordict is deallocated, ensure the tensor memory is not reused for other tensors until all work + queued on stream at the time of deallocation is complete. + + See :meth:`~torch.Tensor.record_stream` for more information.` + + """ + if self._stream is not None and self._stream != stream: + raise RuntimeError( + "A stream is already associated with this TensorDict instance." + ) + self._stream = stream + + def record(tensor): + tensor.record_stream(stream) + + self._fast_apply(record, filter_empty=True) + return self + + def __copy__(self): + """Copies the tensordict without cloning its tensors.""" + return self.copy() + + # point-wise arithmetic ops + def __add__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.add(other) + + def __radd__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.add(other) + + def __iadd__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.add_(other) + + def __truediv__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.div(other) + + def __itruediv__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.div_(other) + + def __rtruediv__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.reciprocal() * other + + def __mul__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.mul(other) + + def __mod__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.mod(other) + + def __rmul__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.mul(other) + + def __imul__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.mul_(other) + + def __sub__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.sub(other) + + def __isub__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.sub_(other) + + def __rsub__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.rsub(other) + + def __pow__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.pow(other) + + def __rpow__(self, other: TensorCollection | torch.Tensor) -> Self: + raise NotImplementedError( + "rpow isn't implemented for tensordict yet. Make sure both elements are wrapped " + "in a tensordict for this to work." + ) + + def __ipow__(self, other: TensorCollection | torch.Tensor) -> Self: + return self.pow_(other) + + def abs(self) -> Self: + """Computes the absolute value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_abs(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def abs_(self) -> Self: + """Computes the absolute value of each element of the TensorDict in-place.""" + torch._foreach_abs_(self._values_list(True, True)) + return self + + def acos(self) -> Self: + """Computes the :meth:`~torch.acos` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_acos(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def acos_(self) -> Self: + """Computes the :meth:`~torch.acos` value of each element of the TensorDict in-place.""" + torch._foreach_acos_(self._values_list(True, True)) + return self + + def exp(self) -> Self: + """Computes the :meth:`~torch.exp` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_exp(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def exp_(self) -> Self: + """Computes the :meth:`~torch.exp` value of each element of the TensorDict in-place.""" + torch._foreach_exp_(self._values_list(True, True)) + return self + + def neg(self) -> Self: + """Computes the :meth:`~torch.neg` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_neg(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def neg_(self) -> Self: + """Computes the :meth:`~torch.neg` value of each element of the TensorDict in-place.""" + torch._foreach_neg_(self._values_list(True, True)) + return self + + def reciprocal(self) -> Self: + """Computes the :meth:`~torch.reciprocal` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_reciprocal(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def reciprocal_(self) -> Self: + """Computes the :meth:`~torch.reciprocal` value of each element of the TensorDict in-place.""" + torch._foreach_reciprocal_(self._values_list(True, True)) + return self + + def sigmoid(self) -> Self: + """Computes the :meth:`~torch.sigmoid` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_sigmoid(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def sigmoid_(self) -> Self: + """Computes the :meth:`~torch.sigmoid` value of each element of the TensorDict in-place.""" + torch._foreach_sigmoid_(self._values_list(True, True)) + return self + + def sign(self) -> Self: + """Computes the :meth:`~torch.sign` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_sign(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def sign_(self) -> Self: + """Computes the :meth:`~torch.sign` value of each element of the TensorDict in-place.""" + torch._foreach_sign_(self._values_list(True, True)) + return self + + def sin(self) -> Self: + """Computes the :meth:`~torch.sin` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_sin(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def sin_(self) -> Self: + """Computes the :meth:`~torch.sin` value of each element of the TensorDict in-place.""" + torch._foreach_sin_(self._values_list(True, True)) + return self + + def sinh(self) -> Self: + """Computes the :meth:`~torch.sinh` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_sinh(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def sinh_(self) -> Self: + """Computes the :meth:`~torch.sinh` value of each element of the TensorDict in-place.""" + torch._foreach_sinh_(self._values_list(True, True)) + return self + + def tan(self) -> Self: + """Computes the :meth:`~torch.tan` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_tan(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def tan_(self) -> Self: + """Computes the :meth:`~torch.tan` value of each element of the TensorDict in-place.""" + torch._foreach_tan_(self._values_list(True, True)) + return self + + def tanh(self) -> Self: + """Computes the :meth:`~torch.tanh` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_tanh(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def tanh_(self) -> Self: + """Computes the :meth:`~torch.tanh` value of each element of the TensorDict in-place.""" + torch._foreach_tanh_(self._values_list(True, True)) + return self + + def trunc(self) -> Self: + """Computes the :meth:`~torch.trunc` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_trunc(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def trunc_(self) -> Self: + """Computes the :meth:`~torch.trunc` value of each element of the TensorDict in-place.""" + torch._foreach_trunc_(self._values_list(True, True)) + return self + + @implement_for("torch", None, "2.4") + def norm( + self, + *, + out=None, + dtype: torch.dtype | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self: + """Computes the norm of each tensor in the tensordict. + + Keyword Args: + out (TensorDict, optional): the output tensordict. + dtype (torch.dtype, optional): the output dtype (torch>=2.4). + key_transform (Callable[[NestedKey], NestedKey], optional): A function to transform key names. + If provided, all keys in the result will be transformed using this function. + For string keys, the function receives a string. For tuple keys, it receives a tuple. + Default: ``None``. + + """ + keys, vals = self._items_list(True, True, collapse=True) + if dtype is not None: + raise RuntimeError("dtype must be None for torch <= 2.3") + vals = torch._foreach_norm(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + result = self._fast_apply( + get, + named=True, + nested_keys=True, + batch_size=[], + propagate_lock=True, + out=out, + ) + if key_transform is not None: + result = result._transform_keys(key_transform) + return result + + @implement_for("torch", "2.4") + def norm( # noqa: F811 + self, + *, + out=None, + dtype: torch.dtype | None = None, + key_transform: Callable[[NestedKey], NestedKey] | None = None, + ) -> Self: + """Computes the norm of each tensor in the tensordict. + + Keyword Args: + out (TensorDict, optional): the output tensordict. + dtype (torch.dtype, optional): the output dtype (torch>=2.4). + key_transform (Callable[[NestedKey], NestedKey], optional): A function to transform key names. + If provided, all keys in the result will be transformed using this function. + For string keys, the function receives a string. For tuple keys, it receives a tuple. + Default: ``None``. + + """ + keys, vals = self._items_list(True, True, collapse=True) + vals = torch._foreach_norm(vals, dtype=dtype) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + result = self._fast_apply( + get, + named=True, + nested_keys=True, + batch_size=[], + propagate_lock=True, + out=out, + ) + if key_transform is not None: + result = result._transform_keys(key_transform) + return result + + def lgamma(self) -> Self: + """Computes the :meth:`~torch.lgamma` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_lgamma(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def lgamma_(self) -> Self: + """Computes the :meth:`~torch.lgamma` value of each element of the TensorDict in-place.""" + torch._foreach_lgamma_(self._values_list(True, True)) + return self + + def frac(self) -> Self: + """Computes the :meth:`~torch.frac` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_frac(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def frac_(self) -> Self: + """Computes the :meth:`~torch.frac` value of each element of the TensorDict in-place.""" + torch._foreach_frac_(self._values_list(True, True)) + return self + + def expm1(self) -> Self: + """Computes the :meth:`~torch.expm1` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_expm1(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def expm1_(self) -> Self: + """Computes the :meth:`~torch.expm1` value of each element of the TensorDict in-place.""" + torch._foreach_expm1_(self._values_list(True, True)) + return self + + def log(self) -> Self: + """Computes the :meth:`~torch.log` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_log(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def log_(self) -> Self: + """Computes the :meth:`~torch.log` value of each element of the TensorDict in-place.""" + torch._foreach_log_(self._values_list(True, True)) + return self + + def logsumexp(self, dim=None, keepdim=False, *, out=None): # noqa: D417 + """Returns the log of summed exponentials of each row of the input tensordict in the given dimension ``dim``. The computation is numerically stabilized. + + If keepdim is ``True``, the output tensor is of the same size as input except in the dimension(s) ``dim`` where it is of size ``1``. + Otherwise, ``dim`` is squeezed (see :func:`~torch.squeeze`), resulting in the output tensor having 1 (or len(dim)) fewer dimension(s). + + Args: + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. If ``None``, all batch dimensions of the + tensordict are reduced. + keepdim (bool): whether the output tensordict has dim retained or not. + + Keyword Args: + out (TensorDictBase, optional): the output tensordict. + + """ + if isinstance(dim, int): + if dim < 0: + new_dim = (self.ndim + dim,) + else: + new_dim = (dim,) + elif dim is not None: + new_dim = tuple(self.ndim + _dim if _dim < 0 else _dim for _dim in dim) + else: + new_dim = tuple(range(self.ndim)) + if new_dim is not None and any((d < 0) or (d >= self.ndim) for d in new_dim): + raise ValueError( + f"The dimension {dim} is incompatible with a tensordict with batch_size {self.batch_size}." + ) + batch_size = self.batch_size + if keepdim: + batch_size = torch.Size( + [b if i not in new_dim else 1 for i, b in enumerate(batch_size)] + ) + else: + batch_size = torch.Size( + [b for i, b in enumerate(batch_size) if i not in new_dim] + ) + if out is not None: + result = self._fast_apply( + lambda x, y: torch.logsumexp(x, dim=new_dim, keepdim=keepdim, out=y), + out, + default=None, + batch_size=batch_size, + ) + return out.update(result) + + return self._fast_apply( + lambda x: torch.logsumexp(x, dim=new_dim, keepdim=keepdim), + batch_size=batch_size, + ) + + def softmax(self, dim: int, dtype: torch.dtype | None = None): # noqa: D417 + """Apply a softmax function to the tensordict elements. + + Args: + dim (int or tuple of ints): A tensordict dimension along which softmax will be computed. + dtype (torch.dtype, optional): the desired data type of returned tensor. + If specified, the input tensor is cast to dtype before the operation is performed. + This is useful for preventing data type overflows. + + """ + if isinstance(dim, int): + dim = _maybe_correct_neg_dim(dim, self.batch_size) + else: + raise ValueError(f"Expected dim of type int, got {type(dim)}.") + return self._fast_apply( + lambda x: torch.softmax(x, dim=dim, dtype=dtype), + ) + + def log10(self) -> Self: + """Computes the :meth:`~torch.log10` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_log10(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def log10_(self) -> Self: + """Computes the :meth:`~torch.log10` value of each element of the TensorDict in-place.""" + torch._foreach_log10_(self._values_list(True, True)) + return self + + def log1p(self) -> Self: + """Computes the :meth:`~torch.log1p` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_log1p(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def log1p_(self) -> Self: + """Computes the :meth:`~torch.log1p` value of each element of the TensorDict in-place.""" + torch._foreach_log1p_(self._values_list(True, True)) + return self + + def log2(self) -> Self: + """Computes the :meth:`~torch.log2` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_log2(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def log2_(self) -> Self: + """Computes the :meth:`~torch.log2` value of each element of the TensorDict in-place.""" + torch._foreach_log2_(self._values_list(True, True)) + return self + + def ceil(self) -> Self: + """Computes the :meth:`~torch.ceil` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_ceil(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def ceil_(self) -> Self: + """Computes the :meth:`~torch.ceil` value of each element of the TensorDict in-place.""" + torch._foreach_ceil_(self._values_list(True, True)) + return self + + def floor(self) -> Self: + """Computes the :meth:`~torch.floor` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_floor(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def floor_(self) -> Self: + """Computes the :meth:`~torch.floor` value of each element of the TensorDict in-place.""" + torch._foreach_floor_(self._values_list(True, True)) + return self + + def round(self) -> Self: + """Computes the :meth:`~torch.round` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_round(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def round_(self) -> Self: + """Computes the :meth:`~torch.round` value of each element of the TensorDict in-place.""" + torch._foreach_round_(self._values_list(True, True)) + return self + + def erf(self) -> Self: + """Computes the :meth:`~torch.erf` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_erf(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def erf_(self) -> Self: + """Computes the :meth:`~torch.erf` value of each element of the TensorDict in-place.""" + torch._foreach_erf_(self._values_list(True, True)) + return self + + def erfc(self) -> Self: + """Computes the :meth:`~torch.erfc` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_erfc(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def erfc_(self) -> Self: + """Computes the :meth:`~torch.erfc` value of each element of the TensorDict in-place.""" + torch._foreach_erfc_(self._values_list(True, True)) + return self + + def asin(self) -> Self: + """Computes the :meth:`~torch.asin` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_asin(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def asin_(self) -> Self: + """Computes the :meth:`~torch.asin` value of each element of the TensorDict in-place.""" + torch._foreach_asin_(self._values_list(True, True)) + return self + + def atan(self) -> Self: + """Computes the :meth:`~torch.atan` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_atan(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def atan_(self) -> Self: + """Computes the :meth:`~torch.atan` value of each element of the TensorDict in-place.""" + torch._foreach_atan_(self._values_list(True, True)) + return self + + def cos(self) -> Self: + """Computes the :meth:`~torch.cos` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_cos(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def cos_(self) -> Self: + """Computes the :meth:`~torch.cos` value of each element of the TensorDict in-place.""" + torch._foreach_cos_(self._values_list(True, True)) + return self + + def cosh(self) -> Self: + """Computes the :meth:`~torch.cosh` value of each element of the TensorDict.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_cosh(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def cosh_(self) -> Self: + """Computes the :meth:`~torch.cosh` value of each element of the TensorDict in-place.""" + torch._foreach_cosh_(self._values_list(True, True)) + return self + + @implement_for("torch", None, "2.5") + def _clone_recurse(self) -> Self: # noqa: D417 + keys, vals = self._items_list(True, True) + items = dict( + _zip_strict( + keys, + (val.clone() if hasattr(val, "clone") else val for val in vals), + ) + ) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=False, + filter_empty=False, + default=None, + ) + if items: + result.update(items) + return result + + @implement_for("torch", "2.5") + def _clone_recurse(self) -> Self: # noqa: F811, D417 + keys, vals = self._items_list(True, True) + foreach_vals = {} + iter_vals = {} + for key, val in zip(keys, vals): + if ( + type(val) is torch.Tensor + and not val.requires_grad + and val.dtype not in (torch.bool,) + ): + foreach_vals[key] = val + else: + iter_vals[key] = val + if foreach_vals: + foreach_vals = dict( + _zip_strict( + foreach_vals.keys(), + torch._foreach_add(tuple(foreach_vals.values()), 0), + ) + ) + if iter_vals: + iter_vals = dict( + _zip_strict( + iter_vals.keys(), + ( + val.clone() if hasattr(val, "clone") else val + for val in iter_vals.values() + ), + ) + ) + + items = foreach_vals + items.update(iter_vals) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=False, + filter_empty=False, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("bitwise_and") + def bitwise_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Performs a bitwise AND operation between ``self`` and :attr:`other`. + + .. math:: + \text{{out}}_i = \text{{input}}_i \land \text{{other}}_i + + Args: + other (TensorDictBase or torch.Tensor): the tensor or TensorDict to perform the bitwise AND with. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + vals = [(v1.bitwise_and(v2)) for v1, v2 in zip(vals, other_val)] + else: + vals = [v.bitwise_and(other) for v in vals] + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("logical_and") + def logical_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Performs a logical AND operation between ``self`` and :attr:`other`. + + .. math:: + \text{{out}}_i = \text{{input}}_i \land \text{{other}}_i + + Args: + other (TensorDictBase or torch.Tensor): the tensor or TensorDict to perform the logical AND with. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + vals = [(v1.logical_and(v2)) for v1, v2 in zip(vals, other_val)] + else: + vals = [v.logical_and(other) for v in vals] + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("add") + def add( + self, + other: TensorCollection | torch.Tensor, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Adds :attr:`other`, scaled by :attr:`alpha`, to ``self``. + + .. math:: + \text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i + + Args: + other (TensorDictBase or torch.Tensor): the tensor or TensorDict to add to ``self``. + + Keyword Args: + alpha (Number, optional): the multiplier for :attr:`other`. + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + if alpha is not None: + vals = torch._foreach_add(vals, other_val, alpha=alpha) + else: + vals = torch._foreach_add(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("add_") + def add_( + self, + other: TensorCollection | torch.Tensor | float, + *, + alpha: float | None = None, + ) -> Self: + """In-place version of :meth:`~.add`. + + .. note:: + In-place ``add`` does not support ``default`` keyword argument. + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + if alpha is not None: + torch._foreach_add_(vals, other_val, alpha=alpha) + else: + torch._foreach_add_(vals, other_val) + return self + + @_maybe_broadcast_other("lerp", 2) + def lerp( + self, + end: TensorCollection | torch.Tensor, + weight: TensorCollection | torch.Tensor | float, + ) -> Self: + r"""Does a linear interpolation of two tensors :attr:`start` (given by ``self``) and :attr:`end` based on a scalar or tensor :attr:`weight`. + + .. math:: + \text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i) + + The shapes of :attr:`start` and :attr:`end` must be + broadcastable. If :attr:`weight` is a tensor, then + the shapes of :attr:`weight`, :attr:`start`, and :attr:`end` must be broadcastable. + + Args: + end (TensorDict): the tensordict with the ending points. + weight (TensorDict, tensor or float): the weight for the interpolation formula. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(end)): + end_val = end._values_list(True, True) + else: + end_val = end + if isinstance(weight, (float, torch.Tensor)): + weight_val = weight + elif _is_tensor_collection(type(weight)): + weight_val = weight._values_list(True, True) + else: + weight_val = weight + vals = torch._foreach_lerp(vals, end_val, weight_val) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + def lerp_( + self, + end: TensorDictBase | torch.Tensor | float, + weight: TensorDictBase | torch.Tensor | float, + ): + """In-place version of :meth:`~.lerp`.""" + if _is_tensor_collection(type(end)): + end_val = end._values_list(True, True) + else: + end_val = end + if isinstance(weight, (float, torch.Tensor)): + weight_val = weight + elif _is_tensor_collection(type(weight)): + weight_val = weight._values_list(True, True) + else: + weight_val = weight + torch._foreach_lerp_(self._values_list(True, True), end_val, weight_val) + return self + + @_maybe_broadcast_other("addcdiv", 2) + def addcdiv( + self, + other1: TensorDictBase | torch.Tensor, + other2: TensorDictBase | torch.Tensor, + value: float | None = 1, + ) -> Self: # noqa: D417 + r"""Performs the element-wise division of :attr:`other1` by :attr:`other2`, multiplies the result by the scalar :attr:`value` and adds it to ``self``. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} + + The shapes of the elements of ``self``, :attr:`other1`, and :attr:`other2` must be + broadcastable. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + other1 (TensorDict or Tensor): the numerator tensordict (or tensor) + tensor2 (TensorDict or Tensor): the denominator tensordict (or tensor) + + Keyword Args: + value (Number, optional): multiplier for :math:`\text{tensor1} / \text{tensor2}` + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other1)): + other1_val = other1._values_list(True, True) + else: + other1_val = other1 + if _is_tensor_collection(type(other2)): + other2_val = other2._values_list(True, True) + else: + other2_val = other2 + vals = torch._foreach_addcdiv(vals, other1_val, other2_val, value=value) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + @_maybe_broadcast_other("addcdiv_", 2) + def addcdiv_(self, other1, other2, *, value: float | None = 1): + """The in-place version of :meth:`~.addcdiv`.""" + if _is_tensor_collection(type(other1)): + other1_val = other1._values_list(True, True) + else: + other1_val = other1 + if _is_tensor_collection(type(other2)): + other2_val = other2._values_list(True, True) + else: + other2_val = other2 + torch._foreach_addcdiv_( + self._values_list(True, True), other1_val, other2_val, value=value + ) + return self + + @_maybe_broadcast_other("addcmul", 2) + def addcmul( + self, + other1: TensorDictBase | torch.Tensor, + other2: TensorDictBase | torch.Tensor, + *, + value: float | None = 1, + ) -> Self: # noqa: D417 + r"""Performs the element-wise multiplication of :attr:`other1` by :attr:`other2`, multiplies the result by the scalar :attr:`value` and adds it to ``self``. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \text{other1}_i \times \text{other2}_i + + The shapes of ``self``, :attr:`other1`, and :attr:`other2` must be + broadcastable. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + other1 (TensorDict or Tensor): the tensordict or tensor to be multiplied + other2 (TensorDict or Tensor): the tensordict or tensor to be multiplied + + Keyword Args: + value (Number, optional): multiplier for :math:`other1 .* other2` + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other1)): + other1_val = other1._values_list(True, True) + else: + other1_val = other1 + if _is_tensor_collection(type(other2)): + other2_val = other2._values_list(True, True) + else: + other2_val = other2 + vals = torch._foreach_addcmul(vals, other1_val, other2_val, value=value) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + @_maybe_broadcast_other("addcmul_", 2) + def addcmul_(self, other1, other2, *, value: float | None = 1): + """The in-place version of :meth:`~.addcmul`.""" + if _is_tensor_collection(type(other1)): + other1_val = other1._values_list(True, True) + else: + other1_val = other1 + if _is_tensor_collection(type(other2)): + other2_val = other2._values_list(True, True) + else: + other2_val = other2 + torch._foreach_addcmul_( + self._values_list(True, True), other1_val, other2_val, value=value + ) + return self + + @_maybe_broadcast_other("sub") + def sub( + self, + other: TensorDictBase | torch.Tensor | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Subtracts :attr:`other`, scaled by :attr:`alpha`, from ``self``. + + .. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i + + Supports broadcasting, + type promotion, and integer, float, and complex inputs. + + Args: + other (TensorDict, Tensor or Number): the tensor or number to subtract from ``self``. + + Keyword Args: + alpha (Number): the multiplier for :attr:`other`. + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + if alpha is not None: + vals = torch._foreach_sub(vals, other_val, alpha=alpha) + else: + vals = torch._foreach_sub(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("sub_") + def sub_( + self, other: TensorDictBase | torch.Tensor | float, alpha: float | None = None + ): + """In-place version of :meth:`~.sub`. + + .. note:: + In-place ``sub`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + if alpha is not None: + torch._foreach_sub_(vals, other_val, alpha=alpha) + else: + torch._foreach_sub_(vals, other_val) + return self + + @_maybe_broadcast_other("rsub") + def rsub( + self, + other: TensorDictBase | torch.Tensor | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Subtracts `self` from :attr:`other`, scaled by :attr:`alpha`, from ``self``. + + .. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i + + Supports broadcasting, + type promotion, and integer, float, and complex inputs. + + Args: + other (TensorDict, Tensor or Number): the tensor or number to subtract from ``self``. + + Keyword Args: + alpha (Number): the multiplier for :attr:`other`. + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + if alpha is not None: + vals = torch._foreach_neg(torch._foreach_sub(vals, other_val, alpha=alpha)) + else: + vals = torch._foreach_neg(torch._foreach_sub(vals, other_val)) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("mod") + def mod(self, other: TensorCollection | torch.Tensor) -> Self: + """Computes the element-wise modulo of ``self`` and :attr:`other`. + + Args: + other (TensorDict or Tensor): the other input tensordict or tensor. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list(True, True, sorting_keys=keys) + else: + other_val = other + if isinstance(other_val, list): + vals = [val % other_val for val, other_val in zip(vals, other_val)] + else: + vals = [val % other_val for val in vals] + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("mul_") + def mul_(self, other: TensorCollection | torch.Tensor) -> Self: + """In-place version of :meth:`~.mul`. + + .. note:: + Inplace ``mul`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + torch._foreach_mul_(vals, other_val) + return self + + @_maybe_broadcast_other("mul") + def mul( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Multiplies :attr:`other` to ``self``. + + .. math:: + \text{{out}}_i = \text{{input}}_i \times \text{{other}}_i + + Supports broadcasting, type promotion, and integer, float, and complex inputs. + + Args: + other (TensorDict, Tensor or Number): the tensor or number to subtract from ``self``. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + vals = torch._foreach_mul(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + def maximum_(self, other: TensorCollection | torch.Tensor) -> Self: + """In-place version of :meth:`~.maximum`. + + .. note:: + Inplace ``maximum`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + torch._foreach_maximum_(vals, other_val) + return self + + @_maybe_broadcast_other("maximum") + def maximum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + """Computes the element-wise maximum of ``self`` and :attr:`other`. + + Args: + other (TensorDict or Tensor): the other input tensordict or tensor. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + vals = torch._foreach_maximum(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + def minimum_(self, other: TensorCollection | torch.Tensor) -> Self: + """In-place version of :meth:`~.minimum`. + + .. note:: + Inplace ``minimum`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + torch._foreach_minimum_(vals, other_val) + return self + + @_maybe_broadcast_other("minimum") + def minimum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + """Computes the element-wise minimum of ``self`` and :attr:`other`. + + Args: + other (TensorDict or Tensor): the other input tensordict or tensor. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + vals = torch._foreach_minimum(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + def clamp_max_(self, other: TensorCollection | torch.Tensor) -> Self: + """In-place version of :meth:`~.clamp_max`. + + .. note:: + Inplace ``clamp_max`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + try: + torch._foreach_clamp_max_(vals, other_val) + except RuntimeError as err: + if "isDifferentiableType" in str(err): + raise RuntimeError( + "Attempted to execute _foreach_clamp_max_ with a differentiable tensor. " + "Use `td.apply(lambda x: x.clamp_max_(val)` instead." + ) + return self + + @_maybe_broadcast_other("clamp_max") + def clamp_max( + self, + other: TensorDictBase | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + """Clamps the elements of ``self`` to :attr:`other` if they're superior to that value. + + Args: + other (TensorDict or Tensor): the other input tensordict or tensor. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + try: + vals = torch._foreach_clamp_max(vals, other_val) + except RuntimeError as err: + if "isDifferentiableType" in str(err): + raise RuntimeError( + "Attempted to execute _foreach_clamp_max with a differentiable tensor. " + "Use `td.apply(lambda x: x.clamp_max(val)` instead." + ) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + def clamp_min_(self, other: TensorDictBase | torch.Tensor) -> Self: + """In-place version of :meth:`~.clamp_min`. + + .. note:: + Inplace ``clamp_min`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + try: + torch._foreach_clamp_min_(vals, other_val) + except RuntimeError as err: + if "isDifferentiableType" in str(err): + raise RuntimeError( + "Attempted to execute _foreach_clamp_min_ with a differentiable tensor. " + "Use `td.apply(lambda x: x.clamp_min_(val)` instead." + ) + + return self + + @_maybe_broadcast_other("clamp_min") + def clamp_min( + self, + other: TensorDictBase | torch.Tensor, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + """Clamps the elements of ``self`` to :attr:`other` if they're inferior to that value. + + Args: + other (TensorDict or Tensor): the other input tensordict or tensor. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + try: + vals = torch._foreach_clamp_min(vals, other_val) + except RuntimeError as err: + if "isDifferentiableType" in str(err): + raise RuntimeError( + "Attempted to execute _foreach_clamp_min with a differentiable tensor. " + "Use `td.apply(lambda x: x.clamp_min(val)` instead." + ) + + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("clamp", 2) + def clamp( + self, + min: TensorDictBase | torch.Tensor = None, + max: TensorDictBase | torch.Tensor = None, + *, + out=None, + ) -> Self: # noqa: D417, W605 + r"""Clamps all elements in :attr:`self` into the range `[` :attr:`min`, :attr:`max` `]`. + + Letting min_value and max_value be :attr:`min` and :attr:`max`, respectively, this returns: + + .. math:: + y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i) + + If :attr:`min` is ``None``, there is no lower bound. + Or, if :attr:`max` is ``None`` there is no upper bound. + + .. note:: + If :attr:`min` is greater than :attr:`max` :func:`torch.clamp(..., min, max) ` + sets all elements in :attr:`input` to the value of :attr:`max`. + + """ + if min is None: + if out is not None: + raise ValueError( + "clamp() with min/max=None isn't implemented with specified output." + ) + return self.clamp_max(max) + if max is None: + if out is not None: + raise ValueError( + "clamp() with min/max=None isn't implemented with specified output." + ) + return self.clamp_min(min) + + is_tc_min = is_tensor_collection(min) + is_tc_max = is_tensor_collection(max) + + if is_tc_min ^ is_tc_max: + raise ValueError( + "Mixed tensordict and non-tensordict min/max values are not authorized." + ) + + if out is None: + if is_tc_min and is_tc_max: + return self._fast_apply( + lambda x, low, high: x.clamp(low, high), min, max, default=None + ) + return self._fast_apply(lambda x: x.clamp(min, max)) + if is_tc_min and is_tc_max: + result = self._fast_apply( + lambda x, y, low, high: x.clamp(low, high, out=y), + out, + min, + max, + default=None, + ) + else: + result = self._fast_apply( + lambda x, y: x.clamp(min, max, out=y), out, default=None + ) + with out.unlock_() if out.is_locked else contextlib.nullcontext(): + return out.update(result) + + @_maybe_broadcast_other("pow_") + def pow_(self, other: TensorDictBase | torch.Tensor) -> Self: + """In-place version of :meth:`~.pow`. + + .. note:: + Inplace ``pow`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + torch._foreach_pow_(vals, other_val) + return self + + @_maybe_broadcast_other("pow") + def pow( + self, + other: TensorDictBase | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Takes the power of each element in ``self`` with :attr:`other` and returns a tensor with the result. + + :attr:`other` can be either a single ``float`` number, a `Tensor` or a ``TensorDict``. + + When :attr:`other` is a tensor, the shapes of :attr:`input` + and :attr:`other` must be broadcastable. + + Args: + other (float, tensor or tensordict): the exponent value + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + vals = torch._foreach_pow(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + @_maybe_broadcast_other("div_") + def div_(self, other: TensorDictBase | torch.Tensor) -> Self: + """In-place version of :meth:`~.div`. + + .. note:: + Inplace ``div`` does not support ``default`` keyword argument. + + """ + if _is_tensor_collection(type(other)): + keys, vals = self._items_list(True, True) + other_val = other._values_list(True, True, sorting_keys=keys) + else: + vals = self._values_list(True, True) + other_val = other + torch._foreach_div_(vals, other_val) + return self + + @_maybe_broadcast_other("div") + def div( + self, + other: TensorDictBase | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: # noqa: D417 + r"""Divides each element of the input ``self`` by the corresponding element of :attr:`other`. + + .. math:: + \text{out}_i = \frac{\text{input}_i}{\text{other}_i} + + Supports broadcasting, type promotion and integer, float, tensordict or tensor inputs. + Always promotes integer types to the default scalar type. + + Args: + other (TensorDict, Tensor or Number): the divisor. + + Keyword Args: + default (torch.Tensor or str, optional): the default value to use for exclusive entries. + If none is provided, the two tensordicts key list must match exactly. + If ``default="intersection"`` is passed, only the intersecting key sets will be considered + and other keys will be ignored. + In all other cases, ``default`` will be used for all missing entries on both sides of the + operation. + + """ + keys, vals = self._items_list(True, True) + if _is_tensor_collection(type(other)): + new_keys, other_val = other._items_list( + True, True, sorting_keys=keys, default=default + ) + if default is not None: + as_dict = dict(zip(keys, vals)) + vals = [as_dict.get(key, default) for key in new_keys] + keys = new_keys + else: + other_val = other + vals = torch._foreach_div(vals, other_val) + items = dict(zip(keys, vals)) + + def pop(name, val): + return items.pop(name, None) + + result = self._fast_apply( + pop, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + filter_empty=True, + default=None, + ) + if items: + result.update(items) + return result + + def sqrt_(self) -> Self: + """In-place version of :meth:`~.sqrt`.""" + torch._foreach_sqrt_(self._values_list(True, True)) + return self + + def sqrt(self) -> Self: + """Computes the element-wise square root of ``self``.""" + keys, vals = self._items_list(True, True) + vals = torch._foreach_sqrt(vals) + items = dict(zip(keys, vals)) + + def get(name, val): + return items.get(name, val) + + return self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + ) + + # Functorch compatibility + @abc.abstractmethod + @cache # noqa: B019 + def _add_batch_dim(self, *, in_dim: int, vmap_level: int) -> Self: + raise NotImplementedError + + @abc.abstractmethod + @cache # noqa: B019 + def _remove_batch_dim(self, vmap_level: int, batch_size: int, out_dim: int) -> Self: + raise NotImplementedError + + @abc.abstractmethod + @cache # noqa: B019 + def _maybe_remove_batch_dim( + self, funcname: str, vmap_level: int, batch_size: int, out_dim: int + ) -> Self: + raise NotImplementedError + + # Validation and checks + def _convert_to_tensor( + self, array: Any + ) -> Tensor | "NonTensorData" | TensorDictBase: # noqa: F821 + # We are sure that array is not a dict or anything in _ACCEPTED_CLASSES + castable = None + if isinstance(array, (float, int, bool)): + castable = True + elif isinstance(array, list) and list_to_stack(): + return _convert_list_to_stack(array)[0] + elif isinstance(array, np.bool_): + castable = True + array = array.item() + elif isinstance(array, (np.ndarray, np.number)): + if array.dtype.names is not None: + return TensorDictBase.from_struct_array(array, device=self.device) + castable = array.dtype.kind in ("c", "i", "f", "b", "u") + elif isinstance(array, (list, tuple)): + array = np.asarray(array) + castable = array.dtype.kind in ("c", "i", "f", "b", "u") + elif hasattr(array, "numpy"): + # tf.Tensor with no shape can't be converted otherwise + array = array.numpy() + castable = array.dtype.kind in ("c", "i", "f", "b", "u") + if castable: + if hasattr(array, "flags") and not array.flags.writeable: + array = array.copy() + return torch.as_tensor(array, device=self.device) + else: + from tensordict.tensorclass import NonTensorData + + return NonTensorData( + data=array, + batch_size=self.batch_size, + device=self.device, + names=self._maybe_names(), + ) + + @abc.abstractmethod + def _convert_to_tensordict(self, dict_value: dict[str, Any]) -> Self: + raise NotImplementedError + + def _check_batch_size(self, *, raise_exception: bool = True) -> None | bool: + batch_dims = self.batch_dims + val = True + for value in self.values(): + if _is_tensor_collection(type(value)): + val &= value._check_batch_size(raise_exception=raise_exception) + if not val: + return False + val &= _shape(value)[:batch_dims] == self.batch_size + if not val: + if raise_exception: + raise RuntimeError( + f"batch_size are incongruent, got value with shape {_shape(value)}, " + f"-- expected {self.batch_size}" + ) + return False + return val + + @abc.abstractmethod + def _check_is_shared(self) -> bool: + raise NotImplementedError + + def _check_new_batch_size(self, new_size: torch.Size) -> None: + batch_dims = len(new_size) + for key, tensor in self.items(): + if _shape(tensor)[:batch_dims] != new_size and not ( + _is_tensor_collection(type(tensor)) and tensor.is_empty() + ): + raise RuntimeError( + f"the {type(tensor).__name__} {key} has shape {_shape(tensor)} which " + f"is incompatible with the batch-size {new_size}." + ) + + @abc.abstractmethod + def _check_device(self, *, raise_exception: bool = True) -> None | bool: + raise NotImplementedError + + def _validate_key(self, key: NestedKey) -> NestedKey: + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + return key + + @property + def _validate_value(self): + if is_compiling(): + return self._validate_value_generic + if self.device: + if self.batch_size: + method_name = "_validate_value_generic" + else: + method_name = "_validate_value_batchfree" + else: + if self.batch_size: + method_name = "_validate_value_devicefree" + else: + method_name = "_validate_value_batchfree_devicefree" + return getattr(self, method_name) + + def _validate_value_generic( + self, + value: CompatibleType | dict[str, CompatibleType], + non_blocking: bool = False, + *, + check_shape: bool = True, + ) -> CompatibleType | dict[str, CompatibleType]: + cls = type(value) + if issubclass(cls, torch.Tensor): + is_tc = False + elif _is_tensor_collection(cls): + is_tc = True + elif issubclass(cls, dict): + # We use non-blocking if someone's watching or if non-blocking is explicitly passed + value = self._convert_to_tensordict( + value, non_blocking=_device_recorder.marked or non_blocking + ) + is_tc = True + elif not issubclass(cls, _ACCEPTED_CLASSES): + # If cls is not a tensor + try: + value = self._convert_to_tensor(value) + except ValueError as err: + raise ValueError( + f"TensorDict conversion only supports tensorclasses, tensordicts," + f" numeric scalars and tensors. Got {type(value)}" + ) from err + is_tc = _is_tensor_collection(cls) + batch_size = self.batch_size + if check_shape and _shape(value)[: self.batch_dims] != batch_size: + # if TensorDict, let's try to map it to the desired shape + if is_tc: + # we must clone the value before not to corrupt the data passed to set() + value = value.clone(recurse=False) + value.batch_size = self.batch_size + else: + raise RuntimeError( + f"batch dimension mismatch, got self.batch_size" + f"={self.batch_size} and value.shape={_shape(value)}." + ) + device = self.device + if device is not None and value.device != device: + if _device_recorder.marked and device.type != "cuda": + _device_recorder.record_transfer(device) + value = value.to(device, non_blocking=non_blocking) + if check_shape: + if not is_tc: + return value + has_names = self._has_names() + # we do our best to match the dim names of the value and the + # container. + if has_names: + if value.names[: self.batch_dims] != self.names: + # we clone not to corrupt the value + value = value.clone(False).refine_names(*self.names) + else: + if value._has_names(): + self._set_names(value.names[: self.batch_dims]) + return value + + def _validate_value_batchfree( + self, + value: CompatibleType | dict[str, CompatibleType], + non_blocking: bool = False, + *, + check_shape: bool = True, + ) -> CompatibleType | dict[str, CompatibleType]: + cls = type(value) + if issubclass(cls, torch.Tensor) or _is_tensor_collection(cls): + pass + elif issubclass(cls, dict): + # We use non-blocking if someone's watching or if non-blocking is explicitly passed + value = self._convert_to_tensordict( + value, non_blocking=_device_recorder.marked or non_blocking + ) + elif not issubclass(cls, _ACCEPTED_CLASSES): + # If cls is not a tensor + try: + value = self._convert_to_tensor(value) + except ValueError as err: + raise ValueError( + f"TensorDict conversion only supports tensorclasses, tensordicts," + f" numeric scalars and tensors. Got {type(value)}" + ) from err + device = self.device + if device is not None and value.device != device: + if _device_recorder.marked and device.type != "cuda": + _device_recorder.record_transfer(device) + value = value.to(device, non_blocking=non_blocking) + return value + + def _validate_value_devicefree( + self, + value: CompatibleType | dict[str, CompatibleType], + non_blocking: bool = False, + *, + check_shape: bool = True, + ) -> CompatibleType | dict[str, CompatibleType]: + cls = type(value) + if issubclass(cls, torch.Tensor): + is_tc = False + elif _is_tensor_collection(cls): + is_tc = True + elif issubclass(cls, dict): + # We use non-blocking if someone's watching or if non-blocking is explicitly passed + value = self._convert_to_tensordict( + value, non_blocking=_device_recorder.marked or non_blocking + ) + is_tc = True + elif not issubclass(cls, _ACCEPTED_CLASSES): + # If cls is not a tensor + try: + value = self._convert_to_tensor(value) + except ValueError as err: + raise ValueError( + f"TensorDict conversion only supports tensorclasses, tensordicts," + f" numeric scalars and tensors. Got {type(value)}" + ) from err + is_tc = _is_tensor_collection(cls) + + batch_size = self.batch_size + if check_shape and _shape(value)[: self.batch_dims] != batch_size: + # if TensorDict, let's try to map it to the desired shape + if is_tc: + # we must clone the value before not to corrupt the data passed to set() + value = value.clone(recurse=False) + value.batch_size = self.batch_size + else: + raise RuntimeError( + f"batch dimension mismatch, got self.batch_size" + f"={self.batch_size} and value.shape={_shape(value)}." + ) + if check_shape: + if not is_tc: + return value + has_names = self._has_names() + # we do our best to match the dim names of the value and the + # container. + if has_names: + if value.names[: self.batch_dims] != self.names: + # we clone not to corrupt the value + value = value.clone(False).refine_names( + *(self.names + value.names[self.batch_dims :]) + ) + else: + if value._has_names(): + self._set_names(value.names[: self.batch_dims]) + return value + + def _validate_value_batchfree_devicefree( + self, + value: CompatibleType | dict[str, CompatibleType], + non_blocking: bool = False, + *, + check_shape: bool = True, + ) -> CompatibleType | dict[str, CompatibleType]: + cls = type(value) + if issubclass(cls, torch.Tensor) or _is_tensor_collection(cls): + pass + elif issubclass(cls, dict): + # We use non-blocking if someone's watching or if non-blocking is explicitly passed + value = self._convert_to_tensordict( + value, non_blocking=_device_recorder.marked or non_blocking + ) + elif not issubclass(cls, _ACCEPTED_CLASSES): + # If cls is not a tensor + try: + value = self._convert_to_tensor(value) + except ValueError as err: + raise ValueError( + f"TensorDict conversion only supports tensorclasses, tensordicts," + f" numeric scalars and tensors. Got {type(value)}" + ) from err + return value + + def __enter__(self): + is_tc = _is_tensorclass(type(self)) + if not hasattr(self, "_last_op_queue"): + if is_tc: + _last_op_queue = self._tensordict._last_op_queue = collections.deque() + else: + _last_op_queue = self._last_op_queue = collections.deque() + else: + _last_op_queue = ( + self._last_op_queue + if not _is_tensorclass(type(self)) + else self._tensordict._last_op_queue + ) + if is_tc: + # get last-op from tensordict - that's where it's written + _last_op = self._tensordict._last_op + else: + _last_op = self._last_op + _last_op_queue.append(_last_op) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + # During exit, updates mustn't be made in-place as the source and dest + # storage location can be identical, resulting in a RuntimeError + if is_compiling(): + self.clear_refs_for_compile_() + if exc_type is not None and issubclass(exc_type, Exception): + return False + is_tc = _is_tensorclass(type(self)) + _last_op = ( + self._last_op_queue.pop() + if not is_tc + else self._tensordict._last_op_queue.pop() + ) + if _last_op is not None: + last_op, (args, kwargs, out_wr) = _last_op + # TODO: transpose, flatten etc. as decorator should lock the content to make sure that no key is + # added or deleted + _inv_caller = LAST_OP_MAPS.get(last_op) + if _inv_caller is not None: + prev_ref = out_wr() + result = _inv_caller(self, args, kwargs, prev_ref) + return result + else: + raise NotImplementedError(f"Unrecognised function {last_op}.") + return self + + def clear_refs_for_compile_(self) -> Self: + """Clears the weakrefs in order for the tensordict to get out of the compile region safely. + + Use this whenever you hit `torch._dynamo.exc.Unsupported: reconstruct: WeakRefVariable()` + before returning a TensorDict. + + Returns: self + """ + self._last_op = None + for v in self.values(True, True, is_leaf=_is_tensor_collection): + if _is_tensorclass(type(v)): + v = v._tensordict + v._last_op = None + return self + + # Clone, select, exclude, empty + def select( + self, *keys: NestedKey, inplace: bool = False, strict: bool = True + ) -> Self: + """Selects the keys of the tensordict and returns a new tensordict with only the selected keys. + + The values are not copied: in-place modifications a tensor of either + of the original or new tensordict will result in a change in both + tensordicts. + + Args: + *keys (str): keys to select + inplace (bool): if True, the tensordict is pruned in place. + Default is ``False``. + strict (bool, optional): whether selecting a key that is not present + will return an error or not. Default: :obj:`True`. + + Returns: + A new tensordict (or the same if ``inplace=True``) with the selected keys only. + + .. note:: + To select keys in a tensordict and return a version of this tensordict + deprived of these keys, see the :meth:`~.split_keys` method. + + Examples: + >>> from tensordict import TensorDict + >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) + >>> td.select("a", ("b", "c")) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.select("a", "b") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.select("this key does not exist", strict=False) + TensorDict( + fields={ + }, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + keys = unravel_key_list(keys) + result = self._select(*keys, inplace=inplace, strict=strict) + if not inplace and (result._is_memmap or result._is_shared): + result.lock_() + return result + + @abc.abstractmethod + def _select( + self, + *keys: NestedKey, + inplace: bool = False, + strict: bool = True, + set_shared: bool = True, + ) -> Self: + raise NotImplementedError + + def exclude(self, *keys: NestedKey, inplace: bool = False) -> Self: + """Excludes the keys of the tensordict and returns a new tensordict without these entries. + + The values are not copied: in-place modifications a tensor of either + of the original or new tensordict will result in a change in both + tensordicts. + + Args: + *keys (str): keys to exclude. + inplace (bool): if True, the tensordict is pruned in place. + Default is ``False``. + + Returns: + A new tensordict (or the same if ``inplace=True``) without the excluded entries. + + Examples: + >>> from tensordict import TensorDict + >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) + >>> td.exclude("a", ("b", "c")) + TensorDict( + fields={ + b: TensorDict( + fields={ + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> td.exclude("a", "b") + TensorDict( + fields={ + }, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + keys = unravel_key_list(keys) + result = self._exclude(*keys, inplace=inplace) + if not inplace and (result._is_memmap or result._is_shared): + result.lock_() + return result + + @abc.abstractmethod + def _exclude( + self, + *keys: NestedKey, + inplace: bool = False, + set_shared: bool = True, + ) -> Self: + raise NotImplementedError + + def _maybe_set_shared_attributes(self, result, lock=False): + # We must use _is_shared to avoid having issues with CUDA tensordicts + if self._is_shared: + result._is_shared = True + if lock: + result.lock_() + elif self._is_memmap: + result._is_memmap = True + if lock: + result.lock_() + + def to_tensordict(self, *, retain_none: bool | None = None) -> Self: + """Returns a regular TensorDict instance from the TensorDictBase. + + Args: + retain_none (bool): if ``True``, the ``None`` values from tensorclass instances + will be written in the tensordict. + Otherwise they will be discarded. Default: ``True``. + + Returns: + a new TensorDict object containing the same values. + + """ + from tensordict import TensorDict + + return TensorDict( + { + key: ( + value.clone() + if not _is_tensor_collection(type(value)) + else ( + value + if is_non_tensor(value) + else value.to_tensordict(retain_none=retain_none) + ) + ) + for key, value in self.items(is_leaf=_is_leaf_nontensor) + }, + device=self.device, + batch_size=self.batch_size, + names=self._maybe_names(), + ) + + def to_lazystack(self, dim: int = 0): + """Converts a TensorDict to a LazyStackedTensorDict or equivalent. + + .. note:: + This method can be used to swap the stack dimension of a LazyStackedTensorDict. + For example, if you have a LazyStackedTensorDict with stack_dim=1, you can use this method to swap it to stack_dim=0: + + >>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3)) + >>> td2 = td.to_lazystack() + >>> td2.batch_size + torch.Size([2, 3]) + >>> assert isinstance(td2, LazyStackedTensorDict) + >>> assert td2.stack_dim == 0 + >>> td3 = td2.to_lazystack(1) + >>> assert td3.stack_dim == 1 + >>> td3.batch_size + torch.Size([2, 3]) + + Args: + dim (int, optional): the dimension along which to stack the tensordict. + Defaults to ``0``. + + Returns: + A LazyStackedTensorDict instance. + + Examples: + >>> from tensordict import TensorDict + >>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3)) + >>> td2 = td.to_lazystack() + >>> td2.batch_size + torch.Size([2, 3]) + >>> assert isinstance(td2, LazyStackedTensorDict) + + """ + from tensordict import lazy_stack, LazyStackedTensorDict + from tensordict.tensorclass import _is_tensorclass + + dim = _maybe_correct_neg_dim(dim, ndim=self.ndim, shape=None) + if (isinstance(self, LazyStackedTensorDict) and self.stack_dim == dim) or ( + _is_tensorclass(type(self)) + and isinstance(self._tensordict, LazyStackedTensorDict) + and self._tensordict.stack_dim == dim + ): + return self + return lazy_stack(self.unbind(dim), dim=dim) + + def clone(self, recurse: bool = True, **kwargs) -> Self: + """Clones a TensorDictBase subclass instance onto a new TensorDictBase subclass of the same type. + + To create a TensorDict instance from any other TensorDictBase subtype, call the :meth:`~.to_tensordict` method + instead. + + Args: + recurse (bool, optional): if ``True``, each tensor contained in the + TensorDict will be copied too. Otherwise only the TensorDict + tree structure will be copied. Defaults to ``True``. + + .. note:: + Unlike many other ops (pointwise arithmetic, shape operations, ...) ``clone`` does not inherit the + original lock attribute. This design choice is made such that a clone can be created to be modified, + which is the most frequent usage. + + """ + result = self._clone(recurse=recurse, **kwargs) + if not recurse and (result._is_shared or result._is_memmap): + result.lock_() + return result + + @abc.abstractmethod + def _clone(self, recurse: bool = False): + raise NotImplementedError + + def to_padded_tensor(self, padding=0.0, mask_key: NestedKey | None = None) -> Self: + """Converts all nested tensors to a padded version and adapts the batch-size accordingly. + + Args: + padding (float): the padding value for the tensors in the tensordict. + Defaults to ``0.0``. + mask_key (NestedKey, optional): if provided, the key where a + mask for valid values will be written. + Will result in an error if the heterogeneous dimension + isn't part of the tensordict batch-size. + Defaults to ``None`` + + """ + batch_size = self.batch_size + if any(shape == -1 for shape in batch_size): + new_batch_size = [] + else: + new_batch_size = None + if mask_key is not None: + raise RuntimeError( + "mask_key should only be provided if the " + "heterogenous dimension is part of the batch-size." + ) + padded_names = [] + + def to_padded(name, x): + if x.is_nested: + padded_names.append(name) + return torch.nested.to_padded_tensor(x, padding=padding) + return x + + result = self._apply_nest( + to_padded, + batch_size=new_batch_size, + named=True, + nested_keys=True, + ) + if new_batch_size is not None: + result = result.auto_batch_size_( + batch_dims=self.batch_dims, keep_compliant_size=True + ) + + if mask_key: + # take the first of the padded keys + padded_key = padded_names[0] + # write the mask + val = self.get(padded_key) + val = torch.nested.to_padded_tensor( + torch.ones_like(val, dtype=torch.bool), padding=False + ) + if val.ndim > result.ndim: + val = val.flatten(result.ndim, -1)[..., -1].clone() + result.set(mask_key, val) + return result + + def as_tensor(self) -> Self: + """Converts every leaf of a tensordict to a plain torch.Tensor.""" + + def as_tensor(tensor): + try: + return tensor.as_tensor() + except AttributeError: + return tensor + + return self._fast_apply(as_tensor, propagate_lock=True) + + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + ) -> dict[str, Any]: + """Returns a dictionary with key-value pairs matching those of the tensordict. + + Args: + retain_none (bool): if ``True``, the ``None`` values from tensorclass instances + will be written in the dictionary. + Otherwise, they will be discarded. Default: ``True``. + convert_tensors (bool, "numpy"): if ``True``, tensors will be converted to lists when creating the dictionary. + If "numpy", tensors will be converted to numpy arrays. + Otherwise, they will remain as tensors. Default: ``False``. + tolist_first (bool): if ``True``, the tensordict will be converted to a list first when + it has batch dimensions. Default: ``False``. + + Returns: + A dictionary representation of the tensordict. + + .. seealso:: :meth:`~tensordict.TensorDictBase.tolist` + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> + >>> td = TensorDict( + ... a=torch.arange(24).view(2, 3, 4), + ... b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)), + ... batch_size=(2, 3) + ... ) + >>> print(td.to_dict()) + {'a': tensor([[[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]], + + [[12, 13, 14, 15], + [16, 17, 18, 19], + [20, 21, 22, 23]]]), 'b': {'c': tensor([[[ 0, 1], + [ 2, 3], + [ 4, 5]], + + [[ 6, 7], + [ 8, 9], + [10, 11]]])}} + >>> print(td.to_dict(convert_tensors=True)) + {'a': [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]], 'b': {'c': [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]}} + + """ + result = {} + for key, value in self.items(): + if _is_tensor_collection(type(value)): + if ( + not retain_none + and _is_non_tensor(type(value)) + and value.data is None + ): + continue + if tolist_first: + value = value.tolist(convert_tensors=convert_tensors) + else: + value = value.to_dict( + retain_none=retain_none, convert_tensors=convert_tensors + ) + elif convert_tensors: + if isinstance(value, torch.Tensor) and convert_tensors == "numpy": + value = value.numpy() + elif hasattr(value, "tolist"): + value = value.tolist() + result[key] = value + return result + + def tolist( + self, + *, + convert_nodes: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ) -> List[Any]: + """Returns a nested list representation of the tensordict. + + If the tensordict has no batch dimensions, this method returns a single list or dictionary. + Otherwise, it returns a nested list where each inner list represents a batch dimension. + + Args: + convert_nodes (bool): if ``True``, leaf nodes will be converted to dictionaries. + Otherwise, they will be returned as lists of values. Default: ``True``. + convert_tensors (bool, "numpy"): if ``True``, tensors will be converted to lists when creating the dictionary. + If "numpy", tensors will be converted to numpy arrays. + Otherwise, they will remain as tensors. Default: ``False``. + tolist_first (bool): if ``True``, the tensordict will be converted to a list first when + it has batch dimensions. Default: ``False``. + as_linked_list (bool): if ``True``, the list will be converted to a :class:`tensordict.utils.LinkedList` + which will automatically update the tensordict when the list is modified. Default: ``False``. + + Returns: + A nested list representation of the tensordict. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> + >>> td = TensorDict( + ... a=torch.arange(24).view(2, 3, 4), + ... b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)), + ... batch_size=(2, 3) + ... ) + >>> print(td.tolist(tolist_first=True)) + [[{'a': tensor([0, 1, 2, 3]), 'b': [{'c': tensor(0)}, {'c': tensor(1)}]}, {'a': tensor([4, 5, 6, 7]), 'b': [{'c': tensor(2)}, {'c': tensor(3)}]}, {'a': tensor([ 8, 9, 10, 11]), 'b': [{'c': tensor(4)}, {'c': tensor(5)}]}], [{'a': tensor([12, 13, 14, 15]), 'b': [{'c': tensor(6)}, {'c': tensor(7)}]}, {'a': tensor([16, 17, 18, 19]), 'b': [{'c': tensor(8)}, {'c': tensor(9)}]}, {'a': tensor([20, 21, 22, 23]), 'b': [{'c': tensor(10)}, {'c': tensor(11)}]}]] + >>> print(td.tolist(tolist_first=False)) + [[{'a': tensor([0, 1, 2, 3]), 'b': {'c': tensor([0, 1])}}, {'a': tensor([4, 5, 6, 7]), 'b': {'c': tensor([2, 3])}}, {'a': tensor([ 8, 9, 10, 11]), 'b': {'c': tensor([4, 5])}}], [{'a': tensor([12, 13, 14, 15]), 'b': {'c': tensor([6, 7])}}, {'a': tensor([16, 17, 18, 19]), 'b': {'c': tensor([8, 9])}}, {'a': tensor([20, 21, 22, 23]), 'b': {'c': tensor([10, 11])}}]] + >>> print(td.tolist(convert_tensors=False)) + [[{'a': [0, 1, 2, 3], 'b': [{'c': 0}, {'c': 1}]}, {'a': [4, 5, 6, 7], 'b': [{'c': 2}, {'c': 3}]}, {'a': [8, 9, 10, 11], 'b': [{'c': 4}, {'c': 5}]}], [{'a': [12, 13, 14, 15], 'b': [{'c': 6}, {'c': 7}]}, {'a': [16, 17, 18, 19], 'b': [{'c': 8}, {'c': 9}]}, {'a': [20, 21, 22, 23], 'b': [{'c': 10}, {'c': 11}]}]] + >>> print(td.tolist(convert_nodes=False)) + [[[tensor([0, 1, 2, 3]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)], [tensor([4, 5, 6, 7]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)], [tensor([ 8, 9, 10, 11]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)]], [[tensor([12, 13, 14, 15]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)], [tensor([16, 17, 18, 19]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)], [tensor([20, 21, 22, 23]), TensorDict( + fields={ + c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)]]] + + """ + if convert_tensors and not convert_nodes: + raise TypeError("convert_tensors requires convert_nodes to be set to True") + if not self.batch_dims: + if convert_nodes: + return self.to_dict( + convert_tensors=convert_tensors, tolist_first=tolist_first + ) + return self + + q = collections.deque() + result = [] + q.append((self, result)) + while len(q): + val, _result = q.popleft() + vals = val.unbind(0) + if val.ndim == 1: + if convert_nodes: + vals = [ + v.to_dict( + convert_tensors=convert_tensors, tolist_first=tolist_first + ) + for v in vals + ] + else: + vals = list(vals) + _result.extend(vals) + else: + for local_val in vals: + local_res = [] + _result.append(local_res) + q.append((local_val, local_res)) + if as_linked_list: + return LinkedList(result, td=self) + return result + + def numpy(self) -> np.ndarray | dict[str, Any]: + """Converts a tensordict to a (possibly nested) dictionary of numpy arrays. + + Non-tensor data is exposed as such. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}}) + >>> print(data) + TensorDict( + fields={ + a: TensorDict( + fields={ + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + c: NonTensorData(data=a string!, batch_size=torch.Size([]), device=None)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(data.numpy()) + {'a': {'b': array(0., dtype=float32), 'c': 'a string!'}} + + seealso: :meth:`~tensordict.TensorDictBase.to_struct_array` to convert to a struct array. + + """ + as_dict = self.to_dict(retain_none=False) + + def to_numpy(x): + if isinstance(x, torch.Tensor): + if x.is_nested: + return tuple(_x.numpy() for _x in x) + return x.numpy() + if hasattr(x, "numpy"): + return x.numpy() + return x + + return torch.utils._pytree.tree_map(to_numpy, as_dict) + + def to_namedtuple(self, dest_cls: type | None = None) -> Any: + """Converts a tensordict to a namedtuple. + + Args: + dest_cls (Type, optional): an optional namedtuple class to use. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> data = TensorDict({ + ... "a_tensor": torch.zeros((3)), + ... "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3]) + >>> data.to_namedtuple() + GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!')) + + """ + + def dict_to_namedtuple(dictionary): + for key, value in dictionary.items(): + if isinstance(value, dict): + dictionary[key] = dict_to_namedtuple(value) + cls = ( + collections.namedtuple("GenericDict", dictionary.keys()) + if dest_cls is None + else dest_cls + ) + return cls(**dictionary) + + return dict_to_namedtuple(self.to_dict(retain_none=False)) + + @classmethod + def from_any( + cls, + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ): + """Recursively converts any object to a TensorDict. + + .. note:: ``from_any`` is less restrictive than the regular TensorDict constructor. It can cast data structures like + dataclasses or tuples to a tensordict using custom heuristics. This approach may incur some extra overhead and + involves more opinionated choices in terms of mapping strategies. + + .. note:: This method recursively converts the input object to a TensorDict. If the object is already a + TensorDict (or any similar tensor collection object), it will be returned as is. + + Args: + obj: The object to be converted. + + Keyword Args: + auto_batch_size (bool, optional): if ``True``, the batch size will be computed automatically. + Defaults to ``False``. + batch_dims (int, optional): If auto_batch_size is ``True``, defines how many dimensions the output tensordict + should have. Defaults to ``None`` (full batch-size at each level). + device (torch.device, optional): The device on which the TensorDict will be created. + batch_size (torch.Size, optional): The batch size of the TensorDict. + Exclusive with ``auto_batch_size``. + + Returns: + A TensorDict representation of the input object. + + Supported objects: + + - Dataclasses through :meth:`~.from_dataclass` (dataclasses will be converted to TensorDict instances, not tensorclasses). + - Namedtuples through :meth:`~.from_namedtuple`. + - Dictionaries through :meth:`~.from_dict`. + - Tuples through :meth:`~.from_tuple`. + - NumPy's structured arrays through :meth:`~.from_struct_array`. + - HDF5 objects through :meth:`~.from_h5`. + + """ + if is_tensor_collection(obj): + # Conversions from non-tensor data must be done manually + # if is_non_tensor(obj): + # from tensordict.tensorclass import LazyStackedTensorDict + # if isinstance(obj, LazyStackedTensorDict): + # return obj + # return cls.from_any(obj.data, auto_batch_size=auto_batch_size) + return obj + if isinstance(obj, dict): + return cls.from_dict( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + if isinstance(obj, UserDict): + return cls.from_dict( + dict(obj), + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + if ( + isinstance(obj, np.ndarray) + and hasattr(obj.dtype, "names") + and obj.dtype.names is not None + ): + return cls.from_struct_array( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + if isinstance(obj, tuple): + if is_namedtuple(obj): + return cls.from_namedtuple( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + return cls.from_tuple( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + if isinstance(obj, list): + if _is_list_tensor_compatible(obj)[0]: + return torch.tensor(obj) + else: + from tensordict.tensorclass import NonTensorStack + + return NonTensorStack.from_list(obj) + if is_dataclass(obj): + return cls.from_dataclass( + obj, + auto_batch_size=auto_batch_size, + device=device, + batch_size=batch_size, + ) + if _has_h5: + import h5py + + if isinstance(obj, h5py.File): + from tensordict.persistent import PersistentTensorDict + + obj = PersistentTensorDict(group=obj) + if auto_batch_size: + obj.auto_batch_size_() + return obj + return obj + + @classmethod + def from_tuple( + cls, + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ): + """Converts a tuple to a TensorDict. + + Args: + obj: The tuple instance to be converted. + + Keyword Args: + auto_batch_size (bool, optional): If ``True``, the batch size will be computed automatically. Defaults to ``False``. + batch_dims (int, optional): If auto_batch_size is ``True``, defines how many dimensions the output tensordict + should have. Defaults to ``None`` (full batch-size at each level). + device (torch.device, optional): The device on which the TensorDict will be created. Defaults to ``None``. + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to ``None``. + + Returns: + A TensorDict representation of the input tuple. + + Examples: + >>> my_tuple = (1, 2, 3) + >>> td = TensorDict.from_tuple(my_tuple) + >>> print(td) + TensorDict( + fields={ + 0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + 2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + from tensordict import TensorDict + + result = TensorDict( + { + str(i): cls.from_any(item, batch_size=batch_size, device=device) + for i, item in enumerate(obj) + }, + batch_size=batch_size, + device=device, + ) + if auto_batch_size: + if batch_size is not None: + raise TypeError(cls._CONFLICTING_BATCH_SIZES.format("from_tuple")) + result.auto_batch_size_(batch_dims=batch_dims) + return result + + _CONFLICTING_BATCH_SIZES = "Conflicting batch sizes in {}: batch_size and auto_batch_size cannot be both specified." + + @classmethod + def from_dataclass( + cls, + dataclass, + *, + dest_cls: Type | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + as_tensorclass: bool = False, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ): + """Converts a dataclass into a TensorDict instance. + + Args: + dataclass: The dataclass instance to be converted. + + Keyword Args: + dest_cls (tensorclass, optional): A tensorclass type to be used to map the data. If not provided, a new + class is created. Without effect if :attr:`obj` is a type or as_tensorclass is `False`. + auto_batch_size (bool, optional): If ``True``, automatically determines and applies batch size to the + resulting TensorDict. Defaults to ``False``. + batch_dims (int, optional): If ``auto_batch_size`` is ``True``, defines how many dimensions the output + tensordict should have. Defaults to ``None`` (full batch-size at each level). + as_tensorclass (bool, optional): If ``True``, delegates the conversion to the free function + :func:`~tensordict.from_dataclass` and returns a tensor-compatible class (:func:`~tensordict.tensorclass`) + or instance instead of a TensorDict. Defaults to ``False``. + device (torch.device, optional): The device on which the TensorDict will be created. + Defaults to ``None``. + batch_size (torch.Size, optional): The batch size of the TensorDict. + Defaults to ``None``. + + Returns: + A TensorDict instance derived from the provided dataclass, unless `as_tensorclass` is True, in which case a tensor-compatible class or instance is returned. + + Raises: + TypeError: If the provided input is not a dataclass instance. + + .. warning:: This method is distinct from the free function `from_dataclass` and serves a different purpose. + While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance. + + .. note:: + - This method creates a new TensorDict instance with keys corresponding to the fields of the input dataclass. + - Each key in the resulting TensorDict is initialized using the `cls.from_any` method. + - The `auto_batch_size` option allows for automatic batch size determination and application to the + resulting TensorDict. + + """ + if as_tensorclass: + from tensordict.tensorclass import from_dataclass + + return from_dataclass( + dataclass, + auto_batch_size=auto_batch_size, + dest_cls=dest_cls, + batch_dims=batch_dims, + batch_size=batch_size, + device=device, + ) + from dataclasses import fields + + from tensordict import TensorDict + + if not is_dataclass(dataclass): + raise TypeError( + f"Expected a dataclass input, got a {type(dataclass)} input instead." + ) + source = {} + for field in fields(dataclass): + source[field.name] = cls.from_any( + getattr(dataclass, field.name), device=device, batch_size=batch_size + ) + result = TensorDict(source, device=device, batch_size=batch_size) + if auto_batch_size: + if batch_size is not None: + raise TypeError(cls._CONFLICTING_BATCH_SIZES.format("from_dataclass")) + result.auto_batch_size_(batch_dims=batch_dims) + return result + + @classmethod + def from_namedtuple( + cls, + named_tuple, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ): + """Converts a namedtuple to a TensorDict recursively. + + Args: + named_tuple: The namedtuple instance to be converted. + + Keyword Args: + auto_batch_size (bool, optional): if ``True``, the batch size will be computed automatically. + Defaults to ``False``. + batch_dims (int, optional): If ``auto_batch_size`` is ``True``, defines how many dimensions the output + tensordict should have. Defaults to ``None`` (full batch-size at each level). + device (torch.device, optional): The device on which the TensorDict will be created. + Defaults to ``None``. + batch_size (torch.Size, optional): The batch size of the TensorDict. + Defaults to ``None``. + + Returns: + A TensorDict representation of the input namedtuple. + + Examples: + >>> from tensordict import TensorDict + >>> import torch + >>> data = TensorDict({ + ... "a_tensor": torch.zeros((3)), + ... "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3]) + >>> nt = data.to_namedtuple() + >>> print(nt) + GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!')) + >>> TensorDict.from_namedtuple(nt, auto_batch_size=True) + TensorDict( + fields={ + a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + nested: TensorDict( + fields={ + a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None), + a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + """ + from tensordict import TensorDict + + def namedtuple_to_dict(namedtuple_obj): + if is_namedtuple(namedtuple_obj): + namedtuple_obj = namedtuple_obj._asdict() + + else: + from torch.return_types import cummax, cummin, max, min + + if isinstance(namedtuple_obj, (min, cummin, max, cummax)): + namedtuple_obj = { + "values": namedtuple_obj.values, + "indices": namedtuple_obj.indices, + } + for key, value in namedtuple_obj.items(): + namedtuple_obj[key] = cls.from_any( + value, device=device, batch_size=batch_size + ) + return dict(namedtuple_obj) + + result = TensorDict( + namedtuple_to_dict(named_tuple), device=device, batch_size=batch_size + ) + if auto_batch_size: + if batch_size is not None: + raise TypeError(cls._CONFLICTING_BATCH_SIZES.format("from_namedtuple")) + result.auto_batch_size_(batch_dims=batch_dims) + return result + + @classmethod + def from_struct_array( + cls, + struct_array: np.ndarray, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ) -> Self: + """Converts a structured numpy array to a TensorDict. + + The resulting TensorDict will share the same memory content as the numpy array (it is a zero-copy operation). + Changing values of the structured numpy array in-place will affect the content of the TensorDict. + + .. note:: This method performs a zero-copy operation, meaning that the resulting TensorDict will share the same memory + content as the input numpy array. Therefore, changing values of the numpy array in-place will affect the content + of the TensorDict. + + Args: + struct_array (np.ndarray): The structured numpy array to be converted. + + Keyword Args: + auto_batch_size (bool, optional): If ``True``, the batch size will be computed automatically. Defaults to ``False``. + batch_dims (int, optional): If ``auto_batch_size`` is ``True``, defines how many dimensions the output + tensordict should have. Defaults to ``None`` (full batch-size at each level). + device (torch.device, optional): The device on which the TensorDict will be created. + Defaults to ``None``. + + .. note:: Changing the device (i.e., specifying any device other than ``None`` or ``"cpu"``) will transfer the data, + resulting in a change to the memory location of the returned data. + + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to None. + + Returns: + A TensorDict representation of the input structured numpy array. + + Examples: + >>> x = np.array( + ... [("Rex", 9, 81.0), ("Fido", 3, 27.0)], + ... dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")], + ... ) + >>> td = TensorDict.from_struct_array(x) + >>> x_recon = td.to_struct_array() + >>> assert (x_recon == x).all() + >>> assert x_recon.shape == x.shape + >>> # Try modifying x age field and check effect on td + >>> x["age"] += 1 + >>> assert (td["age"] == np.array([10, 4])).all() + + """ + if cls is TensorDictBase: + from tensordict._td import TensorDict + + cls = TensorDict + td: Self = cls( + {name: struct_array[name] for name in struct_array.dtype.names}, + batch_size=struct_array.shape if batch_size is None else batch_size, + device=device, + ) + if auto_batch_size: + if batch_size is not None: + raise TypeError( + cls._CONFLICTING_BATCH_SIZES.format("from_struct_array") + ) + td.auto_batch_size_(batch_dims=batch_dims) + return td + + def to_struct_array(self) -> np.ndarray: + """Converts a tensordict to a numpy structured array. + + In a :meth:`.from_struct_array` - :meth:`.to_struct_array` loop, the content of the input and output arrays should match. + However, `to_struct_array` will not keep the memory content of the original arrays. + + .. seealso:: :meth:`.from_struct_array` for more information. + + .. seealso:: :meth:`.numpy` to convert to a dictionary of numpy arrays. + + Returns: + A numpy structured array representation of the input TensorDict. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> td = TensorDict({'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4.0, 5.0, 6.0])}, batch_size=[3]) + >>> arr = td.to_struct_array() + >>> print(arr) + [(1, 4.) (2, 5.) (3, 6.)] + + """ + from .utils import TORCH_TO_NUMPY_DTYPE_DICT + + keys, vals = zip(*self.items()) + _vals = [] + for v in vals: + if is_tensor_collection(v): + if is_non_tensor(v): + from tensordict import NonTensorDataBase + + _vals.append( + v.data if isinstance(v, NonTensorDataBase) else v.tolist() + ) + continue + _vals.append(v.to_struct_array()) + continue + _vals.append(v) + vals = _vals + del _vals + vals = tuple(v if not is_non_tensor(v) else v.data for v in vals) + + # Convert values to numpy arrays and handle string inputs + processed_vals = [] + for v in vals: + if isinstance(v, torch.Tensor): + processed_vals.append(v.detach().cpu().numpy()) + elif isinstance(v, (list, tuple, str)): + # Handle lists/tuples which may contain strings, or strings + processed_vals.append(np.array(v)) + else: + # Keep other types as-is (already numpy arrays, etc.) + processed_vals.append(v) + vals = processed_vals + + def _get_dtype(val): + if isinstance(val, np.ndarray): + if val.dtype.kind in ["U", "S"]: # Unicode or byte strings + # Calculate appropriate string length + if val.size > 0: + max_len = max(len(str(item)) for item in val.flat) + return ( + f"U{max(10, max_len)}" # At least U10, but longer if needed + ) + return "U10" + elif val.ndim > self.ndim: + # For arrays with more dimensions than batch dims, we need to specify shape + extra_shape = val.shape[self.ndim :] + return (val.dtype, extra_shape) + return val.dtype + elif isinstance(val, torch.Tensor): + return TORCH_TO_NUMPY_DTYPE_DICT.get(val.dtype, val.dtype) + else: + return "U10" + + dtype = [(key, _get_dtype(val)) for key, val in zip(keys, vals)] + + if self.ndim: + # For multi-dimensional tensordicts, we need to create structured arrays properly + # Reshape each value to have batch dimensions first, then flatten the batch dimensions + batch_shape = self.shape + batch_size = int(np.prod(batch_shape)) + + # Reshape and prepare data for structured array + reshaped_vals = [] + for val in vals: + if isinstance(val, np.ndarray): + # Ensure the array has the right batch shape + if val.shape[: self.ndim] == batch_shape: + # Flatten batch dimensions + new_shape = (batch_size,) + val.shape[self.ndim :] + reshaped_vals.append(val.reshape(new_shape)) + else: + # If shapes don't match, try to broadcast + try: + reshaped_vals.append( + np.broadcast_to( + val, batch_shape + val.shape[self.ndim :] + ).reshape((batch_size,) + val.shape[self.ndim :]) + ) + except ValueError: + reshaped_vals.append(val) + else: + reshaped_vals.append(val) + + # Create structured array + result = np.empty(batch_size, dtype=dtype) + for key, val in zip(keys, reshaped_vals): + if isinstance(val, np.ndarray) and val.shape[0] == batch_size: + result[key] = val + else: + result[key] = val + + # Reshape back to original batch shape + return result.reshape(batch_shape) + + # For scalar case, create structured array properly + result = np.empty((), dtype=dtype) + for key, val in zip(keys, vals): + if isinstance(val, np.ndarray) and val.ndim == 0: + result[key] = val.item() + elif isinstance(val, (np.ndarray, torch.Tensor)) and val.size == 1: + result[key] = val.item() + else: + result[key] = val + return result + + def to_h5( + self, + filename, + **kwargs, + ) -> Any: + """Converts a tensordict to a PersistentTensorDict with the h5 backend. + + Args: + filename (str or path): path to the h5 file. + **kwargs: kwargs to be passed to :meth:`h5py.File.create_dataset`. + + Returns: + A :class:`~.tensordict.PersitentTensorDict` instance linked to the newly created file. + + Examples: + >>> import tempfile + >>> import timeit + >>> + >>> from tensordict import TensorDict, MemoryMappedTensor + >>> td = TensorDict({ + ... "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)), + ... "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))}, + ... }, [1_000_000]) + >>> + >>> file = tempfile.NamedTemporaryFile() + >>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9) + >>> print(td_h5) + PersistentTensorDict( + fields={ + a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False), + b: PersistentTensorDict( + fields={ + c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([1000000]), + device=None, + is_shared=False)}, + batch_size=torch.Size([1000000]), + device=None, + is_shared=False) + + + """ + from tensordict.persistent import PersistentTensorDict + + out = PersistentTensorDict.from_dict( + self, + filename=filename, + **kwargs, + ) + if self._has_names(): + out.names = self.names + return out + + def empty( + self, recurse=False, *, batch_size=None, device=NO_DEFAULT, names=None + ) -> Self: # noqa: D417 + """Returns a new, empty tensordict with the same device and batch size. + + Args: + recurse (bool, optional): if ``True``, the entire structure of the + ``TensorDict`` will be reproduced without content. + Otherwise, only the root will be duplicated. + Defaults to ``False``. + + Keyword Args: + batch_size (torch.Size, optional): a new batch-size for the tensordict. + device (torch.device, optional): a new device. + names (list of str, optional): dimension names. + + """ + if not recurse: + result = self._select(set_shared=False) + else: + # simply exclude the leaves + result = self._exclude(*self.keys(True, True), set_shared=False) + if batch_size is not None: + result.batch_size = batch_size + if device is not NO_DEFAULT: + if device is None: + result.clear_device_() + else: + result = result.to(device) + if names is not None: + result.names = names + return result + + # Filling + def zero_(self) -> Self: + """Zeros all tensors in the tensordict in-place.""" + + def fn(item): + item.zero_() + + self._fast_apply(fn=fn, call_on_nested=True, propagate_lock=True) + return self + + def fill_(self, key: NestedKey, value: float | bool) -> Self: + """Fills a tensor pointed by the key with a given scalar value. + + Args: + key (str or nested key): entry to be filled. + value (Number or bool): value to use for the filling. + + Returns: + self + + """ + key = _unravel_key_to_tuple(key) + data = self._get_tuple(key, NO_DEFAULT) + if _is_tensor_collection(type(data)): + + def fill(x): + return x.fill_(value) + + data._fast_apply(fill, inplace=True) + else: + data = data.fill_(value) + self._set_tuple(key, data, inplace=True, validated=True, non_blocking=False) + return self + + # Masking + @abc.abstractmethod + def masked_fill_(self, mask: Tensor, value: float | bool) -> Self: + """Fills the values corresponding to the mask with the desired value. + + Args: + mask (boolean torch.Tensor): mask of values to be filled. Shape + must match the tensordict batch-size. + value: value to used to fill the tensors. + + Returns: + self + + Examples: + >>> td = TensorDict(source={'a': torch.zeros(3, 4)}, + ... batch_size=[3]) + >>> mask = torch.tensor([True, False, False]) + >>> td.masked_fill_(mask, 1.0) + >>> td.get("a") + tensor([[1., 1., 1., 1.], + [0., 0., 0., 0.], + [0., 0., 0., 0.]]) + """ + raise NotImplementedError + + @abc.abstractmethod + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: + """Out-of-place version of masked_fill. + + Args: + mask (boolean torch.Tensor): mask of values to be filled. Shape + must match the tensordict batch-size. + value: value to used to fill the tensors. + + Returns: + self + + Examples: + >>> td = TensorDict(source={'a': torch.zeros(3, 4)}, + ... batch_size=[3]) + >>> mask = torch.tensor([True, False, False]) + >>> td1 = td.masked_fill(mask, 1.0) + >>> td1.get("a") + tensor([[1., 1., 1., 1.], + [0., 0., 0., 0.], + [0., 0., 0., 0.]]) + """ + raise NotImplementedError + + def where( + self, + condition: Tensor, + other: Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ) -> Self: # noqa: D417 + """Return a ``TensorDict`` of elements selected from either self or other, depending on condition. + + Args: + condition (BoolTensor): When ``True`` (nonzero), yields ``self``, + otherwise yields ``other``. + other (TensorDictBase or Scalar): value (if ``other`` is a scalar) + or values selected at indices where condition is ``False``. + + Keyword Args: + out (TensorDictBase, optional): the output ``TensorDictBase`` instance. + pad (scalar, optional): if provided, missing keys from the source + or destination tensordict will be written as `torch.where(mask, self, pad)` + or `torch.where(mask, pad, other)`. Defaults to ``None``, ie + missing keys are not tolerated. + update_batch_size (bool, optional): if ``True`` and ``out`` is provided, the batch size of the output will be + updated to match the batch size of the condition. Defaults to ``False``. + + """ + ... + + @abc.abstractmethod + def masked_select(self, mask: Tensor) -> Self: + """Masks all tensors of the TensorDict and return a new TensorDict instance with similar keys pointing to masked values. + + Args: + mask (torch.Tensor): boolean mask to be used for the tensors. + Shape must match the TensorDict ``batch_size``. + + Examples: + >>> td = TensorDict(source={'a': torch.zeros(3, 4)}, + ... batch_size=[3]) + >>> mask = torch.tensor([True, False, False]) + >>> td_mask = td.masked_select(mask) + >>> td_mask.get("a") + tensor([[0., 0., 0., 0.]]) + + """ + raise NotImplementedError + + @abc.abstractmethod + def _change_batch_size(self, new_size: torch.Size) -> None: + raise NotImplementedError + + @abc.abstractmethod + def is_contiguous(self) -> bool: + """Returns a boolean indicating if all the tensors are contiguous.""" + raise NotImplementedError + + @abc.abstractmethod + def contiguous(self) -> Self: + """Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).""" + raise NotImplementedError + + @cache # noqa: B019 + @_as_context_manager() + def flatten_keys( + self, + separator: str = ".", + inplace: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + ) -> Self: + """Converts a nested tensordict into a flat one, recursively. + + The TensorDict type will be lost and the result will be a simple TensorDict instance. + + Args: + separator (str, optional): the separator between the nested items. + inplace (bool, optional): if ``True``, the resulting tensordict will + have the same identity as the one where the call has been made. + Defaults to ``False``. + is_leaf (callable, optional): a callable over a class type returning + a bool indicating if this class has to be considered as a leaf. + + .. note:: The purpose of `is_leaf` is not to prevent recursive calls into nested tensordicts, but + rather to mark certain types as "leaves" for the purpose of filtering when `leaves_only=True`. + Even if `is_leaf(cls)` returns `True`, the nested structure of the tensordict will still be + traversed if `include_nested=True`. + In other words, `is_leaf` does not control the recursion depth, but rather provides a way to filter + out certain types from the result when `leaves_only=True`. This means that a node in the tree can + be both a leaf and a node with children. + In practice, the default value of ``is_leaf`` does exclude tensordict and tensorclass instances + from the leaf set. + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`. + + Examples: + >>> data = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3}, batch_size=[]) + >>> data.flatten_keys(separator=" - ") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b - c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + e - f - g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + This method and :meth:`~.unflatten_keys` are particularly useful when + handling state-dicts, as they make it possible to seamlessly convert + flat dictionaries into data structures that mimic the structure of the + model. + + Examples: + >>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) + >>> ddp_model = torch.ao.quantization.QuantWrapper(model) + >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") + >>> print(state_dict) + TensorDict( + fields={ + module: TensorDict( + fields={ + 0: TensorDict( + fields={ + bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> model_state_dict = state_dict.get("module") + >>> print(model_state_dict) + TensorDict( + fields={ + 0: TensorDict( + fields={ + bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> model.load_state_dict(dict(model_state_dict.flatten_keys("."))) + """ + if inplace: + return self._flatten_keys_inplace(separator=separator, is_leaf=is_leaf) + return self._flatten_keys_outplace(separator=separator, is_leaf=is_leaf) + + def _flatten_keys_outplace(self, separator, is_leaf): + if is_leaf is None: + is_leaf = _is_leaf_nontensor + all_leaves_all_vals = zip( + *self.items(include_nested=True, leaves_only=True, is_leaf=is_leaf) + ) + try: + all_leaves, all_vals = all_leaves_all_vals + except ValueError: + return self.empty() + all_leaves_flat = [ + key if isinstance(key, str) else separator.join(key) for key in all_leaves + ] + + if len(set(all_leaves_flat)) < len(all_leaves_flat): + # find duplicates + seen = set() + conflicts = [] + for leaf, leaf_flat in zip(all_leaves, all_leaves_flat): + if leaf_flat in seen: + conflicts.append(leaf) + else: + seen.add(leaf_flat) + raise KeyError( + f"Flattening keys in tensordict causes keys {conflicts} to collide." + ) + result = self.empty() + _set_dict = getattr(result, "_set_dict", None) + if _set_dict is not None: + _set_dict( + dict(zip(all_leaves_flat, all_vals)), + validated=True, + ) + else: + for val, leaf_flat in zip(all_vals, all_leaves_flat): + result._set_str( + leaf_flat, + val, + validated=True, + inplace=False, + non_blocking=False, + ) + # Uncomment if you want key operations to propagate the shared status + # self._maybe_set_shared_attributes(result) + # if result._is_shared or result._is_memmap: + # result.lock_() + return result + + def _flatten_keys_inplace(self, separator, is_leaf): + if is_leaf is None: + is_leaf = _is_leaf_nontensor + all_leaves = [ + _unravel_key_to_tuple(key) + for key in self.keys(include_nested=True, leaves_only=True, is_leaf=is_leaf) + ] + all_leaves_flat = [separator.join(key) for key in all_leaves] + if len(set(all_leaves_flat)) < len(set(all_leaves)): + # find duplicates + seen = set() + conflicts = [] + for leaf, leaf_flat in zip(all_leaves, all_leaves_flat): + if leaf_flat in seen: + conflicts.append(leaf) + else: + seen.add(leaf_flat) + raise KeyError( + f"Flattening keys in tensordict causes keys {conflicts} to collide." + ) + # we will need to remove the empty tensordicts later on + root_keys = set(self.keys()) + for leaf, leaf_flat in zip(all_leaves, all_leaves_flat): + self.rename_key_(leaf, leaf_flat) + if isinstance(leaf, str): + root_keys.discard(leaf) + self.exclude(*root_keys, inplace=True) + return self + + @cache # noqa: B019 + @_as_context_manager() + def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> Self: + """Converts a flat tensordict into a nested one, recursively. + + The TensorDict type will be lost and the result will be a simple TensorDict instance. + The metadata of the nested tensordicts will be inferred from the root: + all instances across the data tree will share the same batch-size, + dimension names and device. + + Args: + separator (str, optional): the separator between the nested items. + inplace (bool, optional): if ``True``, the resulting tensordict will + have the same identity as the one where the call has been made. + Defaults to ``False``. + + Examples: + >>> data = TensorDict({"a": 1, "b - c": 2, "e - f - g": 3}, batch_size=[]) + >>> data.unflatten_keys(separator=" - ") + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False), + e: TensorDict( + fields={ + f: TensorDict( + fields={ + g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + This method and :meth:`~.unflatten_keys` are particularly useful when + handling state-dicts, as they make it possible to seamlessly convert + flat dictionaries into data structures that mimic the structure of the + model. + + Examples: + >>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) + >>> ddp_model = torch.ao.quantization.QuantWrapper(model) + >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") + >>> print(state_dict) + TensorDict( + fields={ + module: TensorDict( + fields={ + 0: TensorDict( + fields={ + bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> model_state_dict = state_dict.get("module") + >>> print(model_state_dict) + TensorDict( + fields={ + 0: TensorDict( + fields={ + bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> model.load_state_dict(dict(model_state_dict.flatten_keys("."))) + + """ + if not inplace: + result = self._clone(recurse=False).unflatten_keys( + separator=separator, inplace=True + ) + if result._is_shared or result._is_memmap: + result.lock_() + return result + else: + if not is_compiling(): + key_list = list(self.keys()) + else: + key_list = [k for k in self.keys()] # noqa + + for key in key_list: + if separator in key: + new_key = tuple(key.split(separator)) + try: + self.rename_key_(key, new_key, safe=True) + except KeyError: + raise KeyError( + f"Unflattening key(s) in tensordict will override an existing for unflattened key {new_key}." + ) + return self + + def split_keys( + self, + *key_sets, + inplace=False, + default: Any = NO_DEFAULT, + strict: bool = True, + reproduce_struct: bool = False, + ) -> Tuple[T, ...]: + """Splits the tensordict in subsets given one or more set of keys. + + The method will return ``N+1`` tensordicts, where ``N`` is the number of + the arguments provided. + + Args: + key_sets (sequence of Dict[in_key, out_key] or list of keys): the various splits. + inplace (bool, optional): if ``True``, the keys are removed from ``self`` + in-place. Defaults to ``False``. + default (Any, optional): the value to be returned when a key is missing. + If not specified and ``strict=True``, an exception is raised. + strict (bool, optional): if ``True``, an exception is raised when a key + is missing. Defaults to ``True``. + reproduce_struct (bool, optional): if ``True``, all tensordict returned have + the same tree structure as ``self``, even if some sub-tensordicts + contain no leaves. + + .. note:: + ``None`` non-tensor values will be ignored and not returned. + + .. note:: + The method does not check for duplicates in the provided lists. + + Examples: + >>> td = TensorDict( + ... a=0, + ... b=0, + ... c=0, + ... d=0, + ... ) + >>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"]) + >>> print(td_bc) + """ + from tensordict import PersistentTensorDict + + if isinstance(self, PersistentTensorDict): + last_out = self.to_tensordict() + else: + last_out = self.copy() + if strict: + default = NO_DEFAULT + elif default is NO_DEFAULT: + default = None + outs = [] + if inplace: + keys_to_del = set() + for key_set in key_sets: + outs.append(self.empty(recurse=reproduce_struct)) + if not isinstance(key_set, dict): + key_set = {key: key for key in key_set} + for key in key_set: + val = last_out.pop(key, default) + if val is not None: + outs[-1].set(key_set[key], val) + if inplace: + keys_to_del.add(key) + if inplace: + # We update self here because doing it in the loop would + # possibly break people's code when doing a try/except KeyError + # around this method + for key in keys_to_del: + try: + self.pop(key, default=default) + except KeyError: + # We're good if strict is False + if strict: + raise + last_out = self + if not reproduce_struct: + last_out.filter_empty_() + outs.append(last_out) + return tuple(outs) + + def separates( + self, + *keys: NestedKey, + default: Any = NO_DEFAULT, + strict: bool = True, + filter_empty: bool = True, + ) -> Self: + """Separates the specified keys from the tensordict in-place. + + .. seealso:: This method is equivalent to calling :meth:`~tensordict.TensorDictBase.split_keys` with + ``inplace=True`` on a single split. + + .. seealso:: This method is equivalent to calling :meth:`~tensordict.TensorDictBase.exclude` except that it + returns the other split of the data. + + Args: + keys (NestedKey): the keys to separate from the tensordict. + default (Any, optional): the value to be returned when a key is missing. + If not specified and ``strict=True``, an exception is raised. Otherwise, the default of any missing key + will be ``None`` unless specified otherwise. + strict (bool, optional): if ``True``, an exception is raised when a key + is missing. Defaults to ``True``. + filter_empty (bool, optional): if ``True``, empty tensordicts within ``self`` will be removed. + Defaults to ``True``. + + Returns: + T: the separated tensordict. + + Examples: + >>> td = TensorDict( + ... a=0, + ... b=0, + ... c=0, + ... d=0, + ... ) + >>> td_a_c = td.separates("a", "c") + >>> print(td_a_c) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(td) + TensorDict( + fields={ + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + from tensordict import PersistentTensorDict + + if isinstance(self, PersistentTensorDict): + last_out = self.to_tensordict() + else: + last_out = self + strict = strict and default is NO_DEFAULT + if strict: + default = NO_DEFAULT + else: + default = None + key_set = keys + + # We want to keep the metadata such as batch-size etc so we call with recurse=True + out = self.empty(recurse=True) + for key in key_set: + val = last_out.pop(key, default) + out.set(key, val) + out.filter_empty_() + if filter_empty: + self.filter_empty_() + return out + + @abc.abstractmethod + def _index_tensordict( + self, + index: IndexType, + new_batch_size: torch.Size | None = None, + names: List[str] | None = None, + ) -> Self: + raise NotImplementedError + + # Locking functionality + @property + def is_locked(self) -> bool: + return self._is_locked + + @is_locked.setter + def is_locked(self, value: bool) -> None: + if value: + self.lock_() + else: + self.unlock_() + + def _propagate_lock(self, lock_parents_weakrefs=None, *, is_compiling): + """Registers the parent tensordict that handles the lock.""" + self._is_locked = True + if lock_parents_weakrefs is not None: + lock_parents_weakrefs = [ + ref + for ref in lock_parents_weakrefs + if not any(refref is ref for refref in self._lock_parents_weakrefs) + ] + if not is_compiling: + is_root = lock_parents_weakrefs is None + if is_root: + lock_parents_weakrefs = [] + else: + self._lock_parents_weakrefs = ( + self._lock_parents_weakrefs + lock_parents_weakrefs + ) + lock_parents_weakrefs = list(lock_parents_weakrefs) + lock_parents_weakrefs.append(weakref.ref(self)) + + for value in self.values(): + if _is_tensor_collection(type(value)): + value._propagate_lock(lock_parents_weakrefs, is_compiling=is_compiling) + + @property + def _lock_parents_weakrefs(self): + _lock_parents_weakrefs = self.__dict__.get("__lock_parents_weakrefs") + if _lock_parents_weakrefs is None: + self.__dict__["__lock_parents_weakrefs"] = [] + _lock_parents_weakrefs = self.__dict__["__lock_parents_weakrefs"] + return _lock_parents_weakrefs + + @_lock_parents_weakrefs.setter + def _lock_parents_weakrefs(self, value: list): + self.__dict__["__lock_parents_weakrefs"] = value + + @_as_context_manager("is_locked") + def lock_(self) -> Self: + """Locks a tensordict for non in-place operations. + + Functions such as :meth:`~.set`, :meth:`~.__setitem__`, :meth:`~.update`, + :meth:`~.rename_key_` or other operations that add or remove entries + will be blocked. + + This method can be used as a decorator. + + Example: + >>> from tensordict import TensorDict + >>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[]) + >>> with td.lock_(): + ... assert td.is_locked + ... try: + ... td.set("d", 0) # error! + ... except RuntimeError: + ... print("td is locked!") + ... try: + ... del td["d"] + ... except RuntimeError: + ... print("td is locked!") + ... try: + ... td.rename_key_("a", "d") + ... except RuntimeError: + ... print("td is locked!") + ... td.set("a", 0, inplace=True) # No storage is added, moved or removed + ... td.set_("a", 0) # No storage is added, moved or removed + ... td.update({"a": 0}, inplace=True) # No storage is added, moved or removed + ... td.update_({"a": 0}) # No storage is added, moved or removed + >>> assert not td.is_locked + """ + if self.is_locked: + return self + is_comp = is_compiling() + if is_comp: + _lock_warn() + self._propagate_lock(is_compiling=is_comp) + return self + + @erase_cache + def _propagate_unlock(self): + # if we end up here, we can clear the graph associated with this td + self._is_locked = False + + self._is_shared = False + self._is_memmap = False + + # Remove consolidated metadata when unlocking to prevent silent errors + if hasattr(self, "_consolidated"): + delattr(self, "_consolidated") + + sub_tds = [] + for value in self.values(): + if _is_tensor_collection(type(value)): + sub_tds.extend(value._propagate_unlock()) + sub_tds.append(value) + return sub_tds + + def _check_unlock(self, first_attempt=True): + if not first_attempt: + gc.collect() + obj = None + for ref in self._lock_parents_weakrefs: + obj = ref() + # check if the locked parent exists and if it's locked + # we check _is_locked because it can be False or None in the case of Lazy stacks, + # but if we check obj.is_locked it will be True for this class. + if obj is not None and obj._is_locked: + break + + else: + try: + self._lock_parents_weakrefs = [] + except AttributeError: + # Some tds (eg, LazyStack) have an automated way of creating the _lock_parents_weakref + pass + return + + if first_attempt: + del obj + return self._check_unlock(False) + raise RuntimeError( + "Cannot unlock a tensordict that is part of a locked graph. " + "Unlock the root tensordict first. If the tensordict is part of multiple graphs, " + "group the graphs under a common tensordict an unlock this root. " + f"self: {self}, obj: {obj}" + ) + + @_as_context_manager("is_locked") + def unlock_(self) -> Self: + """Unlocks a tensordict for non in-place operations. + + Can be used as a decorator. + + See :meth:`~.lock_` for more details. + """ + try: + sub_tds = self._propagate_unlock() + for sub_td in sub_tds: + sub_td._check_unlock() + + self._check_unlock() + except RuntimeError as err: + self.lock_() + raise err + return self + + # Conversion (device or dtype) + @overload + def to( + self: T, + device: int | device | None = ..., + dtype: torch.dtype | None = ..., + non_blocking: bool = ..., + inplace: bool = False, + ) -> Self: ... + + @overload + def to(self: T, dtype: torch.dtype, non_blocking: bool = ...) -> Self: ... + + @overload + def to(self: T, tensor: Tensor, non_blocking: bool = ...) -> Self: ... + + @overload + def to(self: T, *, other: T, non_blocking: bool = ...) -> Self: ... + + @overload + def to(self: T, *, batch_size: torch.Size) -> Self: ... + + def _to_cuda_with_pin_mem( + self, + *, + num_threads, + device="cuda", + non_blocking=None, + to: Callable, + inplace: bool = False, + ): + if self.is_empty(): + return self.to(device, inplace=inplace) + keys, vals = self._items_list( + leaves_only=True, include_nested=True, is_leaf=_NESTED_TENSORS_AS_LISTS + ) + lkeys = len(keys) + q_in = queue.SimpleQueue() + q_out = queue.SimpleQueue() + threads = [] + items = {} + for key, val in _zip_strict(keys, vals): + q_in.put_nowait((key, val)) + for _ in range(min(num_threads, lkeys)): + thread = Thread(target=_pin_mem, args=(q_in, q_out)) + thread.start() + threads.append(thread) + try: + while len(items) < lkeys: + keyval = q_out.get(timeout=_PIN_MEM_TIMEOUT) + if not isinstance(keyval, tuple): + raise keyval + key, val = keyval + items[key] = to(val) + finally: + for thread in threads: + thread.join(timeout=_PIN_MEM_TIMEOUT) + + def get(name, val): + return items.get(name, val) + + result = self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + device=device, + out=self if inplace else None, + checked=True, + ) + return result + + @_as_context_manager() + def to(self, *args, **kwargs) -> Self: + """Maps a TensorDictBase subclass either on another device, dtype or to another TensorDictBase subclass (if permitted). + + Casting tensors to a new dtype is not allowed, as tensordicts are not bound to contain a single + tensor dtype. + + Args: + device (torch.device, optional): the desired device of the tensordict. + dtype (torch.dtype, optional): the desired floating point or complex dtype of + the tensordict. + tensor (torch.Tensor, optional): Tensor whose dtype and device are the desired + dtype and device for all tensors in this TensorDict. + + Keyword Args: + non_blocking (bool, optional): whether the operations should be blocking. + memory_format (torch.memory_format, optional): the desired memory + format for 4D parameters and buffers in this tensordict. + batch_size (torch.Size, optional): resulting batch-size of the + output tensordict. + other (TensorDictBase, optional): TensorDict instance whose dtype + and device are the desired dtype and device for all tensors + in this TensorDict. + + .. note:: + Since :class:`~tensordict.TensorDictBase` instances do not have + a dtype, the dtype is gathered from the example leaves. + If there are more than one dtype, then no dtype + casting is undertook. + + non_blocking_pin (bool, optional): if ``True``, the tensors are pinned before + being sent to device. This will be done asynchronously but can be + controlled via the ``num_threads`` argument. + + .. note:: + Calling ``tensordict.pin_memory().to("cuda")`` will usually + be much slower than ``tensordict.to("cuda", non_blocking_pin=True)`` as + the pin_memory is called asynchronously in the second case. + Multithreaded ``pin_memory`` will usually be beneficial if the tensors + are large and numerous: when there are too few tensors to be sent, + the overhead of spawning threads and collecting data outweighs the benefits + of multithreading, and if the tensors are small the overhead of iterating + over a long list is also prohibitively large. + + num_threads (int or None, optional): if ``non_blocking_pin=True``, the number + of threads to be used for ``pin_memory``. By default, + ``max(1, torch.get_num_threads())`` threads will be spawn. + ``num_threads=0`` will cancel any + multithreading for the `pin_memory()` calls. + inplace (bool, optional): if ``True``, the data will be written in-place in the same tensordict. + This can be significantly faster whenever building a tensordict is CPU-overhead bound. + Defaults to ``False``. + + Returns: + a new tensordict instance if the device differs from the tensordict + device and/or if the dtype is passed. The same tensordict otherwise. + ``batch_size`` only modifications are done in-place. + + .. note:: + If the TensorDict is consolidated, the resulting TensorDict will be consolidated too. + Each new tensor will be a view on the consolidated storage cast to the desired device. + + This operation can be used as a context manager too. When used as a context manager, + the tensordict is temporarily moved to the target device/dtype, and upon exiting + the context, it is automatically restored to its original device/dtype. This is + particularly useful when working with neural network modules that expect data + on a specific device. + + Examples: + >>> data = TensorDict({"a": 1.0}, [], device=None) + >>> data_cuda = data.to("cuda:0") # casts to cuda + >>> data_int = data.to(torch.int) # casts to int + >>> data_cuda_int = data.to("cuda:0", torch.int) # multiple casting + >>> data_cuda = data.to(torch.randn(3, device="cuda:0")) # using an example tensor + >>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0")) # using a tensordict example + + Using as a context manager for temporary device changes: + >>> from tensordict.nn import TensorDictModule + >>> import torch + >>> + >>> # Create a module and data + >>> mod = TensorDictModule(lambda x: x + 1, in_keys=["x"], out_keys=["y"]) + >>> td = TensorDict(x=torch.zeros(3), batch_size=[3], device="cpu") + >>> + >>> # Use context manager to temporarily move to GPU + >>> with td.to("cuda") as td_gpu: + ... td_gpu.update(mod(td_gpu)) # Process on GPU and update in-place + >>> + >>> # Data is automatically restored to original device + >>> assert td["x"].device.type == "cpu" + >>> assert td["y"].device.type == "cpu" # Output also restored to original device + """ + non_blocking = kwargs.pop("non_blocking", None) + + ( + device, + dtype, + _, + convert_to_format, + batch_size, + non_blocking_pin, + num_threads, + inplace, + ) = _parse_to(*args, **kwargs) + result = self + + if device is not None and dtype is None and device == self.device: + return result + + if self.is_consolidated() and dtype is None: + return self._to_consolidated( + device=device, + pin_memory=non_blocking_pin, + num_threads=num_threads, + non_blocking=non_blocking, + inplace=inplace, + ) + + if non_blocking is None: + sub_non_blocking = True + non_blocking = False + else: + sub_non_blocking = non_blocking + + if convert_to_format is not None: + + def to(tensor): + return tensor.to( + device, + dtype, + non_blocking=sub_non_blocking, + convert_to_format=convert_to_format, + ) + + else: + + def to(tensor): + return tensor.to( + device=device, dtype=dtype, non_blocking=sub_non_blocking + ) + + apply_kwargs = {} + if device is not None or dtype is not None: + if non_blocking_pin and num_threads != 0: + if num_threads is None: + num_threads = max(1, torch.get_num_threads() // 2) + result = self._to_cuda_with_pin_mem( + num_threads=num_threads, to=to, device=device, inplace=inplace + ) + else: + apply_kwargs["device"] = device if device is not None else self.device + apply_kwargs["batch_size"] = batch_size + apply_kwargs["out"] = self if inplace else None + apply_kwargs["checked"] = True + if non_blocking_pin: + + def to_pinmem(tensor, _to=to): + return to(tensor.pin_memory()) + + result = result._fast_apply( + to_pinmem, propagate_lock=True, **apply_kwargs + ) + else: + # result = result._fast_apply(to, propagate_lock=True, **apply_kwargs) + keys, tensors = self._items_list(True, True) + tensors = [to(t) for t in tensors] + items = dict(zip(keys, tensors)) + + def get(name, val): + return items.get(name, val) + + result = self._fast_apply( + get, + named=True, + nested_keys=True, + is_leaf=_NESTED_TENSORS_AS_LISTS, + propagate_lock=True, + **apply_kwargs, + ) + + if batch_size is not None: + result.batch_size = batch_size + if ( + device is not None + and sub_non_blocking + and not non_blocking + and device.type != "cuda" + ): + self._sync_all() + return result + + def _to_consolidated( + self, *, device, pin_memory, num_threads, non_blocking, inplace + ): + if num_threads is None: + # unspecified num_threads should mean 0 + num_threads = 0 + storage = self._consolidated["storage"] + if pin_memory: + storage = storage.pin_memory() + storage_cast = storage.to(device, non_blocking=True) + untyped_storage = storage_cast.untyped_storage() + + def set_(x): + if x.is_nested: + from torch._subclasses.fake_tensor import FakeTensor + from torch._subclasses.functional_tensor import FunctionalTensor + from torch.nested._internal.nested_tensor import ( + _tensor_symint_registry, + NestedTensor, + ) + from torch.nested._internal.ops import extract_kwargs + + if x.layout != torch.jagged: + raise RuntimeError( + "to(device) with nested tensors that do not have a jagged layout is not implemented yet. " + "Please raise an issue on GitHub." + ) + kwargs = extract_kwargs(x) + values = x._values + lengths = x._lengths + offsets = x._offsets + kwargs["offsets"] = set_(offsets) + if lengths is not None: + kwargs["lengths"] = set_(lengths) + ragged_source = lengths + else: + ragged_source = offsets + new_thing = kwargs.get("lengths", kwargs.get("offsets")) + if isinstance(new_thing, (FakeTensor, FunctionalTensor)): + from torch._subclasses.functional_tensor import ( + mb_unwrap_functional_tensor, + ) + + # Temporary hack until we have the union find + tgt = mb_unwrap_functional_tensor(new_thing) + src = mb_unwrap_functional_tensor(ragged_source) + tgt.nested_int_memo = src.nested_int_memo + elif new_thing is not None: + _tensor_symint_registry[new_thing] = _tensor_symint_registry[ + ragged_source + ] + + return NestedTensor( + set_(values), + **kwargs, + ) + storage_offset = x.storage_offset() + stride = x.stride() + return x.new_empty(0, device=device).set_( + untyped_storage, + size=x.shape, + stride=stride, + storage_offset=storage_offset, + ) + + if inplace: + out = self + else: + out = None + + result = self._fast_apply( + set_, + device=torch.device(device), + num_threads=num_threads, + out=out, + checked=True, + ) + result._consolidated = {"storage": storage_cast} + if "metadata" in self._consolidated: + # faster than deepcopy + def copy_dict(d): + return { + k: v if not isinstance(v, dict) else copy_dict(v) + for k, v in d.items() + } + + result._consolidated["metadata"] = copy_dict(self._consolidated["metadata"]) + # Ensure the result remains locked to maintain consolidated state integrity + result.lock_() + if non_blocking in (False, None): + if device.type != "cpu" and non_blocking is False: + # sending to non-cpu device force sync + non_cpu_device = device + elif storage.device.type != "cpu": + # sending from non-cpu device: need sync unless intentionally not asked for + non_cpu_device = storage.device.type + else: + non_cpu_device = None + if non_cpu_device is not None: + device_type = _get_available_device_type() + device_module = _get_device_module(device_type) + if hasattr(device_module, "current_stream"): + device_module.current_stream(non_cpu_device).synchronize() + else: + # Some device modules, such as torch.mps, don't have current_stream attr + device_module.synchronize() + + return result + + @property + def _has_cuda(self): + val = self.__dict__.get("_has_cuda_val") + if val is None: + # Cache this value + val = torch.cuda.is_available() + self.__dict__["_has_cuda_val"] = val + return val + + @property + def _has_mps(self): + val = self.__dict__.get("_has_mps_val") + if val is None: + # Cache this value + val = torch.backends.mps.is_available() + self.__dict__["_has_mps_val"] = val + return val + + def _sync_all(self): + device_type = _get_available_device_type() + if device_type is None: + return + + if device_type == "cuda": + # TODO: dynamo doesn't like torch.cuda.is_initialized + if not is_compiling() and torch.cuda.is_initialized(): + torch.cuda.synchronize() + else: + device_module = _get_device_module(device_type) + device_module.synchronize() + + def is_floating_point(self) -> bool: + """Checks if all tensors in the tensordict are floating point.""" + for item in self.values(include_nested=True, leaves_only=True): + if not item.is_floating_point(): + return False + else: + return True + + def double(self) -> Self: + r"""Casts all tensors to ``torch.bool``.""" + + def dble(x): + return x.double() + + return self._fast_apply(dble, propagate_lock=True) + + def float(self) -> Self: + r"""Casts all tensors to ``torch.float``.""" + + def tofloat(x): + return x.float() + + return self._fast_apply(tofloat, propagate_lock=True) + + def int(self) -> Self: + r"""Casts all tensors to ``torch.int``.""" + + def toint(x): + return x.int() + + return self._fast_apply(toint, propagate_lock=True) + + def bool(self) -> Self: + r"""Casts all tensors to ``torch.bool``.""" + + def tobool(x): + return x.bool() + + return self._fast_apply(tobool, propagate_lock=True) + + def half(self) -> Self: + r"""Casts all tensors to ``torch.half``.""" + + def tohalf(x): + return x.half() + + return self._fast_apply(tohalf, propagate_lock=True) + + def type(self, dst_type: torch.dtype) -> Self: + r"""Casts all tensors to :attr:`dst_type`. + + Args: + dst_type (type or string): the desired type + + """ + + def totype(x): + return x.type(dst_type) + + return self._fast_apply(totype) + + # Gradient compatibility + @property + def requires_grad(self) -> bool: + return any(v.requires_grad for v in self.values()) + + def requires_grad_(self, requires_grad=True) -> Self: + """Change if autograd should record operations on this tensor: sets this tensor's requires_grad attribute in-place. + + Returns this tensordict. + + Args: + requires_grad (bool, optional): whether or not autograd should record operations on this tensordict. + Defaults to ``True``. + + """ + for val in self._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS): + val.requires_grad_(requires_grad) + return self + + @abc.abstractmethod + def detach_(self) -> Self: + """Detach the tensors in the tensordict in-place. + + Returns: + self. + + """ + raise NotImplementedError + + @cache # noqa: B019 + def detach(self) -> Self: + """Detach the tensors in the tensordict. + + Returns: + a new tensordict with no tensor requiring gradient. + + """ + + def detach(x): + return x.detach() + + return self._fast_apply( + detach, + propagate_lock=True, + ) + + @_make_dtype_promotion + def bfloat16(self) -> Self: ... + + @_make_dtype_promotion + def complex128(self) -> Self: ... + + @_make_dtype_promotion + def complex32(self) -> Self: ... + + @_make_dtype_promotion + def complex64(self) -> Self: ... + + @_make_dtype_promotion + def float16(self) -> Self: ... + + @_make_dtype_promotion + def float32(self) -> Self: ... + + @_make_dtype_promotion + def float64(self) -> Self: ... + + @_make_dtype_promotion + def int16(self) -> Self: ... + + @_make_dtype_promotion + def int32(self) -> Self: ... + + @_make_dtype_promotion + def int64(self) -> Self: ... + + @_make_dtype_promotion + def int8(self) -> Self: ... + + @_make_dtype_promotion + def qint32(self) -> Self: ... + + @_make_dtype_promotion + def qint8(self) -> Self: ... + + @_make_dtype_promotion + def quint4x2(self) -> Self: ... + + @_make_dtype_promotion + def quint8(self) -> Self: ... + + @_make_dtype_promotion + def uint16(self) -> Self: ... + + @_make_dtype_promotion + def uint32(self) -> Self: ... + + @_make_dtype_promotion + def uint64(self) -> Self: ... + + @_make_dtype_promotion + def uint8(self) -> Self: ... + + +_ACCEPTED_CLASSES = ( + Tensor, + TensorDictBase, +) + + +def _register_tensor_class(cls): + global _ACCEPTED_CLASSES + _ACCEPTED_CLASSES = set(_ACCEPTED_CLASSES) + _ACCEPTED_CLASSES.add(cls) + _ACCEPTED_CLASSES = tuple(_ACCEPTED_CLASSES) + + +_TENSOR_COLLECTION_MEMO = {} + + +def _is_tensor_collection(datatype: type) -> bool: + is_dynamo = is_compiling() + out = None + if not is_dynamo: + out = _TENSOR_COLLECTION_MEMO.get(datatype) + + if out is None: + out = issubclass(datatype, TensorDictBase) or _is_tensorclass(datatype) + if not is_dynamo: + _TENSOR_COLLECTION_MEMO[datatype] = out + return out + + +def is_tensor_collection(datatype: type | Any) -> bool: + """Checks if a data object or a type is a tensor container from the tensordict lib. + + Returns: + ``True`` if the input is a TensorDictBase subclass, a tensorclass or an istance of these. + ``False`` otherwise. + + Examples: + >>> is_tensor_collection(TensorDictBase) # True + >>> is_tensor_collection(TensorDict()) # True + >>> @tensorclass + ... class MyClass: + ... pass + ... + >>> is_tensor_collection(MyClass) # True + >>> is_tensor_collection(MyClass(batch_size=[])) # True + + """ + # memoizing is 2x faster + if not isinstance(datatype, type): + datatype = type(datatype) + return _is_tensor_collection(datatype) + + +def _default_is_leaf(cls: Type) -> bool: + """Returns ``True`` if a type is not a tensor collection (tensordict or tensorclass). + + Examples: + >>> from tensordict import TensorDict, default_is_leaf + >>> import torch + >>> td = TensorDict(a={}, b="a string!", c=torch.randn(())) + >>> print(td.keys(leaves_only=True, is_leaf=default_is_leaf)) + _TensorDictKeysView(['c'], + include_nested=False, + leaves_only=True) + + .. seealso:: :meth:`~tensordict.is_leaf_nontensor`. + """ + return not _is_tensor_collection(cls) + + +def _is_leaf_nontensor(cls: Type) -> bool: + """Returns ``True`` if a type is not a tensor collection (tensordict or tensorclass) or is a non-tensor. + + Examples: + >>> from tensordict import TensorDict, default_is_leaf + >>> import torch + >>> td = TensorDict(a={}, b="a string!", c=torch.randn(())) + >>> print(td.keys(leaves_only=True, is_leaf=default_is_leaf)) + _TensorDictKeysView(['b', 'c'], + include_nested=False, + leaves_only=True) + + .. seealso:: :meth:`~tensordict.default_is_leaf`. + """ + if _is_tensor_collection(cls): + return _pass_through_cls(cls) + return issubclass(cls, torch.Tensor) + + +def _load_metadata(prefix: Path): + filepath = prefix / "meta.json" + with open(filepath, "rb") as json_metadata: + metadata = json.loads(json_metadata.read()) + return metadata + + +class _NestedTensorsAsLists: + """Class used to iterate over leaves of lazily stacked tensordicts.""" + + def __new__(cls): + if not hasattr(cls, "instance"): + cls.instance = super(cls, cls).__new__(cls) + return cls.instance + + def __bool__(self): + return False + + def __call__(self, val): + return _default_is_leaf(val) + + +class _NestedTensorsAsListsNonTensor: + def __new__(cls): + if not hasattr(cls, "instance"): + cls.instance = super(cls, cls).__new__(cls) + return cls.instance + + def __bool__(self): + return False + + def __call__(self, val): + return _is_leaf_nontensor(val) + + +_NESTED_TENSORS_AS_LISTS = _NestedTensorsAsLists() + + +_NESTED_TENSORS_AS_LISTS_NONTENSOR = _NestedTensorsAsListsNonTensor() + + +def _expand_to_match_shape( + parent_batch_size: torch.Size, + data: Tensor | TensorDictBase, + self_batch_dims: int, + self_device: DeviceType, + index: Any = None, +) -> Tensor | TensorDictBase: + """Creates and empty tensor / tensordict that can host values. + + Given a tensordict with shape ``parent_batch_size``, this function creates an expanded version + of ``data`` such that ``data_expand[index].shape == data.shape``. + + """ + if not parent_batch_size and self_batch_dims == 1: + # This is what happens when indexing an empty tensor with a bool: + # torch.zeros(())[True].shape == torch.Size((1,)) + return data.new_zeros(data.shape[1:]) + if not _is_tensor_collection(type(data)): + result = torch.zeros( + ( + *parent_batch_size, + *_shape(data)[self_batch_dims:], + ), + dtype=data.dtype, + device=self_device, + ) + else: + # tensordict + batch_size = torch.Size([*parent_batch_size, *_shape(data)[self_batch_dims:]]) + result = data.empty(batch_size=batch_size) + return result + + +def from_any( + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +): + """Converts any object to a TensorDict. + + .. seealso:: :meth:`~tensordict.TensorDictBase.from_any` for more information. + """ + return TensorDictBase.from_any( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + + +def from_tuple( + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +) -> "TensorDictBase": + """Converts a tuple to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_tuple` for more information. + """ + return TensorDictBase.from_tuple( + obj, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + + +def from_namedtuple( + named_tuple, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +) -> "TensorDictBase": + """Converts a namedtuple to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_namedtuple` for more information. + """ + from tensordict import TensorDict + + return TensorDict.from_namedtuple( + named_tuple, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + + +def from_struct_array( + struct_array, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +) -> "TensorDictBase": + """Converts a structured numpy array to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_struct_array` for more information. + + Examples: + >>> x = np.array( + ... [("Rex", 9, 81.0), ("Fido", 3, 27.0)], + ... dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")], + ... ) + >>> td = from_struct_array(x) + >>> x_recon = td.to_struct_array() + >>> assert (x_recon == x).all() + >>> assert x_recon.shape == x.shape + >>> # Try modifying x age field and check effect on td + >>> x["age"] += 1 + >>> assert (td["age"] == np.array([10, 4])).all() + + """ + return TensorDictBase.from_struct_array( + struct_array, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + + +def from_list( + input: list[TensorCollection | Mapping], + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + cls: Type | None = None, + lazy_stack: bool = None, +) -> TensorCollection: + """Converts a list of dictionaries or TensorDicts to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_dict` for more information. + """ + if cls is not None: + cls = TensorDictBase + return cls.from_list( + input, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + type=type, + lazy_stack=lazy_stack, + ) + + +def from_dict( + d, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +) -> "TensorDictBase": + """Converts a dictionary to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_dict` for more information. + + + Examples: + >>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)} + >>> print(from_dict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # nested dict: the nested TensorDict can have a different batch-size + >>> # as long as its leading dims match. + >>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}} + >>> print(from_dict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # we can also use this to work out the batch sie of a tensordict + >>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, []) + >>> print( + from_dict(input_td)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + """ + from tensordict import TensorDict + + return TensorDict.from_dict( + d, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) + + +def from_h5( + h5_file, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, +) -> "TensorDictBase": + """Converts an HDF5 file to a TensorDict. + + .. seealso:: :meth:`TensorDictBase.from_h5` for more information. + """ + from tensordict import TensorDict + + return TensorDict.from_h5( + h5_file, + auto_batch_size=auto_batch_size, + batch_dims=batch_dims, + device=device, + batch_size=batch_size, + ) diff --git a/lib/python3.12/site-packages/tensordict/functional.py b/lib/python3.12/site-packages/tensordict/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..7d8a4bd14aa001dc8823cb6374a7203aec4e2852 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/functional.py @@ -0,0 +1,502 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from typing import Any, Callable, Dict, Mapping, Sequence + +import torch + +from tensordict._lazy import LazyStackedTensorDict +from tensordict._td import TensorDict +from tensordict.base import ( + _is_tensor_collection, + CompatibleType, + NestedKey, + T, + TensorDictBase, +) +from tensordict.utils import ( + _check_keys, + _shape, + DeviceType, + is_non_tensor, + is_tensorclass, + unravel_key, +) + + +def pad(tensordict: T, pad_size: Sequence[int], value: float = 0.0) -> T: + """Pads all tensors in a tensordict along the batch dimensions with a constant value, returning a new tensordict. + + Args: + tensordict (TensorDict): The tensordict to pad + pad_size (Sequence[int]): The padding size by which to pad some batch + dimensions of the tensordict, starting from the first dimension and + moving forward. [len(pad_size) / 2] dimensions of the batch size will + be padded. For example to pad only the first dimension, pad has the form + (padding_left, padding_right). To pad two dimensions, + (padding_left, padding_right, padding_top, padding_bottom) and so on. + pad_size must be even and less than or equal to twice the number of batch dimensions. + value (float, optional): The fill value to pad by, default 0.0 + + Returns: + A new TensorDict padded along the batch dimensions + + Examples: + >>> from tensordict import TensorDict, pad + >>> import torch + >>> td = TensorDict({'a': torch.ones(3, 4, 1), + ... 'b': torch.ones(3, 4, 1, 1)}, batch_size=[3, 4]) + >>> dim0_left, dim0_right, dim1_left, dim1_right = [0, 1, 0, 2] + >>> padded_td = pad(td, [dim0_left, dim0_right, dim1_left, dim1_right], value=0.0) + >>> print(padded_td.batch_size) + torch.Size([4, 6]) + >>> print(padded_td.get("a").shape) + torch.Size([4, 6, 1]) + >>> print(padded_td.get("b").shape) + torch.Size([4, 6, 1, 1]) + + """ + if len(pad_size) > 2 * len(tensordict.batch_size): + raise RuntimeError( + "The length of pad_size must be <= 2 * the number of batch dimensions" + ) + + if len(pad_size) % 2: + raise RuntimeError("pad_size must have an even number of dimensions") + + new_batch_size = list(tensordict.batch_size) + for i in range(len(pad_size)): + new_batch_size[i // 2] += pad_size[i] + + reverse_pad = list(pad_size[::-1]) + for i in range(0, len(reverse_pad), 2): + reverse_pad[i], reverse_pad[i + 1] = reverse_pad[i + 1], reverse_pad[i] + + out = TensorDict._new_unsafe( + {}, + torch.Size(new_batch_size), + device=tensordict.device, + ) + for key, tensor in tensordict.items(): + cur_pad = reverse_pad + if len(pad_size) < len(_shape(tensor)) * 2: + cur_pad = [0] * (len(_shape(tensor)) * 2 - len(pad_size)) + reverse_pad + + if _is_tensor_collection(type(tensor)): + padded = pad(tensor, pad_size, value) + else: + padded = torch.nn.functional.pad(tensor, cur_pad, value=value) + out.set(key, padded) + + return out + + +def pad_sequence( + list_of_tensordicts: Sequence[T], + pad_dim: int = 0, + padding_value: float = 0.0, + out: T | None = None, + return_mask: bool | NestedKey = False, +) -> T: + """Pads a list of tensordicts in order for them to be stacked together in a contiguous format. + + Args: + list_of_tensordicts (List[TensorDictBase]): the list of instances to pad and stack. + pad_dim (int, optional): the ``pad_dim`` indicates the dimension to pad all the keys in the tensordict. + Defaults to ``0``. + padding_value (number, optional): the padding value. Defaults to ``0.0``. + out (TensorDictBase, optional): if provided, the destination where the data will be + written. + return_mask (bool or NestedKey, optional): if ``True``, a "masks" entry will be returned. If ``return_mask`` is a nested key (string or tuple of strings), it will be return the masks and be used as the key for the masks entry. + It contains a tensordict with the same structure as the stacked tensordict where every entry contains the mask of valid values with size ``torch.Size([stack_len, *new_shape])``, + where `new_shape[pad_dim] = max_seq_length` and the rest of the `new_shape` matches the previous shape of the contained tensors. + + Examples: + >>> list_td = [ + ... TensorDict({"a": torch.zeros((3, 8)), "b": torch.zeros((6, 8))}, batch_size=[]), + ... TensorDict({"a": torch.zeros((5, 8)), "b": torch.zeros((6, 8))}, batch_size=[]), + ... ] + >>> padded_td = pad_sequence(list_td, return_mask=True) + >>> print(padded_td) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([2, 5, 8]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([2, 6, 8]), device=cpu, dtype=torch.float32, is_shared=False), + masks: TensorDict( + fields={ + a: Tensor(shape=torch.Size([2, 5]), device=cpu, dtype=torch.bool, is_shared=False), + b: Tensor(shape=torch.Size([2, 6]), device=cpu, dtype=torch.bool, is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + """ + if not len(list_of_tensordicts): + raise RuntimeError("list_of_tensordicts cannot be empty") + + if return_mask and is_tensorclass(list_of_tensordicts[0]): + raise RuntimeError( + "Expected 'return_mask=False' when list_of_tensordicts contains " + "tensorclasses, but got 'return_mask=True'. If you want masks, " + "plase convert the tensorclasses to TensorDicts first." + ) + + if not isinstance(return_mask, bool): + masks_key = unravel_key(return_mask) + return_mask = True + else: + masks_key = "masks" + + # check that all tensordict match + update_batch_size = True + max_seq_length = float("-inf") + keys = _check_keys(list_of_tensordicts, leaves_only=True, include_nested=True) + list_of_dicts = [{} for _ in range(len(list_of_tensordicts))] + keys_copy = list(keys) + mask_keys = [] + for i, td in enumerate(list_of_tensordicts): + if is_tensorclass(td): + td = td._tensordict + + for key in keys: + item = td.get(key) + list_of_dicts[i][key] = item + if is_non_tensor(item): + continue + tensor_shape = item.shape + + if len(tensor_shape) == 0: + raise RuntimeError("Cannot pad scalars") + + pos_pad_dim = pad_dim if pad_dim >= 0 else len(tensor_shape) + pad_dim + + # track the maximum sequence length to update batch_size accordingly + if tensor_shape[pos_pad_dim] > max_seq_length: + max_seq_length = tensor_shape[pos_pad_dim] + + # The mask should always contain the batch_size of the TensorDict + mask_shape = td.shape + + # if the pad_dim is past the batch_size of the TensorDict, we need to add the new dimension to the mask + if pos_pad_dim >= td.ndim: + mask_shape += torch.Size([tensor_shape[pos_pad_dim]]) + update_batch_size = False + + if return_mask: + mask_key = unravel_key((masks_key, key)) + mask_keys.append(mask_key) + list_of_dicts[i][mask_key] = torch.ones(mask_shape, dtype=torch.bool) + keys_copy.append(mask_key) + + keys = keys_copy + + old_batch_size = list(list_of_tensordicts[0].batch_size) + if update_batch_size and len(old_batch_size) > 0: + old_batch_size[pad_dim] = max_seq_length + shape = [ + len(list_of_tensordicts), + ] + old_batch_size + + if out is None: + out = list_of_tensordicts[0].empty(recurse=True).reshape(torch.Size(shape)) + + for key in keys: + try: + item0 = list_of_dicts[0][key] + if is_non_tensor(item0): + out.set(key, TensorDict.lazy_stack([d[key] for d in list_of_dicts])) + continue + tensor_shape = item0.shape + pos_pad_dim = ( + (pad_dim if pad_dim >= 0 else len(tensor_shape) + pad_dim) + if len(tensor_shape) > 1 + else 0 # handles the case when the masks are 1-dimensional + ) + out.set( + key, + torch.nn.utils.rnn.pad_sequence( + [d[key].transpose(0, pos_pad_dim) for d in list_of_dicts], + batch_first=True, + padding_value=padding_value if key not in mask_keys else False, + ).transpose(1, pos_pad_dim + 1), + inplace=True, + ) + except Exception as err: + raise RuntimeError(f"pad_sequence failed for key {key}") from err + return out + + +def merge_tensordicts( + *tensordicts: T, + callback_exist: ( + Callable[[Any], Any] | Dict[NestedKey, Callable[[Any], Any]] | None + ) = None, +) -> T: + """Merges tensordicts together. + + Args: + *tensordicts (sequence of TensorDict or equivalent): the list of tensordicts to merge together. + + Keyword Args: + callback_exist (callable or Dict[str, callable], optional): a callable in case an entry exists in each and every tensordict. + If the entry is present in some but not all tensordicts, or if ``callback_exist`` is not passed, + `update` is used and the first non-``None`` value in the tensordict sequence will be used. + If a dictionary of callables is passed, it will contain the associated callback function for some of the + nested keys in the tensordicts passed to the function. + + Examples: + >>> from tensordict import merge_tensordicts, TensorDict + >>> td0 = TensorDict({"a": {"b0": 0}, "c": {"d": {"e": 0}}, "common": 0}) + >>> td1 = TensorDict({"a": {"b1": 1}, "f": {"g": {"h": 1}}, "common": 1}) + >>> td2 = TensorDict({"a": {"b2": 2}, "f": {"g": {"h": 2}}, "common": 2}) + >>> td = merge_tensordicts(td0, td1, td2, callback_exist=lambda *v: torch.stack(list(v))) + >>> print(td) + TensorDict( + fields={ + a: TensorDict( + fields={ + b0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False), + c: TensorDict( + fields={ + d: TensorDict( + fields={ + e: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False), + common: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False), + f: TensorDict( + fields={ + g: TensorDict( + fields={ + h: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(td["common"]) + tensor([0, 1, 2]) + + """ + if len(tensordicts) < 2: + raise RuntimeError( + f"at least 2 tensordicts must be provided, got" f" {len(tensordicts)}" + ) + + out = tensordicts[0].empty(recurse=True) + key_list = set() + + def func(name, *vals): + nonlocal key_list + if name in key_list: + return + key_list.add(name) + cb = ( + callback_exist + if not isinstance(callback_exist, Mapping) + else callback_exist.get(name) + ) + if cb is not None and all(val is not None for val in vals): + out.set(name, cb(*vals)) + return + for val in vals: + if val is not None: + out.set(name, val) + return + + for i in range(len(tensordicts)): + if i > 0: + tds = tensordicts[i + 1 :] + tensordicts[:i] + else: + tds = tensordicts[1:] + tensordicts[i]._fast_apply( + func, *tds, named=True, nested_keys=True, filter_empty=True, default=None + ) + return out + + +def dense_stack_tds( + td_list: Sequence[TensorDictBase] | LazyStackedTensorDict, + dim: int | None = None, +) -> T: + """Densely stack a list of :class:`~tensordict.TensorDictBase` objects (or a :class:`~tensordict.LazyStackedTensorDict`) given that they have the same structure. + + This function is called with a list of :class:`~tensordict.TensorDictBase` (either passed directly or obtrained from + a :class:`~tensordict.LazyStackedTensorDict`). + Instead of calling ``torch.stack(td_list)``, which would return a :class:`~tensordict.LazyStackedTensorDict`, + this function expands the first element of the input list and stacks the input list onto that element. + This works only when all the elements of the input list have the same structure. + The :class:`~tensordict.TensorDictBase` returned will have the same type of the elements of the input list. + + This function is useful when some of the :class:`~tensordict.TensorDictBase` objects that need to be stacked + are :class:`~tensordict.LazyStackedTensorDict` or have :class:`~tensordict.LazyStackedTensorDict` + among entries (or nested entries). + In those cases, calling ``torch.stack(td_list).to_tensordict()`` is infeasible. + Thus, this function provides an alternative for densely stacking the list provided. + + Args: + td_list (List of TensorDictBase or LazyStackedTensorDict): the tds to stack. + dim (int, optional): the dimension to stack them. + If td_list is a LazyStackedTensorDict, it will be retrieved automatically. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict import dense_stack_tds + >>> from tensordict.tensordict import assert_allclose_td + >>> td0 = TensorDict({"a": torch.zeros(3)},[]) + >>> td1 = TensorDict({"a": torch.zeros(4), "b": torch.zeros(2)},[]) + >>> td_lazy = torch.stack([td0, td1], dim=0) + >>> td_container = TensorDict({"lazy": td_lazy}, []) + >>> td_container_clone = td_container.clone() + >>> td_stack = torch.stack([td_container, td_container_clone], dim=0) + >>> td_stack + LazyStackedTensorDict( + fields={ + lazy: LazyStackedTensorDict( + fields={ + a: Tensor(shape=torch.Size([2, 2, -1]), device=cpu, dtype=torch.float32, is_shared=False)}, + exclusive_fields={ + }, + batch_size=torch.Size([2, 2]), + device=None, + is_shared=False, + stack_dim=0)}, + exclusive_fields={ + }, + batch_size=torch.Size([2]), + device=None, + is_shared=False, + stack_dim=0) + >>> td_stack = dense_stack_tds(td_stack) # Automatically use the LazyStackedTensorDict stack_dim + TensorDict( + fields={ + lazy: LazyStackedTensorDict( + fields={ + a: Tensor(shape=torch.Size([2, 2, -1]), device=cpu, dtype=torch.float32, is_shared=False)}, + exclusive_fields={ + 1 -> + b: Tensor(shape=torch.Size([2, 2]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2, 2]), + device=None, + is_shared=False, + stack_dim=1)}, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + # Note that + # (1) td_stack is now a TensorDict + # (2) this has pushed the stack_dim of "lazy" (0 -> 1) + # (3) this has revealed the exclusive keys. + >>> assert_allclose_td(td_stack, dense_stack_tds([td_container, td_container_clone], dim=0)) + # This shows it is the same to pass a list or a LazyStackedTensorDict + + """ + if isinstance(td_list, LazyStackedTensorDict): + dim = td_list.stack_dim + td_list = td_list.tensordicts + elif isinstance(td_list, TensorDict): + # then it is already dense + return td_list + elif dim is None: + raise ValueError( + "If a list of tensordicts is provided, stack_dim must not be None" + ) + shape = list(td_list[0].shape) + shape.insert(dim, len(td_list)) + + return TensorDict.maybe_dense_stack(td_list, dim=dim) + + +def make_tensordict( + input_dict: dict[str, CompatibleType] | None = None, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: DeviceType | None = None, + auto_batch_size: bool | None = None, + **kwargs: CompatibleType, # source +) -> TensorDict: + """Returns a TensorDict created from the keyword arguments or an input dictionary. + + If ``batch_size`` is not specified, returns the maximum batch size possible. + + This function works on nested dictionaries too, or can be used to determine the + batch-size of a nested tensordict. + + Args: + input_dict (dictionary, optional): a dictionary to use as a data source + (nested keys compatible). + **kwargs (TensorDict or torch.Tensor): keyword arguments as data source + (incompatible with nested keys). + batch_size (iterable of int, optional): a batch size for the tensordict. + device (torch.device or compatible type, optional): a device for the TensorDict. + auto_batch_size (bool, optional): if ``True``, the batch size will be computed automatically. + Defaults to ``False``. + + Examples: + >>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)} + >>> print(make_tensordict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # alternatively + >>> td = make_tensordict(**input_dict) + >>> # nested dict: the nested TensorDict can have a different batch-size + >>> # as long as its leading dims match. + >>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}} + >>> print(make_tensordict(input_dict)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> # we can also use this to work out the batch sie of a tensordict + >>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, []) + >>> print(make_tensordict(input_td)) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + b: TensorDict( + fields={ + c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + """ + if input_dict is not None: + kwargs.update(input_dict) + return TensorDict.from_dict( + kwargs, batch_size=batch_size, device=device, auto_batch_size=auto_batch_size + ) diff --git a/lib/python3.12/site-packages/tensordict/memmap.py b/lib/python3.12/site-packages/tensordict/memmap.py new file mode 100644 index 0000000000000000000000000000000000000000..915ee7efa30f1bed2c4f5636a74cc10bab48df92 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/memmap.py @@ -0,0 +1,1130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import functools + +import mmap +import os +import re +import sys +import tempfile +from multiprocessing import reduction, util +from pathlib import Path +from typing import Any, Callable, overload, TYPE_CHECKING + +import numpy as np +import torch +from tensordict.utils import ( + _maybe_correct_neg_dim, + _shape, + _zip_strict, + implement_for, + IndexType, + NESTED_TENSOR_ERR, +) + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +class MemoryMappedTensor(torch.Tensor): + """A Memory-mapped Tensor. + + Supports filenames or file handlers. + + The main advantage of MemoryMappedTensor resides in its serialization methods, + which ensure that the tensor is passed through queues or RPC remote calls without + any copy. + + .. note:: + When used within RPC settings, the filepath should be accessible to both nodes. + If it isn't the behaviour of passing a MemoryMappedTensor from one worker + to another is undefined. + + MemoryMappedTensor supports multiple construction methods. + + Examples: + >>> # from an existing tensor + >>> tensor = torch.randn(3) + >>> with tempfile.NamedTemporaryFile() as file: + ... memmap_tensor = MemoryMappedTensor.from_tensor(tensor, filename=file.name) + ... assert memmap_tensor.filename is not None + >>> # if no filename is passed, a handler is used + >>> tensor = torch.randn(3) + >>> memmap_tensor = MemoryMappedTensor.from_tensor(tensor, filename=file.name) + >>> assert memmap_tensor.filename is None + >>> # one can create an empty tensor too + >>> with tempfile.NamedTemporaryFile() as file: + ... memmap_tensor_empty = MemoryMappedTensor.empty_like(tensor, filename=file.name) + >>> with tempfile.NamedTemporaryFile() as file: + ... memmap_tensor_zero = MemoryMappedTensor.zeros_like(tensor, filename=file.name) + >>> with tempfile.NamedTemporaryFile() as file: + ... memmap_tensor = MemoryMappedTensor.ones_like(tensor, filename=file.name) + """ + + _filename: str | Path = None + _handler: _FileHandler = None + _clear: bool + index: Any + parent_shape: torch.Size + + def __new__( + cls, + source, + *, + dtype=None, + shape=None, + index=None, + device=None, + handler=None, + filename=None, + ): + if device is not None and torch.device(device).type != "cpu": + raise ValueError(f"{cls} device must be cpu!") + if isinstance(source, str): + if filename is not None: + raise TypeError("Duplicated filename argument.") + filename = source + source = None + if filename is not None: + if dtype is not None: + raise TypeError("Cannot pass new dtype if source is provided.") + result = cls.from_tensor( + torch.as_tensor(source), + filename=filename, + # dtype=dtype, + shape=shape, + # index=index, + ) + if index is not None: + return result[index] + return result + elif isinstance(source, torch.StorageBase): + return cls.from_storage( + source, + dtype=dtype, + shape=shape, + index=index, + device=device, + handler=handler, + filename=filename, + ) + elif handler is not None: + return cls.from_handler( + handler, + dtype, + shape, + index, + ) + return super().__new__(cls, source) + + def __init__( + self, + source, + *, + handler=None, + dtype=None, + shape=None, + device=None, + filename=None, + ): ... + + __torch_function__ = torch._C._disabled_torch_function_impl + + @classmethod + def from_tensor( + cls, + input, + *, + filename: Path | str = None, + existsok: bool = False, + copy_existing: bool = False, + copy_data: bool = True, + shape: torch.Size | None = None, + ): # noqa: D417 + """Creates a MemoryMappedTensor with the same content as another tensor. + + If the tensor is already a MemoryMappedTensor the original tensor is + returned if the `filename` argument is `None` or if the two paths match. + In all other cases, a new :class:`MemoryMappedTensor` is produced. + + Args: + input (torch.Tensor): the tensor which content must be copied onto + the MemoryMappedTensor. + + Keyword Args: + filename (path to a file): the path to the file where the tensor + should be stored. If none is provided, a file handler is used + instead. + existsok (bool, optional): if ``True``, the file will overwrite + an existing file. Defaults to ``False``. + copy_existing (bool, optional): if ``True`` and the provided input + is a MemoryMappedTensor with an associated filename, copying + the content to the new location is permitted. Otherwise, an + exception is thrown. This behaviour exists to prevent + inadvertently duplicating data on disk. + copy_data (bool, optional): if ``True``, the content of the tensor + will be copied on the storage. Defaults to ``True``. + shape (torch.Size or torch.Tensor): a shape to override the tensor + shape. If a tensor is passed, it must represent the nested shapes of a + nested tensor. + """ + if isinstance(input, MemoryMappedTensor): + if (filename is None and input._filename is None) or ( + input._filename is not None + and filename is not None + and Path(filename).absolute() == Path(input.filename).absolute() + ): + # either location was not specified, or memmap is already in the + # correct location, so just return the MemmapTensor unmodified + return input + elif not copy_existing and ( + input._filename is not None + and filename is not None + and Path(filename).absolute() != Path(input.filename).absolute() + ): + raise RuntimeError( + f"A filename was provided but the tensor already has a file associated " + f"({input.filename}). " + f"To copy the tensor onto the new location, pass copy_existing=True." + ) + elif isinstance(input, np.ndarray): + raise TypeError( + "Convert input to torch.Tensor before calling MemoryMappedTensor.from_tensor." + ) + if input.requires_grad: + raise RuntimeError( + "MemoryMappedTensor.from_tensor is incompatible with tensor.requires_grad." + ) + if shape is None: + shape = _shape(input, nested_shape=True) + if isinstance(shape, torch.Tensor): + shape_numel = shape.prod(-1).sum() + elif isinstance(shape, torch.Size): + shape_numel = shape.numel() + else: + shape_numel = torch.Size(shape).numel() + if filename is None: + if input.dtype.is_floating_point: + size = torch.finfo(input.dtype).bits // 8 * shape_numel + elif input.dtype.is_complex: + raise ValueError( + "Complex-valued tensors are not supported by MemoryMappedTensor." + ) + elif input.dtype == torch.bool: + size = shape_numel + else: + # assume integer + size = torch.iinfo(input.dtype).bits // 8 * shape_numel + handler = _FileHandler(size) + if isinstance(shape, torch.Tensor): + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError(NESTED_TENSOR_ERR) + result = torch.frombuffer(memoryview(handler.buffer), dtype=input.dtype) + if copy_data: + result.untyped_storage().copy_(input.untyped_storage()) + result = torch._nested_view_from_buffer( + result, + shape, + *offsets_strides, + ) + else: + result = torch.frombuffer(memoryview(handler.buffer), dtype=input.dtype) + result = result.view(shape) + result = cls(result) + else: + handler = None + if not existsok and os.path.exists(str(filename)): + raise RuntimeError(f"The file {filename} already exists.") + result = torch.from_file( + str(filename), + shared=True, + dtype=input.dtype, + size=shape_numel, + # needed when device ctx differs + device=torch.device("cpu"), + ) + if isinstance(shape, torch.Tensor): + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError(NESTED_TENSOR_ERR) + if copy_data: + result.untyped_storage().copy_(input.untyped_storage()) + result = torch._nested_view_from_buffer( + result, + shape, + *offsets_strides, + ) + else: + result = result.view(shape) + result = cls(result) + result._handler = handler + result.filename = filename + result.index = None + result.parent_shape = shape + if copy_data: + if hasattr(input, "full_tensor"): + # for DTensors, cheaper than importing DTensor every time + input = input.full_tensor() + if not result.is_nested: + result.copy_(input) + return result + + @classmethod + def from_storage( + cls, + storage, + *, + shape: torch.Size | None = None, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + index: IndexType | None = None, + filename: Path | str = None, + handler: _handler = None, + ): + if getattr(storage, "filename", None) is not None: + if filename is None: + filename = storage.filename + elif str(storage.filename) != str(filename): + raise RuntimeError( + "Providing a storage with an associated filename that differs from the filename argument is not permitted unless filename=None. " + f"Got filename={str(filename)}, storage.filename={str(storage.filename)}" + ) + tensor = torch.tensor(storage, dtype=dtype, device=device) + if shape is not None: + if isinstance(shape, torch.Tensor): + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError( + "The PyTorch version isn't compatible with memmap " + "nested tensors. Please upgrade to a more recent " + "version." + ) + tensor = torch._nested_view_from_buffer( + tensor, + shape, + *offsets_strides, + ) + else: + tensor = tensor.view(shape) + + tensor = cls(tensor) + if filename is not None: + tensor.filename = filename + elif handler is not None: + tensor._handler = handler + if index is not None: + return tensor[index] + return tensor + + @property + def filename(self): + """The filename of the tensor, if it has one. + + Raises an exception otherwise. + """ + filename = self._filename + if filename is None: + raise RuntimeError("The MemoryMappedTensor has no file associated.") + return filename + + @filename.setter + def filename(self, value): + if value is None and self._filename is None: + return + value = str(Path(value).absolute()) + if self._filename is not None and value != self._filename: + raise RuntimeError( + "the MemoryMappedTensor has already a filename associated." + ) + self._filename = value + + @classmethod + def empty_like(cls, input, *, filename=None): + # noqa: D417 + """Creates a tensor with no content but the same shape and dtype as the input tensor. + + Args: + input (torch.Tensor): the tensor to use as an example. + + Keyword Args: + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + """ + return cls.from_tensor( + torch.zeros((), dtype=input.dtype, device=input.device).expand_as(input), + filename=filename, + copy_data=False, + ) + + @classmethod + def full_like(cls, input, fill_value, *, filename=None): + # noqa: D417 + """Creates a tensor with a single content indicated by the `fill_value` argument, but the same shape and dtype as the input tensor. + + Args: + input (torch.Tensor): the tensor to use as an example. + fill_value (float or equivalent): content of the tensor. + + Keyword Args: + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + """ + return cls.from_tensor( + torch.zeros((), dtype=input.dtype, device=input.device).expand_as(input), + filename=filename, + copy_data=False, + ).fill_(fill_value) + + @classmethod + def zeros_like(cls, input, *, filename=None): + # noqa: D417 + """Creates a tensor with a 0-filled content, but the same shape and dtype as the input tensor. + + Args: + input (torch.Tensor): the tensor to use as an example. + + Keyword Args: + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + """ + return cls.from_tensor( + torch.zeros((), dtype=input.dtype, device=input.device).expand_as(input), + filename=filename, + copy_data=False, + ).fill_(0.0) + + @classmethod + def ones_like(cls, input, *, filename=None): + # noqa: D417 + """Creates a tensor with a 1-filled content, but the same shape and dtype as the input tensor. + + Args: + input (torch.Tensor): the tensor to use as an example. + + Keyword Args: + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + """ + return cls.from_tensor( + torch.ones((), dtype=input.dtype, device=input.device).expand_as(input), + filename=filename, + copy_data=False, + ).fill_(1.0) + + @classmethod + @overload + def ones(cls, *size, dtype=None, device=None, filename=None): ... + + @classmethod + @overload + def ones(cls, shape, *, dtype=None, device=None, filename=None): ... + + @classmethod + def ones(cls, *args, **kwargs): + # noqa: D417 + """Creates a tensor with a 1-filled content, specific shape, dtype and filename. + + Args: + shape (integers or torch.Size): the shape of the tensor. + + Keyword Args: + dtype (torch.dtype): the dtype of the tensor. + device (torch.device): the device of the tensor. Only `None` and `"cpu"` + are accepted, any other device will raise an exception. + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + existsok (bool, optional): whether it is ok to overwrite an existing file. + Defaults to ``False``. + """ + shape, device, dtype, _, filename = _proc_args_const(*args, **kwargs) + if device is not None: + device = torch.device(device) + if device.type != "cpu": + raise RuntimeError("Only CPU tensors are supported.") + result = torch.ones((), dtype=dtype, device=device) + if isinstance(shape, torch.Tensor): + return cls.empty( + shape, device=device, dtype=dtype, filename=filename + ).fill_(1) + if shape: + if isinstance(shape[0], (list, tuple)) and len(shape) == 1: + shape = torch.Size(shape[0]) + else: + shape = torch.Size(shape) + result = result.expand(shape) + return cls.from_tensor( + result, + filename=filename, + existsok=kwargs.pop("existsok", False), + ) + + @classmethod + @overload + def zeros(cls, *size, dtype=None, device=None, filename=None): ... + + @classmethod + @overload + def zeros(cls, shape, *, dtype=None, device=None, filename=None): ... + + @classmethod + def zeros(cls, *args, **kwargs): + # noqa: D417 + """Creates a tensor with a 0-filled content, specific shape, dtype and filename. + + Args: + shape (integers or torch.Size): the shape of the tensor. + + Keyword Args: + dtype (torch.dtype): the dtype of the tensor. + device (torch.device): the device of the tensor. Only `None` and `"cpu"` + are accepted, any other device will raise an exception. + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + existsok (bool, optional): whether it is ok to overwrite an existing file. + Defaults to ``False``. + """ + shape, device, dtype, _, filename = _proc_args_const(*args, **kwargs) + if device is not None: + device = torch.device(device) + if device.type != "cpu": + raise RuntimeError("Only CPU tensors are supported.") + if isinstance(shape, torch.Tensor): + return cls.empty( + shape, device=device, dtype=dtype, filename=filename + ).fill_(0) + result = torch.zeros((), dtype=dtype, device=device) + if shape: + if isinstance(shape[0], (list, tuple)) and len(shape) == 1: + shape = torch.Size(shape[0]) + else: + shape = torch.Size(shape) + result = result.expand(shape) + result = cls.from_tensor( + result, + filename=filename, + existsok=kwargs.pop("existsok", False), + ) + return result + + @classmethod + @overload + def empty(cls, *size, dtype=None, device=None, filename=None): ... + + @classmethod + @overload + def empty(cls, shape, *, dtype=None, device=None, filename=None): ... + + @classmethod + def empty(cls, *args, **kwargs): + # noqa: D417 + """Creates a tensor with empty content, specific shape, dtype and filename. + + Args: + shape (integers or torch.Size): the shape of the tensor. + + Keyword Args: + dtype (torch.dtype): the dtype of the tensor. + device (torch.device): the device of the tensor. Only `None` and `"cpu"` + are accepted, any other device will raise an exception. + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + existsok (bool, optional): whether it is ok to overwrite an existing file. + Defaults to ``False``. + """ + shape, device, dtype, _, filename = _proc_args_const(*args, **kwargs) + if device is not None: + device = torch.device(device) + if device.type != "cpu": + raise RuntimeError("Only CPU tensors are supported.") + result = torch.zeros((), dtype=dtype, device=device) + if isinstance(shape, torch.Tensor): + # nested tensor + shape_numel = shape.prod(-1).sum() + + if filename is None: + if dtype.is_floating_point: + size = torch.finfo(dtype).bits // 8 * shape_numel + elif dtype.is_complex: + raise ValueError( + "Complex-valued tensors are not supported by MemoryMappedTensor." + ) + elif dtype == torch.bool: + size = shape_numel + else: + # assume integer + size = torch.iinfo(dtype).bits // 8 * shape_numel + handler = _FileHandler(size) + + # buffer + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError(NESTED_TENSOR_ERR) + result = torch.frombuffer(memoryview(handler.buffer), dtype=dtype) + result = torch._nested_view_from_buffer( + result, + shape, + *offsets_strides, + ) + result = cls(result) + result._handler = handler + return result + else: + result = torch.from_file( + str(filename), + shared=True, + dtype=dtype, + size=shape_numel, + # needed when device ctx differs + device=torch.device("cpu"), + ) + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError(NESTED_TENSOR_ERR) + result = torch._nested_view_from_buffer( + result, + shape, + *offsets_strides, + ) + result = cls(result) + result.filename = filename + return result + return result + + if shape: + if isinstance(shape[0], (list, tuple)) and len(shape) == 1: + shape = torch.Size(shape[0]) + else: + shape = torch.Size(shape) + result = result.expand(shape) + result = cls.from_tensor( + result, + filename=filename, + copy_data=False, + existsok=kwargs.pop("existsok", False), + ) + return result + + @classmethod + def empty_nested(cls, *args, **kwargs): + # noqa: D417 + """Creates a tensor with empty content, specific shape, dtype and filename. + + Args: + shape (nested_shape): the shapes of the tensors. + + Keyword Args: + dtype (torch.dtype): the dtype of the tensor. + device (torch.device): the device of the tensor. Only `None` and `"cpu"` + are accepted, any other device will raise an exception. + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + existsok (bool, optional): whether it is ok to overwrite an existing file. + Defaults to ``False``. + """ + shape = kwargs.pop("shape", args[0]) + args = (torch.Size([]), *args) + _, device, dtype, _, filename = _proc_args_const(*args, **kwargs) + if device is not None: + device = torch.device(device) + if device.type != "cpu": + raise RuntimeError("Only CPU tensors are supported.") + result = torch.zeros((), dtype=dtype, device=device) + if shape: + if isinstance(shape[0], (list, tuple)) and len(shape) == 1: + shape = torch.Size(shape[0]) + else: + shape = torch.Size(shape) + result = result.expand(shape) + result = cls.from_tensor( + result, + filename=filename, + copy_data=False, + existsok=kwargs.pop("existsok", False), + ) + return result + + @classmethod + @overload + def full(cls, *size, fill_value, dtype=None, device=None, filename=None): ... + + @classmethod + @overload + def full(cls, shape, *, fill_value, dtype=None, device=None, filename=None): ... + + @classmethod + def full(cls, *args, **kwargs): + # noqa: D417 + """Creates a tensor with a single content specified by `fill_value`, specific shape, dtype and filename. + + Args: + shape (integers or torch.Size): the shape of the tensor. + + Keyword Args: + fill_value (float or equivalent): content of the tensor. + dtype (torch.dtype): the dtype of the tensor. + device (torch.device): the device of the tensor. Only `None` and `"cpu"` + are accepted, any other device will raise an exception. + filename (path or equivalent): the path to the file, if any. If none + is provided, a handler is used. + existsok (bool, optional): whether it is ok to overwrite an existing file. + Defaults to ``False``. + """ + shape, device, dtype, fill_value, filename = _proc_args_const(*args, **kwargs) + if device is not None: + device = torch.device(device) + if device.type != "cpu": + raise RuntimeError("Only CPU tensors are supported.") + result = torch.zeros((), dtype=dtype, device=device).fill_(fill_value) + if shape: + if isinstance(shape[0], (list, tuple)) and len(shape) == 1: + shape = torch.Size(shape[0]) + else: + shape = torch.Size(shape) + result = result.expand(shape) + return cls.from_tensor( + result, filename=filename, existsok=kwargs.pop("existsok", False) + ) + + @classmethod + def from_filename(cls, filename, dtype, shape, index=None): + # noqa: D417 + """Loads a MemoryMappedTensor from a given filename. + + Args: + filename (path or equivalent): the path to the file. + dtype (torch.dtype): the dtype of the tensor. + shape (torch.Size or torch.Tensor): the shape of the tensor. If + a tensor is provided, it is assumed that the tensor is a nested_tensor + instance. + index (torch-compatible index type): an index to use to build the + tensor. + + """ + writable = _is_writable(filename) + + if isinstance(shape, torch.Tensor): + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError( + "The PyTorch version isn't compatible with memmap " + "nested tensors. Please upgrade to a more recent " + "version." + ) + tensor = torch.from_file( + str(filename), + shared=writable, + dtype=dtype, + size=shape.prod(-1).sum().int(), + # needed when device ctx differs + device=torch.device("cpu"), + ) + tensor = torch._nested_view_from_buffer( + tensor, + shape, + *offsets_strides, + ) + else: + shape = torch.Size(shape) + # whether the file already existed + tensor = torch.from_file( + str(filename), + shared=writable, + dtype=dtype, + size=shape.numel(), + # needed when device ctx differs + device=torch.device("cpu"), + ) + tensor = tensor.view(shape) + + if index is not None: + tensor = tensor[index] + out = cls(tensor) + out.filename = filename + out._handler = None + out.index = index + out.parent_shape = shape + return out + + @classmethod + def from_handler(cls, handler, dtype, shape, index=None): + # noqa: D417 + """Loads a MemoryMappedTensor from a given handler. + + Args: + handler (compatible file handler): the handler for the tensor. + dtype (torch.dtype): the dtype of the tensor. + shape (torch.Size or torch.Tensor): the shape of the tensor. If + a tensor is provided, it is assumed that the tensor is a nested_tensor + instance. + index (torch-compatible index type, optional): an index to use to build the + tensor. + + """ + out = torch.frombuffer(memoryview(handler.buffer), dtype=dtype) + if isinstance(shape, torch.Tensor): + func_offset_stride = getattr( + torch, "_nested_compute_contiguous_strides_offsets", None + ) + if func_offset_stride is not None: + offsets_strides = func_offset_stride(shape) + else: + raise RuntimeError( + "The PyTorch version isn't compatible with memmap " + "nested tensors. Please upgrade to a more recent " + "version." + ) + out = torch._nested_view_from_buffer( + out, + shape, + *offsets_strides, + ) + else: + shape = torch.Size(shape) + out = torch.reshape(out, shape) + + if index is not None: + out = out[index] + out = cls(out) + out.filename = None + out._handler = handler + out.index = index + out.parent_shape = shape + return out + + @property + def _tensor(self): + raise RuntimeError( + "_tensor property has been removed. MemoryMappedTensor is now a tensor subclass " + "and can be used directly without accessing _tensor." + ) + + def __setstate__(self, state): + if "filename" in state: + self.__dict__ = type(self).from_filename(**state).__dict__ + else: + self.__dict__ = type(self).from_handler(**state).__dict__ + + def __getstate__(self): + if getattr(self, "_handler", None) is not None: + return { + "handler": self._handler, + "dtype": self.dtype, + "shape": list(self.parent_shape), + "index": self.index, + } + elif getattr(self, "_filename", None) is not None: + return { + "filename": self._filename, + "dtype": self.dtype, + "shape": self.parent_shape, + "index": self.index, + } + else: + raise RuntimeError("Could not find handler or filename.") + + def __reduce_ex__(self, protocol): + return self.__reduce__() + + def __reduce__(self): + if getattr(self, "_handler", None) is not None: + return type(self).from_handler, ( + self._handler, + self.dtype, + self.parent_shape, + self.index, + ) + elif getattr(self, "_filename", None) is not None: + return type(self).from_filename, ( + self._filename, + self.dtype, + self.parent_shape, + self.index, + ) + else: + raise RuntimeError("Could not find handler or filename.") + + @implement_for("torch", "2.0", None) + def __getitem__(self, item: IndexType) -> Self | torch.Tensor: + try: + out = super().__getitem__(item) + except ValueError as err: + if "is unbound" in str(err): + raise ValueError( + "Using first class dimension indices with MemoryMappedTensor " + "isn't supported at the moment." + ) from err + raise + try: + out_storage = out.untyped_storage() + except NotImplementedError as err: + if re.search("Cannot access storage of BatchedTensorImpl", str(err)): + raise ValueError( + "Using first class dimension indices with MemoryMappedTensor " + "isn't supported at the moment." + ) from err + raise + if out_storage.data_ptr() == self.untyped_storage().data_ptr(): + out = self._index_wrap(out, item) + return out + + @implement_for("torch", None, "2.0") + def __getitem__(self, item: IndexType) -> Self | torch.Tensor: # noqa: F811 + try: + out = super().__getitem__(item) + except ValueError as err: + if "is unbound" in str(err): + raise ValueError( + "Using first class dimension indices with MemoryMappedTensor " + "isn't supported at the moment." + ) from err + raise + # Check if result is a batched tensor (from functorch/ftdim operations) + # In PyTorch nightlies, ftdim indexing no longer raises ValueError but creates + # BatchedTensorImpl which cannot have its storage accessed + try: + from torch._C._functorch import is_batchedtensor + + if is_batchedtensor(out): + raise ValueError( + "Using first class dimension indices with MemoryMappedTensor " + "isn't supported at the moment." + ) + except ImportError: + pass + if out.storage().data_ptr() == self.storage().data_ptr(): + out = self._index_wrap(out, item) + return out + + def _index_wrap(self, tensor, item, check=False): + if check: + if tensor.untyped_storage().data_ptr() == self.untyped_storage().data_ptr(): + return self._index_wrap(tensor, item) + return tensor + tensor = MemoryMappedTensor(tensor) + tensor._handler = getattr(self, "_handler", None) + tensor.filename = getattr(self, "_filename", None) + tensor.index = item + tensor.parent_shape = getattr(self, "parent_shape", None) + return tensor + + def unbind(self, dim): + out = super().unbind(dim) + if dim < 0: + dim = self.ndim + dim + index_base = (slice(None),) * dim + return tuple( + self._index_wrap(_out, index_base + (i,)) for i, _out in enumerate(out) + ) + + def chunk(self, chunks, dim=0): + dim = _maybe_correct_neg_dim(dim, self.shape) + out = super().chunk(chunks, dim) + slices = [] + i = 0 + for chunk in out: + slices.append( + tuple(slice(None) for _ in range(dim)) + + (slice(i, i + chunk.shape[dim]),) + ) + i += chunk.shape[dim] + return tuple( + self._index_wrap(chunk, _slice, check=True) + for chunk, _slice in _zip_strict(out, slices) + ) + + def split(self, split_size, dim=0): + dim = _maybe_correct_neg_dim(dim, self.shape) + out = super().split(split_size, dim) + slices = [] + i = 0 + for chunk in out: + slices.append( + tuple(slice(None) for _ in range(dim)) + + (slice(i, i + chunk.shape[dim]),) + ) + i += chunk.shape[dim] + return tuple( + self._index_wrap(split, _slice, check=True) + for split, _slice in _zip_strict(out, slices) + ) + + +##################### +# File handler +# borrowed from mp.heap + +if sys.platform == "win32": + import _winapi + + class _FileHandler: + _rand = tempfile._RandomNameSequence() + + def __init__(self, size): + self.size = size + for _ in range(100): + name = "pym-%d-%s" % (os.getpid(), next(self._rand)) + buf = mmap.mmap(-1, size, tagname=name) + if _winapi.GetLastError() == 0: + break + # We have reopened a preexisting mmap. + buf.close() + else: + raise FileExistsError("Cannot find name for new mmap") + self.name = name + self.buffer = buf + self._state = (self.size, self.name) + + def __getstate__(self): + from multiprocessing.context import assert_spawning + + assert_spawning(self) + return self._state + + def __setstate__(self, state): + self.size, self.name = self._state = state + # Reopen existing mmap + self.buffer = mmap.mmap(-1, self.size, tagname=self.name) + # XXX Temporarily preventing buildbot failures while determining + # XXX the correct long-term fix. See issue 23060 + # assert _winapi.GetLastError() == _winapi.ERROR_ALREADY_EXISTS + +else: + + class _FileHandler: + if sys.platform == "linux": + _dir_candidates = ["/dev/shm"] + else: + _dir_candidates = [] + + def __init__(self, size, fd=-1): + self.size = size + self.fd = fd + if fd == -1: + self.fd, name = tempfile.mkstemp( + prefix="pym-%d-" % os.getpid(), dir=self._choose_dir(size) + ) + os.unlink(name) + util.Finalize(self, os.close, (self.fd,)) + os.ftruncate(self.fd, size) + self.buffer = mmap.mmap(self.fd, self.size) + + def _choose_dir(self, size): + # Choose a non-storage backed directory if possible, + # to improve performance + for d in self._dir_candidates: + st = os.statvfs(d) + if st.f_bavail * st.f_frsize >= size: # enough free space? + return d + return util.get_temp_dir() + + def _reduce_handler(handler): + if handler.fd == -1: + raise ValueError( + "Handler is unpicklable because " + "forking was enabled when it was created" + ) + return _rebuild_handler, (handler.size, reduction.DupFd(handler.fd)) + + def _rebuild_handler(size, dupfd): + detached = dupfd.detach() + return _FileHandler(size, detached) + + reduction.register(_FileHandler, _reduce_handler) + + +def _reduce_memmap(memmap_tensor): + return memmap_tensor.__reduce__() + + +reduction.register(MemoryMappedTensor, _reduce_memmap) + + +def _proc_args_const(*args, **kwargs): + if len(args) > 0: + # then the first (or the N first) args are the shape + if len(args) == 1 and isinstance(args[0], torch.Tensor): + shape = args[0] + elif len(args) == 1 and not isinstance(args[0], int): + shape = torch.Size(args[0]) + else: + shape = torch.Size(args) + else: + # we should have a "shape" keyword arg + shape = kwargs.pop("shape", None) + if shape is None: + raise TypeError("Could not find the shape argument in the arguments.") + if not isinstance(shape, torch.Tensor): + shape = torch.Size(shape) + return ( + shape, + kwargs.pop("device", None), + kwargs.pop("dtype", None), + kwargs.pop("fill_value", None), + kwargs.pop("filename", None), + ) + + +# Torch functions + +MEMMAP_HANDLED_FUNCTIONS: dict[Callable, Callable] = {} + + +def implements_for_memmap(torch_function: Callable) -> Callable[[Callable], Callable]: + """Register a torch function override for MemoryMappedTensor.""" + + @functools.wraps(torch_function) + def decorator(func: Callable) -> Callable: + MEMMAP_HANDLED_FUNCTIONS[torch_function] = func + return func + + return decorator + + +@implements_for_memmap(torch.unbind) +def _unbind(tensor, dim): + return tensor.unbind(dim) + + +@implements_for_memmap(torch.chunk) +def _chunk(input, chunks, dim=0): + return input.chunk(chunks, dim=dim) + + +def _is_writable(file_path): + file_path = str(file_path) + if os.path.exists(file_path): + return os.access(file_path, os.W_OK) + # Assume that the file can be written in the directory + return True diff --git a/lib/python3.12/site-packages/tensordict/nn/__init__.py b/lib/python3.12/site-packages/tensordict/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..08bf63616e5cbd9be43ec0eecfedc01ae468e097 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/__init__.py @@ -0,0 +1,97 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from tensordict.nn.common import ( + dispatch, + make_tensordict, + TensorDictModule, + TensorDictModuleBase, + TensorDictModuleWrapper, + WrapModule, +) +from tensordict.nn.distributions import ( + AddStateIndependentNormalScale, + CompositeDistribution, + NormalParamExtractor, + OneHotCategorical, + rand_one_hot, + TruncatedNormal, +) +from tensordict.nn.ensemble import EnsembleModule +from tensordict.nn.functional_modules import ( + get_functional, + is_functional, + make_functional, + repopulate_module, +) +from tensordict.nn.params import TensorDictParams +from tensordict.nn.probabilistic import ( + InteractionType, + ProbabilisticTensorDictModule, + ProbabilisticTensorDictSequential, + set_interaction_type, +) +from tensordict.nn.sequence import TensorDictSequential +from tensordict.nn.tensorclass_module import ( + TensorClassModuleBase, + TensorClassModuleWrapper, +) +from tensordict.nn.utils import ( + add_custom_mapping, + biased_softplus, + inv_softplus, + mappings, + set_skip_existing, + skip_existing, +) + +from .common import as_tensordict_module + +from .cudagraphs import CudaGraphModule +from .utils import composite_lp_aggregate, set_composite_lp_aggregate + +__all__ = [ + # Core modules + "TensorDictModule", + "TensorDictModuleBase", + "TensorDictModuleWrapper", + "WrapModule", + "TensorDictSequential", + "EnsembleModule", + "CudaGraphModule", + "TensorClassModuleBase", + "TensorClassModuleWrapper", + # Probabilistic modules + "ProbabilisticTensorDictModule", + "ProbabilisticTensorDictSequential", + "InteractionType", + "set_interaction_type", + # Functional modules + "make_functional", + "get_functional", + "is_functional", + "repopulate_module", + # Distributions + "AddStateIndependentNormalScale", + "CompositeDistribution", + "NormalParamExtractor", + "OneHotCategorical", + "rand_one_hot", + "TruncatedNormal", + # Parameters + "TensorDictParams", + # Utilities + "dispatch", + "make_tensordict", + "as_tensordict_module", + "add_custom_mapping", + "biased_softplus", + "inv_softplus", + "mappings", + "skip_existing", + "set_skip_existing", + "composite_lp_aggregate", + "set_composite_lp_aggregate", +] diff --git a/lib/python3.12/site-packages/tensordict/nn/tensorclass_module.py b/lib/python3.12/site-packages/tensordict/nn/tensorclass_module.py new file mode 100644 index 0000000000000000000000000000000000000000..6c0d64888575dde32ed31dde855629fabd1bb4b1 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/tensorclass_module.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterable +from dataclasses import Field +from typing import Any, cast, Generic, get_args, get_origin, TypeVar, Union + +from tensordict._td import TensorDict +from tensordict.nn.common import dispatch, TensorDictModuleBase +from tensordict.tensorclass import TensorClass +from torch import nn, Tensor + +__all__ = ["TensorClassModuleBase", "TensorClassModuleWrapper"] + + +def _tensor_class_keys(tensorclass_type: type[TensorClass]) -> list[tuple[str, ...]]: + """Extract all keys from a TensorClass type, including nested keys. + + Args: + tensorclass_type (type[TensorClass]): The TensorClass type to extract keys from. + + Returns: + list[tuple[str, ...]]: A list of key tuples representing all fields in the TensorClass. + + """ + fields = cast("Iterable[Field[Any]]", tensorclass_type.fields()) + keys: list[tuple[str, ...]] = [] + for field in fields: + key = field.name + if issubclass(field.type, TensorClass): + subkeys = _tensor_class_keys(cast(type[TensorClass], field.type)) + for subkey in subkeys: + keys.append((key,) + subkey) + else: + keys.append((key,)) + return keys + + +InputTensorClass = TypeVar("InputTensorClass", bound=TensorClass) +OutputTensorClass = TypeVar("OutputTensorClass", bound=TensorClass) + + +class TensorClassModuleWrapper(TensorDictModuleBase): + """Wrapper class for TensorClassModuleBase objects. + + This wrapper allows TensorClassModuleBase instances to be used in TensorDict-based + workflows by handling the conversion between TensorDict and TensorClass representations. + When called with a TensorDict, the wrapper converts it to a TensorClass, passes it through + the wrapped module, and converts the output back to a TensorDict. + + Args: + module (TensorClassModuleBase): The TensorClassModuleBase instance to wrap. + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.tensorclass import TensorClass + >>> from tensordict.nn import TensorClassModuleBase + >>> import torch + >>> + >>> class InputTC(TensorClass): + ... x: torch.Tensor + ... + >>> class OutputTC(TensorClass): + ... y: torch.Tensor + ... + >>> class MyModule(TensorClassModuleBase[InputTC, OutputTC]): + ... def forward(self, input: InputTC) -> OutputTC: + ... return OutputTC(y=input.x + 1, batch_size=input.batch_size) + ... + >>> module = MyModule() + >>> td_module = module.as_td_module() + >>> td = TensorDict({"x": torch.zeros(3)}, batch_size=[3]) + >>> result = td_module(td) + >>> assert "y" in result + + """ + + def __init__( + self, module: TensorClassModuleBase[InputTensorClass, OutputTensorClass] + ) -> None: + super().__init__() + self.tc_module = module + self.in_keys = _tensor_class_keys(cast(type[TensorClass], module.input_type)) + self.out_keys = _tensor_class_keys(cast(type[TensorClass], module.output_type)) + + @dispatch(auto_batch_size=False) + def forward(self, tensordict: TensorDict, *args, **kwargs) -> TensorDict: + """Forward pass converting TensorDict to TensorClass and back. + + Args: + tensordict (TensorDict): Input tensordict. + *args: Additional positional arguments. + **kwargs: Additional keyword arguments. + + Returns: + TensorDict: Output tensordict. + + """ + return self.tc_module( + self.tc_module.input_type.from_tensordict(tensordict) + ).to_tensordict() + + +InputClass = TypeVar("InputClass", bound=Union[TensorClass, Tensor]) +OutputClass = TypeVar("OutputClass", bound=Union[TensorClass, Tensor]) + + +class TensorClassModuleBase(Generic[InputClass, OutputClass], ABC, nn.Module): + """A TensorClassModuleBase is a base class for modules that operate on TensorClass instances. + + TensorClassModuleBase subclasses provide a type-safe way to define modules that work with TensorClass + inputs and outputs. The class automatically extracts input and output type information from the + generic type parameters. + + The module can be converted to a TensorDictModule using the :meth:`as_td_module` + method, allowing it to be used in TensorDict-based workflows. + + Type Parameters: + InputClass: The input type, must be a TensorClass or Tensor. + OutputClass: The output type, must be a TensorClass or Tensor. + + Attributes: + input_type (type[InputClass]): The input type class. + output_type (type[OutputClass]): The output type class. + + Examples: + >>> from tensordict.tensorclass import TensorClass + >>> from tensordict.nn import TensorClassModuleBase + >>> import torch + >>> + >>> class InputTC(TensorClass): + ... a: torch.Tensor + ... b: torch.Tensor + ... + >>> class OutputTC(TensorClass): + ... result: torch.Tensor + ... + >>> class AddModule(TensorClassModuleBase[InputTC, OutputTC]): + ... def forward(self, x: InputTC) -> OutputTC: + ... return OutputTC( + ... result=x.a + x.b, + ... batch_size=x.batch_size + ... ) + ... + >>> module = AddModule() + >>> input_tc = InputTC(a=torch.tensor([1.0]), b=torch.tensor([2.0]), batch_size=[1]) + >>> output = module(input_tc) + >>> assert output.result == torch.tensor([3.0]) + + """ + + input_type: type[InputClass] + output_type: type[OutputClass] + + def __init_subclass__(cls) -> None: + """Initialize subclass by extracting type information from generic parameters.""" + super().__init_subclass__() + for base in cls.__orig_bases__: # type:ignore[attr-defined] + origin = get_origin(base) + if origin is TensorClassModuleBase: + generic_args = get_args(base) + if generic_args: + cls.input_type, cls.output_type = generic_args + else: + raise ValueError( + "Generic input/output types not set in TensorClassModuleBase" + ) + + @abstractmethod + def forward(self, x: InputClass) -> OutputClass: + """Forward pass of the module. + + Args: + x (InputClass): Input instance. + + Returns: + OutputClass: Output instance. + + """ + ... + + def __call__(self, x: InputClass) -> OutputClass: + """Call the module's forward method. + + Args: + x (InputClass): Input instance. + + Returns: + OutputClass: Output instance. + + """ + return cast("OutputClass", super().__call__(x)) + + def as_td_module(self) -> TensorClassModuleWrapper: + """Convert this module to a TensorDictModule. + + This method wraps the TensorClassModuleBase in a TensorClassModuleWrapper, + allowing it to be used with TensorDict inputs and outputs. + + Returns: + TensorClassModuleWrapper: A wrapper that converts between TensorDict + and TensorClass representations. + + Raises: + ValueError: If either input_type or output_type is not a TensorClass. + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.tensorclass import TensorClass + >>> from tensordict.nn import TensorClassModuleBase + >>> import torch + >>> + >>> class InputTC(TensorClass): + ... x: torch.Tensor + ... + >>> class OutputTC(TensorClass): + ... y: torch.Tensor + ... + >>> class MyModule(TensorClassModuleBase[InputTC, OutputTC]): + ... def forward(self, input: InputTC) -> OutputTC: + ... return OutputTC(y=input.x * 2, batch_size=input.batch_size) + ... + >>> module = MyModule() + >>> td_module = module.as_td_module() + >>> td = TensorDict({"x": torch.ones(3)}, batch_size=[3]) + >>> result = td_module(td) + >>> assert (result["y"] == 2).all() + + """ + if not ( + issubclass(self.input_type, TensorClass) + and issubclass(self.output_type, TensorClass) + ): + raise ValueError( + "Only TensorClassModuleBase implementations with both input and " + "output type as TensorClass can be converted to TensorDictModule" + ) + return TensorClassModuleWrapper(self) # type:ignore[arg-type,type-var] diff --git a/lib/python3.12/site-packages/tensordict/persistent.py b/lib/python3.12/site-packages/tensordict/persistent.py new file mode 100644 index 0000000000000000000000000000000000000000..3c5a73e30138891389ae2163ecdae13b6256e869 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/persistent.py @@ -0,0 +1,1534 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +"""Persistent tensordicts (H5 and others).""" +from __future__ import annotations + +import importlib +import os + +import tempfile +import warnings +import weakref +from functools import wraps +from pathlib import Path +from typing import Any, Callable, Tuple, Type, TYPE_CHECKING + +import numpy as np +import torch + +from tensordict._td import ( + _TensorDictKeysView, + _unravel_key_to_tuple, + CompatibleType, + NO_DEFAULT, + TensorDict, +) +from tensordict._tensorcollection import TensorCollection +from tensordict.base import ( + _default_is_leaf, + _is_leaf_nontensor, + _register_tensor_class, + is_tensor_collection, + T, + TensorDictBase, +) +from tensordict.memmap import MemoryMappedTensor +from tensordict.utils import ( + _as_context_manager, + _CloudpickleWrapper, + _KEY_ERROR, + _LOCK_ERROR, + _parse_to, + _proc_init, + _split_tensordict, + _zip_strict, + cache, + erase_cache, + expand_right, + IndexType, + is_non_tensor, + lock_blocked, + NestedKey, + NUMPY_TO_TORCH_DTYPE_DICT, + unravel_key, +) +from torch import multiprocessing as mp, Tensor + + +_has_h5 = importlib.util.find_spec("h5py", None) is not None + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +class _Visitor: + def __init__(self, fun=None): + self.elts = [] + self.fun = fun + + def __call__(self, name): + self.elts.append(name) + + def __iter__(self): + if self.fun is None: + yield from self.elts + else: + for elt in self.elts: + yield self.fun(elt) + + +class _PersistentTDKeysView(_TensorDictKeysView): + def __iter__(self): + # For consistency with tensordict where currently a non-tensor is stored in a + # tensorclass and hence can be seen as a nested tensordict + # that situation should be clarified + read_non_tensor = self.is_leaf is _is_leaf_nontensor or not self.leaves_only + if self.include_nested: + visitor = _Visitor(lambda key: unravel_key(tuple(key.split("/")))) + self.tensordict.file.visit(visitor) + else: + visitor = self.tensordict.file.keys() + for key in visitor: + metadata = self.tensordict._get_metadata(key) + if metadata.get("non_tensor"): + if read_non_tensor: + yield key + else: + continue + elif metadata.get("array"): + yield key + elif not self.leaves_only and ( + not isinstance(key, tuple) or self.include_nested + ): + yield key + + def __contains__(self, key): + key = unravel_key(key) + return key in list(self) + + +class PersistentTensorDict(TensorDictBase): + """Persistent TensorDict implementation. + + :class:`PersistentTensorDict` instances provide an interface with data stored + on disk such that access to this data is made easy while still taking advantage + from the fast access provided by the backend. + + Like other :class:`TensorDictBase` subclasses, :class:`PersistentTensorDict` + has a ``device`` attribute. This does *not* mean that the data is being stored + on that device, but rather that when loaded, the data will be cast onto + the desired device. + + Keyword Args: + batch_size (torch.Size or compatible): the tensordict batch size. + Defaults to ``torch.Size(())``. + filename (str, optional): the path to the h5 file. Exclusive with ``group``. + group (h5py.Group, optional): a file or a group that contains data. Exclusive with ``filename``. + mode (str, optional): Reading mode. Defaults to ``"r"``. + backend (str, optional): storage backend. Currently only ``"h5"`` is supported. + device (torch.device or compatible, optional): device of the tensordict. + Defaults to ``None`` (ie. default PyTorch device). + **kwargs: kwargs to be passed to :meth:`h5py.File.create_dataset`. + + .. note:: + Currently, PersistentTensorDict instances are not closed when getting out-of-scope. + This means that it is the responsibility of the user to close them if necessary. + + Examples: + >>> import tempfile + >>> with tempfile.NamedTemporaryFile() as f: + ... data = PersistentTensorDict(file=f, batch_size=[3], mode="w") + ... data["a", "b"] = torch.randn(3, 4) + ... print(data) + + """ + + _td_dim_names = None + LOCKING = None + + def __init__( + self, + *, + batch_size=None, + filename=None, + group=None, + mode="r", + backend="h5", + device=None, + **kwargs, + ): + if batch_size is None: + batch_size = torch.Size(()) + self._locked_tensordicts = [] + self._lock_id = set() + if not _has_h5: + raise ModuleNotFoundError("Could not load h5py.") + import h5py + + self.filename = filename + self.mode = mode + if backend != "h5": + raise NotImplementedError + if filename is not None and group is None: + self.file = h5py.File(filename, mode, locking=self.LOCKING) + elif group is not None: + self.file = group + else: + raise RuntimeError( + f"Either group or filename must be provided, and not both. Got group={group} and filename={filename}." + ) + self._batch_size = torch.Size(batch_size) + self._device = torch.device(device) if device is not None else None + self._is_shared = False + self._is_memmap = False + self.kwargs = kwargs + + # we use this to allow nested tensordicts to have a different batch-size + self._nested_tensordicts = {} + self._pin_mem = False + + # this must be kept last + self._check_batch_size(self._batch_size) + + @classmethod + def from_h5(cls, filename, *, mode="r", batch_size: torch.size | None = None): + """Creates a PersistentTensorDict from a h5 file. + + This function will automatically determine the batch-size for each nested tensordict (unless ``batch_size`` + is provided). + + Args: + filename (str): The path to the h5 file. + + Keyword Args: + mode (str, optional): Reading mode. Defaults to ``"r"``. + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to None (batch-size automatically + determined). + + Returns: + A PersistentTensorDict representation of the input h5 file. + + Examples: + >>> ptd = PersistentTensorDict.from_h5("path/to/file.h5") + >>> print(ptd) + PersistentTensorDict( + fields={ + key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + out = cls(filename=filename, mode=mode, batch_size=batch_size) + if batch_size is None: + # determine batch size + _set_max_batch_size(out) + return out + + @classmethod + def from_dict( + cls, + input_dict, + filename, + *, + auto_batch_size: bool = False, + batch_size=None, + device=None, + **kwargs, + ): + """Converts a dictionary or a TensorDict to a h5 file. + + Args: + input_dict (dict, TensorDict or compatible): data to be stored as h5. + filename (str or path): path to the h5 file. + + Keyword Args: + auto_batch_size (bool, optional): if ``True``, the batch size will be computed automatically. + Defaults to ``False``. + batch_size (tensordict batch-size, optional): if provided, batch size + of the tensordict. If not, the batch size will be gathered from the + input structure (if present) or determined automatically. + device (torch.device or compatible, optional): the device where to + expect the tensor once they are returned. Defaults to ``None`` + (on cpu by default). + **kwargs: kwargs to be passed to :meth:`h5py.File.create_dataset`. + + Returns: + A :class:`PersitentTensorDict` instance linked to the newly created file. + + """ + import h5py + + file = h5py.File(filename, "w", locking=cls.LOCKING) + _has_batch_size = True + if batch_size is None: + if is_tensor_collection(input_dict): + batch_size = input_dict.batch_size + else: + _has_batch_size = False + batch_size = torch.Size([]) + + # let's make a tensordict first + out = cls(group=file, batch_size=batch_size, device=device, **kwargs) + if is_tensor_collection(input_dict): + out.update(input_dict) + else: + out.update(TensorDict(input_dict, batch_size=batch_size)) + if not _has_batch_size: + _set_max_batch_size(out) + return out + + def close(self): + """Closes the persistent tensordict.""" + self.file.close() + + def _process_key(self, key): + key = _unravel_key_to_tuple(key) + return "/".join(key) + + def _check_batch_size(self, batch_size) -> None: + for key in self.keys(include_nested=True, leaves_only=True): + key = self._process_key(key) + array = self.file[key] + if _is_non_tensor_h5(array): + continue + size = array.shape + if torch.Size(size[: len(batch_size)]) != batch_size: + raise ValueError( + f"batch size and array size mismatch: array.shape={size}, batch_size={batch_size}." + ) + + def _get_array(self, key, default=NO_DEFAULT): + try: + key = self._process_key(key) + array = self.file[key] + return array + except KeyError: + if default is not NO_DEFAULT: + return default + raise KeyError(f"key {key} not found in PersistentTensorDict {self}") + + def _process_array(self, key, array): + import h5py + + if isinstance(array, (h5py.Dataset,)): + if self.device is not None: + device = self.device + else: + device = torch.device("cpu") + # we convert to an array first to avoid "Creating a tensor from a list of numpy.ndarrays is extremely slow." + if not _is_non_tensor_h5(array): + array = array[()] + out = torch.as_tensor(array, device=device) + if self._pin_mem: + out = out.pin_memory() + else: + from tensordict.tensorclass import NonTensorData + + array = array[()] + out = NonTensorData( + data=array, device=device, batch_size=self.batch_size + ) + return out + else: + out = self._nested_tensordicts.get(key) + if out is None: + out = self._nested_tensordicts[key] = PersistentTensorDict( + group=array, + batch_size=self.batch_size, + device=self.device, + ) + return out + + @cache # noqa: B019 + def _get_str(self, key: NestedKey, default, **kwargs): + key = _unravel_key_to_tuple(key) + array = self._get_array(key, default) + if array is default: + return array + return self._process_array(key, array) + + _get_tuple = _get_str + + def get_at( + self, key: NestedKey, idx: IndexType, default: CompatibleType = NO_DEFAULT + ) -> CompatibleType: + import h5py + + array = self._get_array(key, default) + if isinstance(array, (h5py.Dataset,)): + if self.device is not None: + device = self.device + else: + device = torch.device("cpu") + # indexing must be done before converting to tensor. + idx = self._process_index(idx, array) + # `get_at` is there to save us. + try: + out = torch.as_tensor(array[idx], device=device) + except TypeError as err: + if "Boolean indexing array has incompatible shape" in str(err): + # Known bug in h5py: cannot broadcast boolean mask on the right as + # done in np and torch. Therefore we put a performance warning + # and convert to torch tensor first. + warnings.warn( + "Indexing an h5py.Dataset object with a boolean mask " + "that needs broadcasting does not work directly. " + "tensordict will cast the entire array in memory and index it using the mask. " + "This is suboptimal and may lead to performance issue." + ) + out = torch.as_tensor(np.asarray(array), device=device)[idx] + else: + raise err + if self._pin_mem: + return out.pin_memory() + return out + elif array is not default: + out = self._nested_tensordicts.get(key) + if out is None: + out = self._nested_tensordicts[key] = PersistentTensorDict( + group=array, + batch_size=self.batch_size, + device=self.device, + ) + return out._get_sub_tensordict(idx) + else: + return default + + def _get_metadata(self, key): + """Gets the metadata for an entry. + + This method avoids creating a tensor from scratch, and just reads the metadata of the array. + """ + import h5py + + array = self._get_array(key) + if ( + isinstance(array, (h5py.Dataset,)) + and array.dtype in NUMPY_TO_TORCH_DTYPE_DICT + ): + shape = torch.Size(array.shape) + return { + "dtype": NUMPY_TO_TORCH_DTYPE_DICT[array.dtype], + "shape": shape, + "dim": len(shape), + "array": True, + } + elif ( + isinstance(array, (h5py.Dataset,)) + and array.dtype not in NUMPY_TO_TORCH_DTYPE_DICT + ): + return {"non_tensor": True} + else: + val = self.get(key) + shape = val.shape + return { + "dtype": None, + "shape": shape, + "dim": len(shape), + "array": False, + } + + @classmethod + def _process_index(cls, idx, array=None): + if isinstance(idx, tuple): + return tuple(cls._process_index(_idx, array) for _idx in idx) + if isinstance(idx, torch.Tensor): + return idx.cpu().detach().numpy() + if isinstance(idx, (range, list)): + return np.asarray(idx) + return idx + + def __getitem__(self, item: IndexType) -> Self | Tensor | TensorCollection | Any: + if isinstance(item, str) or ( + isinstance(item, tuple) and _unravel_key_to_tuple(item) + ): + result = self.get(item, default=NO_DEFAULT) + if is_non_tensor(result): + result_data = getattr(result, "data", NO_DEFAULT) + if result_data is NO_DEFAULT: + return result.tolist(as_linked_list=True) + return result_data + return result + if isinstance(item, list): + # convert to tensor + item = torch.tensor(item) + return self._get_sub_tensordict(item) + + __getitems__ = __getitem__ + + def __setitem__(self, index: IndexType, value: Any): + index_unravel = _unravel_key_to_tuple(index) + if index_unravel: + return self.set(index_unravel, value, inplace=True) + + if isinstance(index, list): + # convert to tensor + index = torch.tensor(index) + sub_td = self._get_sub_tensordict(index) + err_set_batch_size = None + if not isinstance(value, TensorDictBase): + value = TensorDict.from_dict(value, batch_size=[]) + # try to assign the current shape. If that does not work, we can + # try to expand + try: + value.batch_size = sub_td.batch_size + except RuntimeError as err0: + err_set_batch_size = err0 + if value.shape != sub_td.shape: + try: + value = value.expand(sub_td.shape) + except RuntimeError as err: + if err_set_batch_size is not None: + raise err from err_set_batch_size + raise RuntimeError( + f"Cannot broadcast the tensordict {value} to the shape of the indexed persistent tensordict {self}[{index}]." + ) from err + sub_td.update(value, inplace=True) + + @cache # noqa: B019 + def _valid_keys(self): + keys = [] + for key in self.file.keys(): + metadata = self._get_metadata(key) + if not metadata.get("non_tensor"): + keys.append(key) + return keys + + # @cache # noqa: B019 + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> _PersistentTDKeysView: + if is_leaf not in (None, _default_is_leaf, _is_leaf_nontensor): + raise ValueError( + f"is_leaf {is_leaf} is not supported within tensordicts of type {type(self).__name__}." + ) + return _PersistentTDKeysView( + tensordict=self, + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=is_leaf, + sort=sort, + ) + + def _items_metadata(self, include_nested=False, leaves_only=False): + """Iterates over the metadata of the PersistentTensorDict.""" + for key in self.keys(include_nested, leaves_only): + yield (key, self._get_metadata(key)) + + def _values_metadata(self, include_nested=False, leaves_only=False): + """Iterates over the metadata of the PersistentTensorDict.""" + for key in self.keys(include_nested, leaves_only): + yield self._get_metadata(key) + + def _change_batch_size(self, value): + raise NotImplementedError + + def _stack_onto_( + self, list_item: list[CompatibleType], dim: int + ) -> PersistentTensorDict: + for key in self.keys(): + vals = [td._get_str(key, None) for td in list_item] + if all(v is None for v in vals): + continue + stacked = torch.stack(vals, dim=dim) + self.set_(key, stacked) + return self + + @property + def batch_size(self): + return self._batch_size + + @batch_size.setter + def batch_size(self, value): + _batch_size = self._batch_size + try: + self._batch_size = torch.Size(value) + self._check_batch_size(self._batch_size) + except ValueError: + self._batch_size = _batch_size + + _erase_names = TensorDict._erase_names + _has_names = TensorDict._has_names + _set_names = TensorDict._set_names + names = TensorDict.names + + def _rename_subtds(self, names): + if names is None: + names = [None] * self.ndim + for item in self._nested_tensordicts.values(): + if is_tensor_collection(item): + td_names = list(names) + [None] * (item.ndim - self.ndim) + item.rename_(*td_names) + + def contiguous(self): + """Materializes a PersistentTensorDict on a regular TensorDict.""" + return self.to_tensordict() + + @lock_blocked + def del_(self, key): + key = self._process_key(key) + del self.file[key] + return self + + def detach_(self): + # PersistentTensorDict do not carry gradients. This is a no-op + return self + + @property + def device(self): + return self._device + + def empty( + self, recurse=False, *, batch_size=None, device=NO_DEFAULT, names=None + ) -> T: + if recurse: + out = self.empty( + recurse=False, batch_size=batch_size, device=device, names=names + ) + for key, val in self.items(): + if is_tensor_collection(val): + out._set_str( + key, + val.empty( + recurse=True, + batch_size=batch_size, + device=device, + names=names, + ), + inplace=False, + validated=True, + non_blocking=False, + ) + return out + return TensorDict( + {}, + device=self.device if device is NO_DEFAULT else device, + batch_size=self.batch_size if batch_size is None else batch_size, + names=self.names if names is None and self._has_names() else names, + ) + + def _propagate_lock(self, lock_parents_weakrefs=None, *, is_compiling): + """Registers the parent tensordict that handles the lock.""" + self._is_locked = True + if lock_parents_weakrefs is not None: + lock_parents_weakrefs = [ + ref + for ref in lock_parents_weakrefs + if not any(refref is ref for refref in self._lock_parents_weakrefs) + ] + if not is_compiling: + is_root = lock_parents_weakrefs is None + if is_root: + lock_parents_weakrefs = [] + else: + self._lock_parents_weakrefs = ( + self._lock_parents_weakrefs + lock_parents_weakrefs + ) + lock_parents_weakrefs = list(lock_parents_weakrefs) + lock_parents_weakrefs.append(weakref.ref(self)) + + for _td in self._nested_tensordicts.values(): + _td._propagate_lock(lock_parents_weakrefs, is_compiling=is_compiling) + + @erase_cache + def _propagate_unlock(self): + # if we end up here, we can clear the graph associated with this td + self._is_locked = False + + self._is_shared = False + self._is_memmap = False + + sub_tds = [] + for _td in self._nested_tensordicts.values(): + sub_tds.extend(_td._propagate_unlock()) + sub_tds.append(_td) + return sub_tds + + def zero_(self) -> Self: + for key in self.keys(): + self.fill_(key, 0) + return self + + def entry_class(self, key: NestedKey) -> type: + entry_class = self._get_metadata(key) + is_array = entry_class.get("array") + if is_array: + return torch.Tensor + elif is_array is False: + return PersistentTensorDict + else: + raise RuntimeError(f"Encountered a non-numeric data {key}.") + + def is_contiguous(self): + return False + + def masked_fill(self, mask, value): + return self.to_tensordict().masked_fill(mask, value) + + def where( + self, + condition: torch.Tensor, + other: torch.Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ): + return self.to_tensordict().where( + condition=condition, + other=other, + out=out, + pad=pad, + update_batch_size=update_batch_size, + ) + + def masked_fill_(self, mask, value): + for key in self.keys(include_nested=True, leaves_only=True): + array = self._get_array(key) + array[expand_right(mask, array.shape).cpu().numpy()] = value + return self + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't allowed for persistent tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't allowed for persistent tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_tensor( + self, + key: NestedKey, + tensor: torch.Tensor, + *, + copy_data: bool = True, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't allowed for persistent tensordicts." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + num_threads: int = 0, + ) -> PersistentTensorDict: + raise RuntimeError( + "Cannot build a memmap TensorDict in-place from a PersistentTensorDict. Use `td.memmap()` instead." + ) + + def _memmap_( + self, + *, + prefix: str | None, + copy_existing: bool, + executor, + futures, + inplace, + like, + share_non_tensor, + existsok, + robust_key, + ) -> T: + if inplace: + raise RuntimeError("Cannot call memmap inplace in a persistent tensordict.") + + # re-implements this to make it faster using the meta-data + def save_metadata(data: TensorDictBase, filepath, metadata=None): + if metadata is None: + metadata = {} + metadata.update( + { + "shape": list(data.shape), + "device": str(data.device), + "_type": str(type(self)), + } + ) + with open(filepath, "wb") as json_metadata: + from tensordict.utils import json_dumps + + json_str = json_dumps(metadata) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + json_metadata.write(json_str.encode("utf-8")) + else: + json_metadata.write(json_str) + + if prefix is not None: + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + metadata = {} + if not self.keys(): + raise Exception( + "memmap_like() must be called when the TensorDict is (partially) " + "populated. Set a tensor first." + ) + dest = TensorDict( + {}, + batch_size=self.batch_size, + names=self.names if self._has_names() else None, + device=torch.device("cpu"), + ) + dest._is_memmap = True + for key, value in self._items_metadata(): + if not value["array"]: + value = self._get_str(key, default=NO_DEFAULT) + dest._set_str( + key, + value._memmap_( + prefix=prefix / key if prefix is not None else None, + executor=executor, + like=like, + copy_existing=copy_existing, + futures=futures, + inplace=inplace, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ), + inplace=False, + validated=True, + non_blocking=False, + ) + if prefix is not None: + metadata[key] = { + "type": type(value).__name__, + } + continue + else: + value = self._get_str(key, default=NO_DEFAULT) + if prefix is not None: + metadata[key] = { + "dtype": str(value.dtype), + "shape": list(value.shape), + "device": str(value.device), + } + + def _populate( + tensordict=dest, key=key, value=value, prefix=prefix, like=like + ): + val = MemoryMappedTensor.from_tensor( + value, + filename=( + str(prefix / f"{key}.memmap") + if prefix is not None + else None + ), + copy_data=not like, + copy_existing=copy_existing, + existsok=existsok, + ) + tensordict._set_str( + key, + val, + inplace=False, + validated=True, + non_blocking=False, + ) + + if executor is None: + _populate() + else: + futures.append(executor.submit(_populate)) + + if prefix is not None: + if executor is None: + save_metadata(dest, prefix / "meta.json", metadata) + else: + futures.append( + executor.submit(save_metadata, dest, prefix / "meta.json", metadata) + ) + return dest + + _load_memmap = TensorDict._load_memmap + + def pin_memory(self, *args, **kwargs): + raise RuntimeError( + f"Cannot pin memory of a {type(self).__name__}. Call to_tensordict() before making this call." + ) + + @lock_blocked + def popitem(self) -> Tuple[NestedKey, CompatibleType]: + raise NotImplementedError( + f"popitem not implemented for class {type(self).__name__}." + ) + + def map( + self, + fn: Callable, + dim: int = 0, + num_workers: int | None = None, + *, + out: TensorDictBase = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = False, + pbar: bool = False, + mp_start_method: str | None = None, + ): + if pool is None: + if num_workers is None: + num_workers = mp.cpu_count() # Get the number of CPU cores + if generator is None: + generator = torch.Generator() + seed = ( + torch.empty((), dtype=torch.int64).random_(generator=generator).item() + ) + if mp_start_method is not None: + ctx = mp.get_context(mp_start_method) + else: + ctx = mp.get_context() + + queue = ctx.Queue(maxsize=num_workers) + for i in range(num_workers): + queue.put(i) + with ctx.Pool( + processes=num_workers, + initializer=_proc_init, + initargs=(seed, queue, worker_threads), + maxtasksperchild=max_tasks_per_child, + ) as pool: + return self.map( + fn, + dim=dim, + chunksize=chunksize, + pool=pool, + index_with_generator=index_with_generator, + ) + num_workers = pool._processes + dim_orig = dim + if dim < 0: + dim = self.ndim + dim + if dim < 0 or dim >= self.ndim: + raise ValueError(f"Got incompatible dimension {dim_orig}") + + self_split = _split_tensordict( + self, + chunksize, + num_chunks, + num_workers, + dim, + use_generator=index_with_generator, + to_tensordict=True, + ) + if not index_with_generator: + length = len(self_split) + self_split = tuple(split.to_tensordict() for split in self_split) + else: + length = None + + if out is not None and (out.is_shared() or out.is_memmap()): + + def wrap_fn_with_out(fn, out): + @wraps(fn) + def newfn(item_and_out): + item, out = item_and_out + result = fn(item) + out.update_(result) + return + + out_split = _split_tensordict( + out, + chunksize, + num_chunks, + num_workers, + dim, + use_generator=index_with_generator, + ) + return _CloudpickleWrapper(newfn), _zip_strict(self_split, out_split) + + fn, self_split = wrap_fn_with_out(fn, out) + out = None + + call_chunksize = 1 + imap = pool.imap(fn, self_split, call_chunksize) + + if pbar and importlib.util.find_spec("tqdm", None) is not None: + import tqdm + + imap = tqdm.tqdm(imap, total=length) + + imaplist = [] + start = 0 + for item in imap: + if item is not None: + if out is not None: + if chunksize != 0: + end = start + item.shape[dim] + chunk = slice(start, end) + out[chunk].update_(item) + start = end + else: + out[start].update_(item) + start += 1 + else: + imaplist.append(item) + del imap + + # support inplace modif + if imaplist: + if chunksize == 0: + out = torch.stack(imaplist, dim) + else: + out = torch.cat(imaplist, dim) + return out + + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> PersistentTensorDict: + old_key = self._process_key(old_key) + new_key = self._process_key(new_key) + try: + self.file.move(old_key, new_key) + except ValueError as err: + raise KeyError(f"key {new_key} already present in TensorDict.") from err + return self + + def fill_(self, key: NestedKey, value: float | bool) -> TensorDictBase: + """Fills a tensor pointed by the key with the a given value. + + Args: + key (str): key to be remaned + value (Number, bool): value to use for the filling + + Returns: + self + + """ + md = self._get_metadata(key) + if md.get("array"): + array = self._get_array(key) + array[:] = value + else: + nested = self.get(key) + for subkey in nested.keys(): + nested.fill_(subkey, value) + return self + + def _create_nested_str(self, key): + self.file.create_group(key) + target_td = self._get_str(key, default=NO_DEFAULT) + return target_td + + def _select( + self, *keys: NestedKey, inplace: bool = False, strict: bool = True + ) -> PersistentTensorDict: + raise NotImplementedError( + "Cannot call select on a PersistentTensorDict. " + "Create a regular tensordict first using the `to_tensordict` method." + ) + + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> PersistentTensorDict: + raise NotImplementedError( + "Cannot call exclude on a PersistentTensorDict. " + "Create a regular tensordict first using the `to_tensordict` method." + ) + + @_as_context_manager() + def flatten_keys(self, separator: str = ".", inplace: bool = False) -> T: + if inplace: + raise ValueError( + "Cannot call flatten_keys in_place with a PersistentTensorDict." + ) + return self.to_tensordict().flatten_keys(separator=separator) + + @_as_context_manager() + def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> T: + if inplace: + raise ValueError( + "Cannot call unflatten_keys in_place with a PersistentTensorDict." + ) + return self.to_tensordict().unflatten_keys(separator=separator) + + def share_memory_(self): + raise NotImplementedError( + "Cannot call share_memory_ on a PersistentTensorDict. " + "Create a regular tensordict first using the `to_tensordict` method." + ) + + def to(self, *args, **kwargs: Any) -> PersistentTensorDict: + ( + device, + dtype, + non_blocking, + convert_to_format, + batch_size, + non_blocking_pin, + num_threads, + inplace, + ) = _parse_to(*args, **kwargs) + if inplace: + raise TypeError(f"Cannot use inplace=True with {type(self).__name__}.to().") + + if non_blocking_pin: + raise RuntimeError( + f"Cannot use non_blocking_pin=True {type(self).__name__}.to(). Call " + f"`to_tensordict()` before executing this code." + ) + result = self + if device is not None and dtype is None and device == self.device: + return result + if dtype is not None: + return self.to_tensordict().to(*args, **kwargs) + result = self + if device is not None: + result = result.clone(False) + result._device = device + for key, nested in list(result._nested_tensordicts.items()): + result._nested_tensordicts[key] = nested.to(device) + if batch_size is not None: + result.batch_size = batch_size + return result + + def _to_numpy(self, value): + if hasattr(value, "requires_grad") and value.requires_grad: + raise RuntimeError("Cannot set a tensor that has requires_grad=True.") + if isinstance(value, torch.Tensor): + out = value.cpu().detach().numpy() + elif isinstance(value, dict): + out = TensorDict(value, []) + elif is_non_tensor(value): + value = value.data + if isinstance(value, str): + return value + import h5py + + out = np.array(value) + out = out.astype(h5py.opaque_dtype(out.dtype)) + elif is_tensor_collection(value): + out = value + elif isinstance(value, (np.ndarray,)): + out = value + else: + raise NotImplementedError( + f"Cannot set values of type {value} in a PersistentTensorDict." + ) + return out + + def _set( + self, + key: NestedKey, + value: Any, + *, + inplace: bool = False, + idx=None, + validated: bool = False, + ignore_lock: bool = False, + non_blocking: bool = False, + ) -> PersistentTensorDict: + if not validated: + value = self._validate_value(value, check_shape=idx is None) + value = self._to_numpy(value) + if not inplace: + if idx is not None: + raise RuntimeError("Cannot pass an index to _set when inplace=False.") + elif self.is_locked and not ignore_lock: + raise RuntimeError(_LOCK_ERROR) + # shortcut set if we're placing a tensordict + key = _unravel_key_to_tuple(key) + first_key, subkey = key[0], key[1:] + if is_tensor_collection(value): + target_td = self._get_str(first_key, default=None) + if target_td is None: + self.file.create_group(first_key) + target_td = self._get_str(first_key, default=NO_DEFAULT) + target_td.batch_size = value.batch_size + elif not is_tensor_collection(target_td): + raise RuntimeError( + f"cannot set a tensor collection in place of a non-tensor collection in {type(self).__name__}. " + f"Got self.get({first_key})={target_td} and value={value}." + ) + if idx is None: + if len(subkey): + target_td.set(subkey, value, inplace=inplace) + else: + target_td.update(value, inplace=inplace) + else: + if len(subkey): + target_td.set_at_(subkey, value, idx=idx) + else: + target_td.update_at_(value, idx=idx) + + return self + + if inplace: + # could be called before but will go under further refactoring of set + key = self._process_key(key) + array = self.file[key] + if idx is None: + idx = () + else: + idx = self._process_index(idx, array) + try: + array[idx] = value + except TypeError as err: + if "Boolean indexing array has incompatible shape" in str(err): + # Known bug in h5py: cannot broadcast boolean mask on the right as + # done in np and torch. Therefore we put a performance warning + # and convert to torch tensor first. + warnings.warn( + "Indexing an h5py.Dataset object with a boolean mask " + "that needs broadcasting does not work directly. " + "tensordict will cast the entire array in memory and index it using the mask. " + "This is suboptimal and may lead to performance issue." + ) + idx = tuple( + ( + expand_right(torch.as_tensor(_idx), array.shape).numpy() + if _idx.dtype == np.dtype("bool") + else _idx + ) + for _idx in idx + ) + array[idx] = torch.as_tensor(value) + else: + raise err + + else: + key = self._process_key(key) + try: + self.file.create_dataset(key, data=value, **self.kwargs) + except (ValueError, OSError) as err: + if "name already exists" in str(err): + warnings.warn( + "Replacing an array with another one is inefficient. " + "Consider using different names or populating in-place using `inplace=True`." + ) + del self.file[key] + self.file.create_dataset(key, data=value, **self.kwargs) + # If we have a nested key, let's make sure we have the corresponding TD registered + if subkey: + self._get_tuple((first_key, *subkey[:-1]), default=NO_DEFAULT) + return self + + def _convert_inplace(self, inplace, key): + key = self._process_key(key) + if inplace is not False: + has_key = key in self.file + if inplace is True and not has_key: # inplace could be None + raise KeyError( + _KEY_ERROR.format(key, type(self).__name__, sorted(self.keys())) + ) + inplace = has_key + return inplace + + def _set_non_tensor(self, key: NestedKey, value: Any): + raise NotImplementedError( + f"set_non_tensor is not compatible with the tensordict type {type(self).__name__}." + ) + + def _set_str( + self, + key: str, + value: Any, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, + ): + inplace = self._convert_inplace(inplace, key) + return self._set( + key, + value, + inplace=inplace, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + + def _set_tuple(self, key, value, *, inplace, validated, non_blocking): + key = _unravel_key_to_tuple(key) + if len(key) == 1: + return self._set_str( + key[0], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + elif key[0] in self.keys(): + return self._get_str(key[0], NO_DEFAULT)._set_tuple( + key[1:], + value, + inplace=inplace, + validated=validated, + non_blocking=non_blocking, + ) + inplace = self._convert_inplace(inplace, key) + return self._set( + key, value, inplace=inplace, validated=validated, non_blocking=non_blocking + ) + + def _set_at_str(self, key, value, idx, *, validated, non_blocking): + return self._set( + key, + value, + inplace=True, + idx=idx, + validated=validated, + non_blocking=non_blocking, + ) + + def _set_at_tuple(self, key, value, idx, *, validated, non_blocking): + return self._set( + key, + value, + inplace=True, + idx=idx, + validated=validated, + non_blocking=non_blocking, + ) + + def _set_metadata(self, orig_metadata_container: PersistentTensorDict): + for key, td in orig_metadata_container._nested_tensordicts.items(): + array = self._get_array(key) + self._nested_tensordicts[key] = PersistentTensorDict( + group=array, + batch_size=td.batch_size, + device=td.device, + ) + self._nested_tensordicts[key].names = td._td_dim_names + self._nested_tensordicts[key]._set_metadata(td) + + def _clone(self, recurse: bool = True, newfile=None) -> PersistentTensorDict: + import h5py + + if recurse: + # this should clone the h5 to a new location indicated by newfile + if newfile is None: + warnings.warn( + "A destination should be provided when cloning a " + "PersistentTensorDict. A temporary file will be used " + "instead. Use `recurse=False` to keep track of the original data " + "with a new PersistentTensorDict instance." + ) + tmpfile = tempfile.NamedTemporaryFile(delete=False) + newfile = tmpfile.name + tmpfile.close() # Close file handle before h5py opens it + f_dest = h5py.File(newfile, "w", locking=self.LOCKING) + f_src = self.file + for key in self.keys(include_nested=True, leaves_only=True): + key = self._process_key(key) + f_dest.create_dataset(key, data=f_src[key], **self.kwargs) + # f_src.copy(f_src[key], f_dest[key], "DataSet") + # create a non-recursive copy and update the file + # this way, we can keep the batch-size of every nested tensordict + clone = self.clone(False) + clone.file = f_dest + clone.filename = newfile + clone._pin_mem = False + names = self._td_dim_names + if names is not None: + names = list(names) + clone.names = names + clone._nested_tensordicts = {} + clone._set_metadata(self) + return clone + else: + # we need to keep the batch-size of nested tds, which we do manually + nested_tds = { + key: td.clone(False) for key, td in self._nested_tensordicts.items() + } + filename = self.filename + file = self.file if filename is None else None + clone = PersistentTensorDict( + filename=filename, + group=file, + mode=self.mode, + backend="h5", + device=self.device, + batch_size=self.batch_size, + ) + clone._nested_tensordicts = nested_tds + clone._pin_mem = False + clone.names = self._td_dim_names + return clone + + def __getstate__(self): + state = self.__dict__.copy() + filename = state["file"].file.filename + group_name = state["file"].name + state["file"] = None + state["filename"] = filename + state["group_name"] = group_name + state["__lock_parents_weakrefs"] = None + return state + + def __setstate__(self, state): + import h5py + + state["file"] = h5py.File( + state["filename"], mode=state["mode"], locking=self.LOCKING + ) + if state["group_name"] != "/": + state["file"] = state["file"][state["group_name"]] + del state["group_name"] + self.__dict__.update(state) + if self._is_locked: + # this can cause avoidable overhead, as we will be locking the leaves + # then locking their parent, and the parent of the parent, every + # time re-locking tensordicts that have already been locked. + # To avoid this, we should lock only at the root, but it isn't easy + # to spot what the root is... + self._is_locked = False + self.lock_() + + def _add_batch_dim(self, *, in_dim, vmap_level): + raise RuntimeError("Persistent tensordicts cannot be used with vmap.") + + def _remove_batch_dim(self, vmap_level, batch_size, out_dim): ... + + def _maybe_remove_batch_dim(self, funcname, vmap_level, batch_size, out_dim): ... + + def _view(self, *args, **kwargs): + raise RuntimeError( + "Cannot call `view` on a persistent tensordict. Call `reshape` instead." + ) + + def _transpose(self, dim0, dim1): + raise RuntimeError( + "Cannot call `transpose` on a persistent tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _permute( + self, + *args, + **kwargs, + ): + raise RuntimeError( + "Cannot call `permute` on a persistent tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _squeeze(self, dim=None): + raise RuntimeError( + "Cannot call `squeeze` on a persistent tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def _unsqueeze(self, dim: int): + raise RuntimeError( + "Cannot call `unsqueeze` on a persistent tensordict. Make it dense before calling this method by calling `to_tensordict`." + ) + + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]: + splits = -(self.batch_size[dim] // -chunks) + return self.split(splits, dim) + + __eq__ = TensorDict.__eq__ + __ne__ = TensorDict.__ne__ + __xor__ = TensorDict.__xor__ + __or__ = TensorDict.__or__ + __ge__ = TensorDict.__ge__ + __gt__ = TensorDict.__gt__ + __le__ = TensorDict.__le__ + __lt__ = TensorDict.__lt__ + + _apply_nest = TensorDict._apply_nest + _cast_reduction = TensorDict._cast_reduction + _check_device = TensorDict._check_device + _check_is_shared = TensorDict._check_is_shared + _convert_to_tensordict = TensorDict._convert_to_tensordict + _get_names_idx = TensorDict._get_names_idx + _index_tensordict = TensorDict._index_tensordict + _multithread_apply_flat = TensorDict._multithread_apply_flat + _multithread_rebuild = TensorDict._multithread_rebuild + _to_module = TensorDict._to_module + _unbind = TensorDict._unbind + all = TensorDict.all + any = TensorDict.any + expand = TensorDict.expand + from_dict_instance = TensorDict.from_dict_instance + masked_select = TensorDict.masked_select + _repeat = TensorDict._repeat + _repeat = TensorDict._repeat + repeat_interleave = TensorDict.repeat_interleave + reshape = TensorDict.reshape + split = TensorDict.split + + +_register_tensor_class(PersistentTensorDict) + + +def _set_max_batch_size(source: PersistentTensorDict): + """Updates a tensordict with its maximium batch size.""" + tensor_data = list(source._items_metadata()) + for key, val in tensor_data: + if not val.get("non_tensor", None) and not val.get("array", None): + _set_max_batch_size(source.get(key, None)) + + batch_size = [] + if not tensor_data: # when source is empty + source.batch_size = batch_size + return + + curr_dim = 0 + # We need to reload this list because the value have changed + tensor_data = list(source._items_metadata()) + tensor_keys, tensor_data = zip(*tensor_data) + # Filter out the non-tensor data + tensor_data = [data for data in tensor_data if not data.get("non_tensor")] + while True: + if tensor_data[0]["dim"] > curr_dim: + curr_dim_size = tensor_data[0]["shape"][curr_dim] + else: + source.batch_size = batch_size + return + for tensor in tensor_data[1:]: + if tensor["dim"] <= curr_dim or tensor["shape"][curr_dim] != curr_dim_size: + source.batch_size = batch_size + return + batch_size.append(curr_dim_size) + curr_dim += 1 + + +def _is_non_tensor_h5(val): + import h5py + + dt = val.dtype + if ( + h5py.check_string_dtype(dt) + or h5py.check_vlen_dtype(dt) + or h5py.check_enum_dtype(dt) + or h5py.check_opaque_dtype(dt) + ): + return True + return False diff --git a/lib/python3.12/site-packages/tensordict/tensorclass.py b/lib/python3.12/site-packages/tensordict/tensorclass.py new file mode 100644 index 0000000000000000000000000000000000000000..f209dc0647518e94f0f4a9a4a16ed5a924be85eb --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/tensorclass.py @@ -0,0 +1,4781 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import abc +import concurrent +import ctypes + +import dataclasses +import functools +import inspect + +import multiprocessing.managers +import multiprocessing.sharedctypes +import numbers +import os +import pickle +import shutil +import sys +import warnings +from copy import copy, deepcopy +from dataclasses import dataclass +from pathlib import Path +from textwrap import indent + +from typing import ( + Any, + Callable, + get_args, + get_origin, + get_type_hints, + List, + Literal, + Sequence, + Type, + TYPE_CHECKING, + TypeVar, + Union, +) + +import numpy as np + +import tensordict as tensordict_lib + +import torch +from tensordict._lazy import LazyStackedTensorDict +from tensordict._nestedkey import NestedKey +from tensordict._pytree import _register_td_node +from tensordict._td import is_tensor_collection, NO_DEFAULT, TensorDict, TensorDictBase +from tensordict._tensorcollection import TensorCollection +from tensordict._torch_func import TD_HANDLED_FUNCTIONS +from tensordict.base import ( + _ACCEPTED_CLASSES, + _GET_DEFAULTS_TO_NONE, + _is_leaf_nontensor, + _is_tensor_collection, + _register_tensor_class, + _UNSET, + CompatibleType, +) +from tensordict.utils import ( # @manual=//pytorch/tensordict:_C + _GENERIC_NESTED_ERR, + _get_shape_from_args, + _is_dataclass as is_dataclass, + _is_json_serializable, + _is_tensorclass, + _LOCK_ERROR, + _td_fields, + _TENSORCLASS_MEMO, + _unravel_key_to_tuple, + _zip_strict, + capture_non_tensor_stack, + DeviceType, + IndexType, + is_tensorclass, + KeyDependentDefaultDict, + LinkedList, + list_to_stack, + set_capture_non_tensor_stack, +) +from torch import multiprocessing as mp, Tensor +from torch.multiprocessing import Manager +from torch.utils._pytree import tree_map + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +def _identity(cls): + return cls + + +try: + from typing import dataclass_transform +except ImportError: + + def dataclass_transform(*args, **kwargs): + """No-op. + + Placeholder for dataclass_transform (python<3.11). + """ + return _identity + + +T = TypeVar("T", bound=TensorCollection) +# We use an abstract AnyType instead of Any because Any isn't recognised as a type for python < 3.10 +major, minor = sys.version_info[:2] +if (major, minor) < (3, 10): + from typing import Union # noqa + + NonType = type(None) + UnionType = type(Union) +else: + from types import NoneType, UnionType +if (major, minor) < (3, 11): + + class _AnyType: + def __subclasscheck__(self, subclass): + return False + +else: + _AnyType = Any + +_TensorTypes = ( + torch.FloatTensor, + torch.DoubleTensor, + torch.IntTensor, + torch.LongTensor, + torch.ByteTensor, + torch.BoolTensor, + torch.Tensor, # The base tensor class +) +_TENSOR_ONLY_TYPE_ERR = TypeError( + "tensor_only requires types to be Tensor, Tensor-subtrypes or None." +) +# methods where non_tensordict data should be cleared in the return value +_CLEAR_METADATA = {"all", "any"} +# torch functions where we can wrap the corresponding TensorDict version +_TD_PASS_THROUGH = { + torch.atleast_1d: True, + torch.atleast_2d: True, + torch.atleast_3d: True, + torch.broadcast_to: True, + torch.cat: True, + torch.clone: True, + torch.empty_like: True, + torch.flatten: True, + torch.flip: True, + torch.fliplr: True, + torch.flipud: True, + torch.full_like: True, + torch.gather: True, + torch.movedim: True, + torch.moveaxis: True, + torch.narrow: True, + torch.ones_like: True, + torch.permute: True, + torch.rand_like: True, + torch.randn_like: True, + torch.roll: True, + torch.rot90: True, + torch.split: True, + torch.squeeze: True, + torch.stack: True, + torch.swapaxes: True, + torch.swapdims: True, + torch.tile: True, + torch.unbind: True, + torch.unflatten: True, + torch.unsqueeze: True, + torch.zeros_like: True, + torch.autograd.grad: True, +} +# Methods to be executed from tensordict, any ref to self means 'tensorclass' +_METHOD_FROM_TD = [ + "__enter__", + "__exit__", + "dumps", + "load_", + "memmap", + "memmap_", + "memmap_like", + "memmap_refresh_", + "save", +] +# Methods to be executed from tensordict, any ref to self means 'self._tensordict', no wrap of result +_FALLBACK_METHOD_FROM_TD_NOWRAP = [ + "_check_batch_size", + "_check_device", + "_check_dim_name", + "_check_unlock", + "_default_get", + "_get_at_str", + "_get_at_tuple", + "_get_names_idx", # no wrap output + "_get_str", + "_get_tuple", + "_get_tuple_maybe_non_tensor", + "_has_names", + "_items_list", + "_maybe_names", + "_multithread_apply_flat", + "_multithread_apply_nest", + "_multithread_rebuild", # rebuild checks if self is a non tensor + "_propagate_lock", + "_propagate_unlock", + "_reduce_get_metadata", + "_set_device", + "_set_names", + "_values_list", + "batch_dims", + "batch_size", + "bytes", + "cat_tensors", + "clear_refs_for_compile_", + "data_ptr", + "depth", + "dim", + "dtype", + "entry_class", + "get_item_shape", + "get_non_tensor", + "init_remote", + "irecv", + "is_consolidated", + "is_contiguous", + "is_cpu", + "is_cuda", + "is_empty", + "is_floating_point", + "is_locked", + "is_memmap", + "is_meta", + "is_shared", + "isend", + "items", + "keys", + "make_memmap", + "make_memmap_from_tensor", + "names", + "ndim", + "ndimension", + "numel", + "numpy", + "param_count", + "pop", + "recv", + "reduce", + "requires_grad", + "saved_path", + "send", + "shape", + "size", + "sorted_keys", + "to_mds", + "to_struct_array", + "tolist", + "values", +] + +# Methods to be executed from tensordict, any ref to self means 'self._tensordict' +_FALLBACK_METHOD_FROM_TD_FORCE = [ + "__ge__", + "__gt__", + "__le__", + "__lt__", + "__ror__", +] +_FALLBACK_METHOD_FROM_TD = [ + "__abs__", + "__add__", + "__and__", + "__bool__", + "__eq__", + "__iadd__", + "__imul__", + "__invert__", + "__ipow__", + "__isub__", + "__itruediv__", + "__mod__", + "__mul__", + "__ne__", + "__neg__", + "__or__", + "__pow__", + "__radd__", + "__rand__", + "__rmul__", + "__rpow__", + "__rsub__", + "__rtruediv__", + "__rxor__", + "__sub__", + "__truediv__", + "__xor__", + "_add_batch_dim", + "_apply_nest", + "_clone", + "_clone_recurse", + "_data", + "_erase_names", # TODO: must be specialized + "_exclude", # TODO: must be specialized + "_fast_apply", + "_flatten_keys_inplace", + "_flatten_keys_outplace", + "_get_sub_tensordict", + "_grad", + "_map", + "_maybe_remove_batch_dim", + "_memmap_", + "_permute", + "_remove_batch_dim", + "_repeat", + "_select", # TODO: must be specialized + "_set_at_tuple", + "_set_tuple", + "_to_module", + "abs", + "abs_", + "acos", + "acos_", + "add", + "add_", + "addcdiv", + "addcdiv_", + "addcmul", + "addcmul_", + "all", + "amax", + "amin", + "any", + "apply", + "apply_", + "as_tensor", + "asin", + "asin_", + "atan", + "atan_", + "atleast_1d", + "atleast_2d", + "atleast_3d", + "auto_batch_size_", + "auto_device_", + "bfloat16", + "bitwise_and", + "bool", + "broadcast_to", + "cat", + "cat_from_tensordict", + "ceil", + "ceil_", + "chunk", + "clamp", + "clamp_max", + "clamp_max_", + "clamp_min", + "clamp_min_", + "clear", + "clear_device_", + "complex128", + "complex32", + "complex64", + "consolidate", + "contiguous", + "copy_", + "copy_at_", + "cos", + "cos_", + "cosh", + "cosh_", + "cpu", + "create_nested", + "cuda", + "cummax", + "cummin", + "densify", + "detach", + "detach_", + "div", + "div_", + "double", + "empty", + "erf", + "erf_", + "erfc", + "erfc_", + "exclude", + "exp", + "exp_", + "expand", + "expand_as", + "expm1", + "expm1_", + "extend", + "fill_", + "filter_empty_", + "filter_non_tensor_data", + "flatten", + "flatten_keys", + "flip", + "fliplr", + "flipud", + "float", + "float16", + "float32", + "float64", + "floor", + "floor_", + "frac", + "frac_", + "from_any", + "from_consolidated", + "from_dataclass", + "from_h5", + "from_list", + "from_modules", + "from_namedtuple", + "from_pytree", + "from_remote_init", + "from_struct_array", + "from_tuple", + "fromkeys", + "gather", + "gather_and_stack", + "half", + "int", + "int16", + "int32", + "int64", + "int8", + "isfinite", + "isnan", + "isneginf", + "isposinf", + "isreal", + "lazy_stack", + "lerp", + "lerp_", + "lgamma", + "lgamma_", + "load_memmap_", + "lock_", + "log", + "log10", + "log10_", + "log1p", + "log1p_", + "log2", + "log2_", + "log_", + "logical_and", + "logsumexp", + "make_memmap_from_storage", + "map", + "map_iter", + "masked_fill", + "masked_fill_", + "masked_select", + "max", + "maximum", + "maximum_", + "maybe_dense_stack", + "mean", + "min", + "minimum", + "minimum_", + "mod", + "moveaxis", + "movedim", + "mul", + "mul_", + "named_apply", + "nanmean", + "nansum", + "narrow", + "neg", + "neg_", + "new_empty", + "new_full", + "new_ones", + "new_tensor", + "new_zeros", + "norm", + "permute", + "pin_memory", + "pin_memory_", + "popitem", + "pow", + "pow_", + "prod", + "qint32", + "qint8", + "quantile", + "quint4x2", + "quint8", + "reciprocal", + "reciprocal_", + "record_stream", + "refine_names", + "rename", + "rename_", # TODO: must be specialized + "rename_key_", + "repeat", + "repeat_interleave", + "replace", + "requires_grad_", + "reshape", + "roll", + "rot90", + "round", + "round_", + "rsub", + "select", + "separates", + "set_", + "set_non_tensor", + "setdefault", + "sigmoid", + "sigmoid_", + "sign", + "sign_", + "sin", + "sin_", + "sinh", + "sinh_", + "softmax", + "split", + "split_keys", + "sqrt", + "sqrt_", + "squeeze", + "stack", + "stack_from_tensordict", + "stack_tensors", + "std", + "sub", + "sub_", + "sum", + "swapaxes", + "swapdims", + "tan", + "tan_", + "tanh", + "tanh_", + "tensor_split", + "tile", + "to", + "to_h5", + "to_lazystack", + "to_module", + "to_namedtuple", + "to_padded_tensor", + "to_pytree", + "transpose", + "trunc", + "trunc_", + "type", + "uint16", + "uint32", + "uint64", + "uint8", + "unflatten", + "unflatten_keys", + "unlock_", + "unsqueeze", + "var", + "view", + "where", + "zero_", + "zero_grad", +] + +# These methods require a copy of the non tensor data +_FALLBACK_METHOD_FROM_TD_COPY = [ + "_clone", # TODO: must be specialized + "clone", # TODO: must be specialized + "copy", # TODO: must be specialized +] + + +def is_non_tensor(obj) -> bool: + """A local implementation of is_non_tensor. + + The utils implementation does an attribute check, but here we have access to the classes + which is more immediate. + + """ + return isinstance(obj, (NonTensorDataBase, NonTensorStack)) + + +class _tensorclass_dec: + autocast: bool + frozen: bool + nocast: bool + shadow: bool + tensor_only: bool + + def __new__( + cls, + autocast: bool = False, + frozen: bool = False, + nocast: bool = False, + shadow: bool = False, + tensor_only: bool = False, + ): + if not isinstance(autocast, bool): + clz = autocast + self = super().__new__(cls) + self.__init__( + autocast=False, + frozen=False, + nocast=False, + shadow=False, + tensor_only=False, + ) + return self.__call__(clz) + return super().__new__(cls) + + def __init__( + self, + autocast: bool = False, + frozen: bool = False, + nocast: bool = False, + shadow: bool = False, + tensor_only: bool = False, + ): + if autocast and nocast: + raise TypeError("autocast is exclusive with nocast.") + if tensor_only and (autocast or nocast): + raise TypeError( + "tensor_only and autocast or nocast are exclusive features." + ) + self.autocast = autocast + self.frozen = frozen + self.nocast = nocast + self.shadow = shadow + self.tensor_only = tensor_only + + @dataclass_transform() + def __call__(self, cls: T) -> T: + clz = _tensorclass( + cls, frozen=self.frozen, shadow=self.shadow, tensor_only=self.tensor_only + ) + clz._autocast = self.autocast + clz._nocast = self.nocast + clz._shadow = self.shadow + clz._frozen = self.frozen + clz._tensor_only = self.tensor_only + return clz + + +def from_dataclass( + obj: Any, + *, + dest_cls: Type | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + batch_size: torch.Size | None = None, + frozen: bool = False, + autocast: bool = False, + nocast: bool = False, + inplace: bool = False, + shadow: bool = False, + tensor_only: bool = False, + device: torch.device | None = None, +) -> Any: + """Converts a dataclass instance or a type into a tensorclass instance or type, respectively. + + This function takes a dataclass instance or a dataclass type and converts it into a tensor-compatible class, + optionally applying various configurations such as auto-batching, immutability, and type casting. + + Args: + obj (Any): The dataclass instance or type to be converted. If a type is provided, a new class is returned. + + Keyword Args: + dest_cls (tensorclass, optional): A tensorclass type to be used to map the data. If not provided, a new + class is created. Without effect if :attr:`obj` is a type. + auto_batch_size (bool, optional): If ``True``, automatically determines and applies batch size to the resulting object. Defaults to ``False``. + batch_dims (int, optional): If auto_batch_size is ``True``, defines how many dimensions the output tensordict should have. Defaults to ``None`` (full batch-size at each level). + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to ``None``. + frozen (bool, optional): If ``True``, the resulting class or instance will be immutable. Defaults to ``False``. + autocast (bool, optional): If ``True``, enables automatic type casting for the resulting class or instance. Defaults to ``False``. + nocast (bool, optional): If ``True``, disables any type casting for the resulting class or instance. Defaults to ``False``. + tensor_only (bool, optional): if ``True``, it is expected that all items in tensorclass will be + tensor instances (tensor-compatible, since non-tensor data is converted to tensors if possible). + This can bring significant speed-ups at the cost of flexible interactions with non-tensor data. + Defaults to ``False``. + inplace (bool, optional): If ``True``, the dataclass type passed will be modified in-place. Defaults to ``False``. + Without effect if an instance is provided. + device (torch.device, optional): The device on which the TensorDict will be created. Defaults to ``None``. + shadow (bool, optional): Disables the validation of field names against TensorDict's reserved attributes. + Use with caution, as this may cause unintended consequences. Defaults to False. + + Returns: + A tensor-compatible class or instance derived from the provided dataclass. + + Raises: + TypeError: If the provided input is not a dataclass instance or type. + + Examples: + >>> from dataclasses import dataclass + >>> import torch + >>> from tensordict.tensorclass import from_dataclass + >>> + >>> @dataclass + >>> class X: + ... a: int + ... b: torch.Tensor + ... + >>> x = X(0, 0) + >>> x2 = from_dataclass(x) + >>> print(x2) + X( + a=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + b=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> X2 = from_dataclass(X, autocast=True) + >>> print(X2(a=0, b=0)) + X( + a=NonTensorData(data=0, batch_size=torch.Size([]), device=None), + b=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), + batch_size=torch.Size([]), + device=None, + is_shared=False) + + .. note:: + If a dataclass type is provided, a new class is returned with the specified configurations. + If a dataclass instance is provided, a new instance of the tensor-compatible class is returned. + The `auto_batch_size`, `frozen`, `autocast`, and `nocast` options allow for flexible configuration of the resulting class or instance. + + .. warning:: Whereas :meth:`~tensordict.TensorDict.from_dataclass` will return a :class:`~tensordict.TensorDict` instance + by default, this method will return a tensorclass instance or type. + + """ + from dataclasses import asdict, make_dataclass + + if isinstance(obj, type): + if is_tensorclass(obj): + return obj + if not inplace: + cls = make_dataclass( + obj.__name__ + "_tc", + fields=obj.__dataclass_fields__, + bases=obj.__bases__, + ) + else: + cls = obj + clz = _tensorclass(cls, frozen=frozen, shadow=shadow, tensor_only=tensor_only) + clz._type_hints = get_type_hints(obj) + clz._autocast = autocast + clz._nocast = nocast + clz._shadow = shadow + clz._frozen = frozen + clz._tensor_only = tensor_only + return clz + + if not is_dataclass(obj): + raise TypeError(f"Expected a obj input, got a {type(obj)} input instead.") + name = obj.__class__.__name__ + "_tc" + if dest_cls is None: + if tensor_only and (autocast or nocast): + raise TypeError( + "tensor_only and autocast or nocast are exclusive features." + ) + clz = _tensorclass( + make_dataclass(name, fields=obj.__dataclass_fields__), + frozen=frozen, + shadow=shadow, + tensor_only=tensor_only, + ) + clz._autocast = autocast + clz._nocast = nocast + clz._shadow = shadow + clz._frozen = frozen + clz._tensor_only = tensor_only + else: + clz = dest_cls + result = clz(**asdict(obj), batch_size=batch_size, device=device) + if auto_batch_size: + if batch_size is not None: + raise TypeError( + TensorDictBase._CONFLICTING_BATCH_SIZES.format("from_dataclass") + ) + result = result.auto_batch_size_(batch_dims=batch_dims) + return result + + +@dataclass_transform() +def tensorclass( + cls: T = None, + /, + *, + autocast: bool = False, + frozen: bool = False, + nocast: bool = False, + shadow: bool = False, + tensor_only: bool = False, +) -> T | None: + """A decorator to create :obj:`tensorclass` classes. + + ``tensorclass`` classes are specialized :func:`dataclasses.dataclass` instances that + can execute some pre-defined tensor operations out of the box, such as + indexing, item assignment, reshaping, casting to device or storage and many + others. + + Keyword Args: + autocast (bool, optional): if ``True``, the types indicated will be enforced when an argument is set. + This argument is exclusive with ``nocast`` (both cannot be true at the same time). Defaults to ``False``. + frozen (bool, optional): if ``True``, the content of the tensorclass cannot be modified. This argument is + provided to dataclass-compatibility, a similar behavior can be obtained through the `lock` argument in + the class constructor. Defaults to ``False``. + nocast (bool, optional): if ``True``, Tensor-compatible types such as ``int``, ``np.ndarray`` and the like + will not be cast to a tensor type. This argument is exclusive with ``autocast`` (both cannot be true + at the same time). Defaults to ``False``. + shadow (bool, optional): Disables the validation of field names against TensorDict's reserved attributes. + Use with caution, as this may cause unintended consequences. Defaults to False. + tensor_only (bool, optional): if ``True``, it is expected that all items in tensorclass will be + tensor instances (tensor-compatible, since non-tensor data is converted to tensors if possible). + This can bring significant speed-ups at the cost of flexible interactions with non-tensor data. + Defaults to ``False``. + + tensorclass can be used with or without arguments: + + Examples: + >>> @tensorclass + ... class X: + ... y: int + >>> X(torch.ones(())).y + tensor(1.) + >>> @tensorclass(autocast=False) + ... class X: + ... y: int + >>> X(torch.ones(())).y + tensor(1.) + >>> @tensorclass(autocast=True) + ... class X: + ... y: int + >>> X(torch.ones(())).y + 1 + >>> @tensorclass(nocast=True) + ... class X: + ... y: Any + >>> X(1).y + 1 + >>> @tensorclass(nocast=False) + ... class X: + ... y: Any + >>> X(1).y + tensor(1) + + Examples: + >>> from tensordict import tensorclass + >>> import torch + >>> from typing import Optional + >>> + >>> @tensorclass + ... class MyData: + ... X: torch.Tensor + ... y: torch.Tensor + ... z: str + ... def expand_and_mask(self): + ... X = self.X.unsqueeze(-1).expand_as(self.y) + ... X = X[self.y] + ... return X + ... + >>> data = MyData( + ... X=torch.ones(3, 4, 1), + ... y=torch.zeros(3, 4, 2, 2, dtype=torch.bool), + ... z="test" + ... batch_size=[3, 4]) + >>> print(data) + MyData( + X=Tensor(torch.Size([3, 4, 1]), dtype=torch.float32), + y=Tensor(torch.Size([3, 4, 2, 2]), dtype=torch.bool), + z="test" + batch_size=[3, 4], + device=None, + is_shared=False) + >>> print(data.expand_and_mask()) + tensor([]) + + It is also possible to nest tensorclasses instances within each other: + Examples: + >>> from tensordict import tensorclass + >>> import torch + >>> from typing import Optional + >>> + >>> @tensorclass + ... class NestingMyData: + ... nested: MyData + ... + >>> nesting_data = NestingMyData(nested=data, batch_size=[3, 4]) + >>> # although the data is stored as a TensorDict, the type hint helps us + >>> # to appropriately cast the data to the right type + >>> assert isinstance(nesting_data.nested, type(data)) + + + """ + + def wrap(cls): + return _tensorclass_dec( + autocast=autocast, + frozen=frozen, + nocast=nocast, + shadow=shadow, + tensor_only=tensor_only, + )(cls) + + # See if we're being called as @tensorclass or @tensorclass(). + if cls is None: + # We're called with parens. + return wrap + + # We're called as @tensorclass without parens. + return wrap(cls) + + +@dataclass_transform() +def _tensorclass(cls: T, *, frozen, shadow: bool, tensor_only: bool) -> T: + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + if func not in _TD_PASS_THROUGH or not all( + issubclass(t, (Tensor, cls, TensorDictBase)) for t in types + ): + return NotImplemented + + if kwargs is None: + kwargs = {} + + # get the output type from the arguments / keyword arguments + if len(args) > 0: + tensorclass_instance = args[0] + else: + tensorclass_instance = kwargs.get("input", kwargs["tensors"]) + if isinstance(tensorclass_instance, (tuple, list)): + tensorclass_instance = tensorclass_instance[0] + args = tuple(_arg_to_tensordict(arg) for arg in args) + kwargs = {key: _arg_to_tensordict(value) for key, value in kwargs.items()} + + result = TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + if isinstance(result, (list, tuple)): + return type(result)( + _from_tensordict_with_copy(tensorclass_instance, tensordict_result) + for tensordict_result in result + ) + return _from_tensordict_with_copy(tensorclass_instance, result) + + _is_non_tensor = getattr(cls, "_is_non_tensor", False) + + # Breaks some tests, don't do that: + # if not dataclasses.is_dataclass(cls): + cls = dataclass(cls, frozen=frozen) + _TENSORCLASS_MEMO[cls] = True + + # Use dataclasses.fields() to get the correct field list (excludes ClassVar fields) + # rather than __dataclass_fields__ which may incorrectly include them + expected_keys = cls.__expected_keys__ = { + field.name for field in dataclasses.fields(cls) + } + + if not shadow: + for attr in expected_keys: + if attr in dir(TensorDict) and attr not in ("_is_non_tensor", "data"): + raise AttributeError( + f"Attribute name {attr} can't be used with @tensorclass or TensorClass. To allow it, please indicate " + f"that builtin names can be overwritten by using the allow_names keyword argument (@tensorclass(shadow=True) " + f"or TensorClass['shadow']." + ) + + cls.fields = classmethod(dataclasses.fields) + for field in cls.fields(): + if hasattr(cls, field.name): + # if we have used Cls(TensorClass["shadow"]), we have a subclass of Cls(TensorClass) + # so we cannot directly delete the attribute + try: + delattr(cls, field.name) + except AttributeError: + pass + _get_type_hints(cls, tensor_only=tensor_only) + cls.__init__ = _init_wrapper(cls.__init__, cls, frozen, shadow, tensor_only) + cls._from_tensordict = classmethod(_from_tensordict) + cls.from_tensordict = cls._from_tensordict + if not hasattr(cls, "__torch_function__"): + cls.__torch_function__ = classmethod(__torch_function__) + cls.__getstate__ = _getstate + cls.__setstate__ = _setstate + + if tensor_only: + cls.__getattr__ = _getattr_tensor_only + else: + cls.__getattr__ = _getattr + + cls.__setattr_parent__ = object.__setattr__ + if "__setattr__" not in cls.__dict__: + if not tensor_only: + cls.__setattr__ = _setattr + else: + cls.__setattr__ = _setattr_tensor_only + if "__getitem__" not in cls.__dict__: + cls.__getitem__ = _getitem + if "__getitems__" not in cls.__dict__: + cls.__getitems__ = _getitem + if "__setitem__" not in cls.__dict__: + cls.__setitem__ = _setitem + if not _is_non_tensor: + cls.__repr__ = _repr + if "__len__" not in cls.__dict__: + cls.__len__ = _len + + cls.__eq__ = _eq + cls.__ne__ = _ne + cls.__or__ = _or + cls.__xor__ = _xor + cls.__bool__ = _bool + + if not hasattr(cls, "_new_unsafe"): + cls._new_unsafe = classmethod(_new_unsafe) + if not hasattr(cls, "non_tensor_items") and "non_tensor_items" not in expected_keys: + cls.non_tensor_items = _non_tensor_items + if not hasattr(cls, "set") and "set" not in expected_keys: + cls.set = _set + if not hasattr(cls, "set_at_") and "set_at_" not in expected_keys: + cls.set_at_ = _set_at_ + if not hasattr(cls, "_set_str"): + cls._set_str = _set_str + if not hasattr(cls, "_set_at_str"): + cls._set_at_str = _set_at_str + if not hasattr(cls, "del_") and "del_" not in expected_keys: + cls.del_ = _del_ + if "__delattr__" not in cls.__dict__: + cls.__delattr__ = _delattr + if not hasattr(cls, "get") and "get" not in expected_keys: + cls.get = _get + if not hasattr(cls, "get_at") and "get_at" not in expected_keys: + cls.get_at = _get_at + if not hasattr(cls, "unbind") and "unbind" not in expected_keys: + cls.unbind = _unbind + cls._unbind = _unbind + if not hasattr(cls, "state_dict") and "state_dict" not in expected_keys: + cls.state_dict = _state_dict + if not hasattr(cls, "load_state_dict") and "load_state_dict" not in expected_keys: + cls.load_state_dict = _load_state_dict + if not hasattr(cls, "_memmap_") and "_memmap_" not in expected_keys: + cls._memmap_ = _memmap_ + if not hasattr(cls, "share_memory_") and "share_memory_" not in expected_keys: + cls.share_memory_ = _share_memory_ + if not hasattr(cls, "update") and "update" not in expected_keys: + cls.update = _update + if not hasattr(cls, "update_") and "update_" not in expected_keys: + cls.update_ = _update_ + if not hasattr(cls, "update_at_") and "update_at_" not in expected_keys: + cls.update_at_ = _update_at_ + for method_name in _METHOD_FROM_TD: + if not hasattr(cls, method_name): + setattr(cls, method_name, getattr(TensorDict, method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD: + if not hasattr(cls, method_name): + setattr(cls, method_name, _wrap_td_method(method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD_FORCE: + setattr(cls, method_name, _wrap_td_method(method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD_NOWRAP: + if not hasattr(cls, method_name) and method_name not in expected_keys: + is_property = isinstance( + getattr(TensorDictBase, method_name, None), property + ) + setattr( + cls, + method_name, + _wrap_td_method(method_name, no_wrap=True, is_property=is_property), + ) + for method_name in _FALLBACK_METHOD_FROM_TD_COPY: + if not hasattr(cls, method_name): + setattr( + cls, + method_name, + _wrap_td_method(method_name, copy_non_tensor=True), + ) + + # if not hasattr(cls, "batch_size") and "batch_size" not in expected_keys: + # cls.batch_size = property(_batch_size, _batch_size_setter) + + # Memmap + if not hasattr(cls, "load_memmap") and "load_memmap" not in expected_keys: + cls.load_memmap = TensorDictBase.load_memmap + if not hasattr(cls, "load") and "load" not in expected_keys: + cls.load = TensorDictBase.load + if not hasattr(cls, "_load_memmap"): + cls._load_memmap = classmethod(_load_memmap) + if not hasattr(cls, "from_dict") and "from_dict" not in expected_keys: + cls.from_dict = classmethod(_from_dict) + if ( + not hasattr(cls, "from_dict_instance") + and "from_dict_instance" not in expected_keys + ): + cls.from_dict_instance = _from_dict_instance + + for attr in set(TensorDict.__dict__.keys()).union(TensorDictBase.__dict__.keys()): + if attr in ("__torch_function__",): + continue + func = getattr(TensorDict, attr) + if inspect.ismethod(func) and attr not in cls.__dict__: + tdcls = func.__self__ + if issubclass(tdcls, TensorDictBase): # detects classmethods + setattr(cls, attr, _wrap_classmethod(tdcls, cls, func)) + + if not hasattr(cls, "to_tensordict") and "to_tensordict" not in expected_keys: + cls.to_tensordict = _to_tensordict + if not hasattr(cls, "device") and "device" not in expected_keys: + cls.device = property(_device, _device_setter) + if not _is_non_tensor and not hasattr(cls, "data") and "data" not in expected_keys: + cls.data = property(_data, _data_setter) + if not hasattr(cls, "grad") and "grad" not in expected_keys: + cls.grad = property(_grad) + if not hasattr(cls, "to_dict") and "to_dict" not in expected_keys: + cls.to_dict = _to_dict + + cls.__doc__ = f"{cls.__name__}{inspect.signature(cls)}" + + _register_tensor_class(cls) + try: + _register_td_node(cls) + except ValueError: + # The class may already be registered as a pytree node + pass + + # faster than doing instance checks + cls._is_non_tensor = _is_non_tensor + cls._is_tensorclass = True + + from tensordict import _pytree + + _pytree._CONSTRUCTORS[cls] = _pytree._tensorclass_constructor + return cls + + +# def _batch_size(self): +# return self.__dict__["_tensordict"]._batch_size +# def _batch_size_setter(self, value): +# self.__dict__["_tensordict"].batch_size = value + + +def _arg_to_tensordict(arg): + # if arg is a tensorclass or sequence of tensorclasses, extract the underlying + # tensordicts and return those instead + + # since arg can be anything (e.g. callable etc) we can't use pytree + # def convert(x): + # if _is_tensorclass(type(x)): + # return x._tensordict + # return x + # return torch.utils._pytree.tree_map(convert, arg) + + if _is_tensorclass(type(arg)): + return arg._tensordict + elif isinstance(arg, (tuple, list)) and all( + _is_tensorclass(type(item)) for item in arg + ): + return type(arg)(item._tensordict for item in arg) + return arg + + +def _from_tensordict_with_copy(tc, tensordict): + # creates a new tensorclass with the same type as tc, and a copy of the + # non_tensordict data + return type(tc)._from_tensordict( + tensordict=tensordict, non_tensordict=dict(tc._non_tensordict) + ) + + +def _from_tensordict_with_none(tc, tensordict): + # creates a new tensorclass with the same type as tc, and all non_tensordict entries + # set to None + return type(tc)._from_tensordict( + tensordict=tensordict, + non_tensordict={key: None for key in tc._non_tensordict}, + ) + + +def _init_wrapper( + __init__: Callable, cls, frozen: bool, shadow: bool, tensor_only: bool +) -> Callable: + init_sig = inspect.signature(__init__) + params = list(init_sig.parameters.values()) + # drop first entry of params which corresponds to self and isn't passed by the user + required_params = [p.name for p in params[1:] if p.default is inspect._empty] + # if not required_params and hasattr(cls, "__init_parent__"): + # init_sig_parent = inspect.signature(cls.__init_parent__) + # params_parent = list(init_sig_parent.parameters.values()) + # # drop first entry of params which corresponds to self and isn't passed by the user + # required_params = [p.name for p in params_parent[1:] if p.default is inspect._empty] + + @functools.wraps(__init__) + def wrapper( + self, + *args: Any, + **kwargs, + ): + if "batch_size" in required_params: + batch_size = torch.Size(()) + else: + batch_size = kwargs.pop("batch_size", torch.Size(())) + if isinstance(batch_size, int): + batch_size = (batch_size,) + elif batch_size is None: + batch_size = torch.Size(()) + + if "names" in required_params: + names = None + else: + names = kwargs.pop("names", None) + if "device" in required_params: + device = None + else: + device = kwargs.pop("device", None) + if "lock" in required_params: + lock = None + else: + lock = kwargs.pop("lock", None) + if lock is None: + lock = frozen + if not is_compiling(): + # zip not supported by dynamo + # Use __dataclass_fields__ but filter out ClassVar fields to preserve order + expected_keys_list = [ + key + for key in type(self).__dataclass_fields__ + if key in self.__expected_keys__ + ] + + # Check that we don't have too many positional arguments + if len(args) > len(expected_keys_list): + raise TypeError( + f"{type(self).__name__}.__init__() takes {len(expected_keys_list) + 1} positional arguments but {len(args) + 1} were given" + ) + + for value, key in zip(args, expected_keys_list): + if key in kwargs: + raise ValueError(f"The key {key} is already set in kwargs") + kwargs[key] = value + else: + if args: + raise RuntimeError( + "dynamo doesn't support arguments when building a tensorclass, pass the keyword explicitly." + ) + + if not is_compiling(): + for key, field in type(self).__dataclass_fields__.items(): + # Only process fields that are in __expected_keys__ (excludes ClassVar fields) + if key in self.__expected_keys__: + if field.default_factory not in ( + dataclasses.MISSING, + ) and not isinstance( + field.default_factory, + getattr( + dataclasses, "_MISSING_TYPE", type(dataclasses.MISSING) + ), + ): + default = field.default_factory() + else: + default = field.default + if default not in (None, dataclasses.MISSING): + kwargs.setdefault(key, default) + else: + # TODO: Decide what to do here + pass + + missing_params = [p for p in required_params if p not in kwargs] + if missing_params: + n_missing = len(missing_params) + raise TypeError( + f"{type(self).__name__}.__init__() missing {n_missing} " + f"required positional argument{'' if n_missing == 1 else 's'}: " + f"""{", ".join(f"'{name}'" for name in missing_params)}""" + ) + + super(type(self), self).__setattr__( + "_tensordict", + TensorDict._new_unsafe( + {}, + batch_size=torch.Size(batch_size), + device=device, + names=names, + ), + ) + # super(type(self), self).__setattr__("_non_tensordict", {}) + # super(type(self), self).__setattr__("_is_initialized", True) + self.__setattr_parent__("_non_tensordict", {}) + self.__setattr_parent__("_is_initialized", True) + + # convert the non tensor data in a regular data + __init__(self, **kwargs) + if frozen: + local_setattr = _setattr + for key, val in kwargs.items(): + local_setattr(self, key, val) + del self.__dict__[key] + if lock: + self._tensordict.lock_() + + if not shadow: + new_params = [ + inspect.Parameter("batch_size", inspect.Parameter.KEYWORD_ONLY), + inspect.Parameter("device", inspect.Parameter.KEYWORD_ONLY, default=None), + inspect.Parameter("names", inspect.Parameter.KEYWORD_ONLY, default=None), + ] + else: + new_params = [] + for p in params: + if p._name == "batch_size": + break + else: + new_params.append( + inspect.Parameter("batch_size", inspect.Parameter.KEYWORD_ONLY) + ) + for p in params: + if p._name == "device": + break + else: + new_params.append( + inspect.Parameter( + "device", inspect.Parameter.KEYWORD_ONLY, default=None + ) + ) + for p in params: + if p._name == "names": + break + else: + new_params.append( + inspect.Parameter("names", inspect.Parameter.KEYWORD_ONLY, default=None) + ) + wrapper.__signature__ = init_sig.replace(parameters=params + new_params) + + return wrapper + + +_cast_funcs = KeyDependentDefaultDict(_identity) +_cast_funcs[torch.Tensor] = torch.as_tensor +_cast_funcs[np.ndarray] = np.asarray + + +def _new_unsafe(cls, *args, **kwargs) -> T: + return _from_tensordict(cls, TensorDict._new_unsafe(*args, **kwargs)) + + +def _get_type_hints(cls, with_locals=False, tensor_only=False): + ####### + # Set proper type annotations for autocasting to tensordict/tensorclass + # + # by updating locals, we can allow this to be used within a function + # local-cross referencing will not work though + # def foo(): + # @tensorclass + # class MyOtherClass: + # x: torch.Tensor + # @tensorclass + # class MyClass: + # x: MyClass # works + # y: MyOtherClass # fails + # + # In this case, we will use the get_parent_local function to get the locals + # from the parent frame and so recursively until we can find the class. + + if with_locals: + # This function gets the parent frame recursively until we can find the current class. + # Any exception leads to this to be None and auto-casting will be disabled + localns = locals() + localns = copy(localns) + + def get_parent_locals(cls, localns=localns): + # Get the current frame + frame = inspect.currentframe() + try: + parent_locs = localns + while cls.__name__ not in parent_locs: + # Get the parent frame + parent_frame = frame.f_back + # Get the locals dictionary of the parent frame + parent_locs = parent_frame.f_locals + frame = parent_frame + except Exception: + localns.setdefault(cls.__name__, cls) + return localns + finally: + # Clean up the frame reference + del frame + return copy(parent_locs) + + localns = get_parent_locals(cls) + else: + localns = None + + globalns = None + + try: + cls._type_hints = get_type_hints( + cls, + localns=localns, + # globalns=globals(), + ) + if tensor_only: + + def is_tensor_or_optional_tensor(type_hint): + # Check if the type hint is exactly torch.Tensor + if isinstance(type_hint, type): + return issubclass(type_hint, _TensorTypes) or _is_tensor_collection( + type_hint + ) + if isinstance(type_hint, type(Any)): + return False + if isinstance(type_hint, UnionType): + args = get_args(type_hint) + return all( + t is None or t is NoneType or is_tensor_or_optional_tensor(t) + for t in args + ) + # Check if the type hint is a Union (e.g., Tensor | None) + origin = get_origin(type_hint) + + if origin is Union: + args = get_args(type_hint) + return all( + t is None or t is NoneType or is_tensor_or_optional_tensor(t) + for t in args + ) + return False + + for key, val in cls._type_hints.items(): + if key not in cls.__expected_keys__: + continue + if not is_tensor_or_optional_tensor(val): + raise _TENSOR_ONLY_TYPE_ERR + cls._type_hints = { + key: val if isinstance(val, type) else _AnyType + for key, val in cls._type_hints.items() + } + except NameError: + if not with_locals: + return _get_type_hints(cls, with_locals=True, tensor_only=tensor_only) + cls._set_dict_warn_msg = ( + "A NameError occurred while trying to retrieve a type annotation. " + "This can occur when a tensorclass references another locally defined " + "tensorclass. " + f"As a result type hints cannot be read and {cls}.from_dict(...) " + f"or `{cls}.set` will not attempt to map dictionaries to " + "the relevant tensorclass. To resolve this issue, consider defining " + "your tensorclass globally." + ) + cls._type_hints = None + except TypeError as err: + if err is _TENSOR_ONLY_TYPE_ERR: + raise err + # This is a rather common case where type annotation is like + # class MyClass: + # x: int | str + # in which case get_type_hints doesn't work (it does work + # however with old-school Optional or Union...) + # We simply differ the warning till _set() is called + cls._set_dict_warn_msg = ( + "A TypeError occurred when trying to retrieve a type annotation. " + "This may be caused by annotations that use plain `|` instead of typing.Union " + "or typing.Optional which are supported. If you wish to use the feature " + "of setting dict as attributes with automapping to tensordict/tensorclass " + "(`my_obj.attr = dict(...)`), consider re-writing the tensorclass with " + "traditional type annotations." + ) + cls._type_hints = None + + +def _from_tensordict( + cls, + tensordict: TensorDictBase, + non_tensordict: dict | None = None, + safe: bool = True, +) -> Self: # noqa: D417 + """Tensor class wrapper to instantiate a new tensor class object. + + Args: + tensordict (TensorDictBase): Dictionary of tensor types + non_tensordict (dict): Dictionary with non-tensor and nested tensor class objects + safe (bool): Whether to raise an error if the tensordict is not a TensorDictBase instance + + """ + if safe and not isinstance(tensordict, TensorDictBase): + raise RuntimeError( + f"Expected a TensorDictBase instance but got {type(tensordict)}" + ) + # Validating keys of tensordict + # tensordict = tensordict.copy() + tensor_keys = tensordict.keys() + # TODO: compile doesn't like set() over an arbitrary object + if is_compiling(): + tensor_keys = {k for k in tensor_keys} # noqa: C416 + exp_keys = {k for k in cls.__expected_keys__} # noqa: C416 + if non_tensordict is not None: + nontensor_keys = {k for k in non_tensordict.keys()} # noqa: C416 + else: + nontensor_keys = set() + non_tensordict = {} + # TODO: Makes compile unhappy + # total_keys = tensor_keys.union(nontensor_keys) + total_keys = set(tensor_keys) + total_keys.update(nontensor_keys) + else: + tensor_keys = set(tensor_keys) + exp_keys = set(cls.__expected_keys__) + if non_tensordict is not None: + nontensor_keys = set(non_tensordict.keys()) + total_keys = tensor_keys.union(nontensor_keys) + else: + nontensor_keys = set() + non_tensordict = {} + total_keys = tensor_keys + for key in nontensor_keys: + if key not in tensor_keys: + continue + if non_tensordict[key] is None: + del non_tensordict[key] + continue + raise KeyError(f"{key} is present in both tensor and non-tensor dicts.") + if total_keys - exp_keys: + raise ValueError( + f"Keys from the tensordict ({set(tensordict.keys())}) must " + f"correspond to the class attributes ({cls.__expected_keys__}). Got the set of " + f"keys {{{total_keys - exp_keys}}} which do not belong to the class." + ) + else: + to_add = exp_keys - total_keys + for key in to_add: + non_tensordict[key] = None + + if not is_compiling(): + # bypass initialisation. this means we don't incur any overhead creating an + # empty tensordict and writing values to it. we can skip this because we already + # have a tensordict to use as the underlying tensordict + tc = cls.__new__(cls) + tc.__dict__.update( + {"_tensordict": tensordict, "_non_tensordict": non_tensordict} + ) + # since we aren't calling the dataclass init method, we need to manually check + # whether a __post_init__ method has been defined and invoke it if so + if hasattr(cls, "__post_init__"): + tc.__post_init__() + return tc + else: + # TODO: things that did NOT work: **tensordict, dict(tensordict) + kwargs = dict(tensordict.items()) + kwargs.update(non_tensordict) + kwargs["batch_size"] = tensordict.batch_size + kwargs["device"] = tensordict.device + kwargs["names"] = tensordict._maybe_names() + return cls(**kwargs) + + +def _memmap_( + self, + *, + prefix: str | None = None, + copy_existing: bool = False, + executor=None, + futures=None, + inplace=True, + like=False, + memmaped: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key, +): + _non_tensordict = dict(self._non_tensordict) + cls = type(self) + + if not memmaped and prefix is not None: + prefix = Path(prefix) + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + + def save_metadata(cls=cls, _non_tensordict=_non_tensordict, prefix=prefix): + with open(prefix / "meta.json", "wb") as f: + metadata = {"_type": str(cls)} + to_pickle = {} + for key, value in _non_tensordict.items(): + value = _from_shared_nontensor(value) + if _is_json_serializable(value): + metadata[key] = value + else: + to_pickle[key] = value + from tensordict.utils import json_dumps + + json_str = json_dumps(metadata) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + f.write(json_str.encode("utf-8")) + else: + f.write(json_str) + if to_pickle: + with open(prefix / "other.pickle", "wb") as pickle_file: + pickle.dump(to_pickle, pickle_file) + + if executor is None: + save_metadata() + else: + futures.append(executor.submit(save_metadata)) + + prefix = prefix / "_tensordict" + new_futures = [] + if not isinstance(self, NonTensorDataBase): + # TODO: We can't execute this using multiple threads because from_tensordict expects + # the tensordict and non_tensordict to be complete + td = self._tensordict._memmap_( + prefix=prefix, + # executor=None, + # futures=[], + executor=executor, + futures=new_futures, + inplace=inplace, + like=like, + copy_existing=copy_existing, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + if new_futures: + futures += new_futures + td._device = torch.device("cpu") + else: + # For non-tensor data, we don't create an empty _tensordict dir + td = self._tensordict.empty() + td._is_memmap = True + td._is_locked = True + td._memmap_prefix = prefix + if inplace: + self.__dict__["_tensordict"] = td + if not inplace: + if new_futures: + concurrent.futures.wait(new_futures) + result = cls._from_tensordict(td, _non_tensordict) + else: + result = self + return result + + +def _share_memory_(self): + self._tensordict.share_memory_() + return self + + +def _load_memmap(cls, prefix: Path, metadata: dict, *, robust_key, **kwargs): + non_tensordict = dict(metadata) + del non_tensordict["_type"] + if os.path.exists(prefix / "other.pickle"): + with open(prefix / "other.pickle", "rb") as pickle_file: + non_tensordict.update(pickle.load(pickle_file)) + if os.path.exists(prefix / "_tensordict"): + td = TensorDict.load_memmap( + prefix / "_tensordict", **kwargs, non_blocking=False, robust_key=robust_key + ) + else: + if not issubclass(cls, NonTensorDataBase): + raise ValueError("The _tensordict directory seems to be missing.") + td = TensorDict(device="cpu") + return cls._from_tensordict(td, non_tensordict) + + +def _getstate(self) -> dict[str, Any]: + """Returns a state dict which consists of tensor and non_tensor dicts for serialization. + + Returns: + dictionary of state of tensor class + + """ + return {"tensordict": self._tensordict, "non_tensordict": self._non_tensordict} + + +def _setstate(self, state: dict[str, Any]) -> None: # noqa: D417 + """Used to set the state of an object using state parameter. + + Args: + state (dict): State parameter to set the object + """ + self._tensordict = state.get("tensordict") + self._non_tensordict = state.get("non_tensordict") + + +def _getattr_tensor_only(self, item: str, **kwargs) -> Any: + # Use _UNSET sentinel instead of try/except for torch.compile compatibility + out = self._tensordict._get_str(item, _UNSET, **kwargs) + if out is not _UNSET: + return out + out = self._non_tensordict.get(item, _UNSET) + if out is not _UNSET: + return out + out = getattr(self._tensordict, item, NO_DEFAULT) + if out is not NO_DEFAULT: + if not callable(out) and not is_non_tensor(out): + return out + if is_non_tensor(out): + return out.data if hasattr(out, "data") else out.tolist(as_linked_list=True) + return _wrap_method(self, item, out) + raise AttributeError(item) + + +def _getattr(self, item: str, **kwargs) -> Any: + __dataclass_fields__ = type(self).__expected_keys__ + + if item in __dataclass_fields__: + _non_tensordict = self._non_tensordict + if _non_tensordict and item not in self._tensordict.keys(): + out = _non_tensordict.get(item, NO_DEFAULT) + if out is not NO_DEFAULT: + if ( + isinstance(self, NonTensorDataBase) + and item == "data" + and (self._is_shared or self._is_memmap) + ): + return _from_shared_nontensor(out) + return out + out = self._tensordict._get_str(item, NO_DEFAULT, **kwargs) + if is_non_tensor(out): + return ( + out.data + if not isinstance(out, NonTensorStack) + else out.tolist(as_linked_list=True) + ) + return out + + out = getattr(self._tensordict, item, NO_DEFAULT) + if out is not NO_DEFAULT: + if not callable(out) and not is_non_tensor(out): + return out + if is_non_tensor(out): + return out.data if hasattr(out, "data") else out.tolist(as_linked_list=True) + return _wrap_method(self, item, out) + raise AttributeError(item) + + +SET_ATTRIBUTES = ( + "batch_size", + "device", + "_locked_tensordicts", + "names", + "_is_initialized", +) + + +def _setattr(self, key: str, value: Any) -> None: # noqa: D417 + __dict__ = self.__dict__ + if ( + "_tensordict" not in __dict__ + or "_non_tensordict" not in __dict__ + or ( + not self._shadow + and ( + key in SET_ATTRIBUTES + or ( + key in type(self).__dict__ + and key not in getattr(self, "__expected_keys__", set()) + ) + ) + ) + ): + # if we ever decide to allow anything to be written in a tc + # or key not in self.__dataclass_fields__): + return self.__setattr_parent__(key, value) + + if key not in self.__expected_keys__: + raise AttributeError( + f"Cannot set the attribute {key} in {self} as this entry is not amongst the expected ones ({self.__expected_keys__})." + ) + out = self.set(key, value) + if out is not self: + raise RuntimeError( + "Cannot set the attribute on a locked tensorclass, even if " + "clone_on_set is set to True. Use my_obj.set(...) instead." + ) + + +def _setattr_tensor_only(self, key: str, value: Any) -> None: # noqa: D417 + if not is_compiling(): + __dict__ = self.__dict__ + if ( + "_tensordict" not in __dict__ + or "_non_tensordict" not in __dict__ + or ( + not self._shadow + and (key in SET_ATTRIBUTES or key in type(self).__dict__) + ) + ): + return self.__setattr_parent__(key, value) + else: + # Pass? + if key in SET_ATTRIBUTES: + return self.__setattr_parent__(key, value) + if key not in self.__expected_keys__: + raise AttributeError( + f"Cannot set attribute {key} in {self} as this entry is not amongst the expected ones ({self.__expected_keys__})." + ) + if value is None: + self._non_tensordict[key] = None + return + out = self._set_str(key, value, inplace=False, validated=False, ignore_lock=False) + if out is not self: + raise RuntimeError( + "Cannot set attribute on a locked tensorclass, even if " + "clone_on_set is set to True. Use my_obj.set(...) instead." + ) + + +def _wrap_td_method( + funcname, *, copy_non_tensor=False, no_wrap=False, is_property=False +): + def deliver_result(self, result, kwargs): + if result is None: + return + if isinstance(result, TensorDictBase) and kwargs.get("out") is not result: + if not is_compiling(): + non_tensordict = super(type(self), self).__getattribute__( + "_non_tensordict" + ) + else: + non_tensordict = self._non_tensordict + non_tensordict = dict(non_tensordict) + if copy_non_tensor and non_tensordict: + # use tree_map to copy + non_tensordict = tree_map(_identity, non_tensordict) + return self._from_tensordict(result, non_tensordict, safe=False) + return result + + if not is_property: + + def wrapped_func(self, *args, **kwargs): + if not is_compiling(): + td = super(type(self), self).__getattribute__("_tensordict") + else: + td = self._tensordict + result = getattr(td, funcname)(*args, **kwargs) + if no_wrap: + return result + + if result is td: + return self + + if isinstance(result, tuple): + return tuple(deliver_result(self, r, kwargs) for r in result) + return deliver_result(self, result, kwargs) + + return wrapped_func + + def wrapped_func(self): + if not is_compiling(): + td = super(type(self), self).__getattribute__("_tensordict") + else: + td = self._tensordict + result = getattr(td, funcname) + + if no_wrap: + return result + + if result is td: + return self + + if isinstance(result, tuple): + return tuple(deliver_result(self, r, {}) for r in result) + return deliver_result(self, result, {}) + + def wrapped_func_setter(self, value): + if not is_compiling(): + td = super(type(self), self).__getattribute__("_tensordict") + else: + td = self._tensordict + return setattr(td, funcname, value) + + return property(wrapped_func, wrapped_func_setter) + + +def _wrap_method(self, attr, func, nowarn=False): + if not nowarn: + warnings.warn( + f"The method {func} wasn't explicitly implemented for tensorclass. " + f"This fallback will be deprecated in future releases because it is inefficient " + f"and non-compilable. Please raise an issue in tensordict repo to support this method!" + ) + + @functools.wraps(func) + def wrapped_func(*args, **kwargs): + args = tuple(_arg_to_tensordict(arg) for arg in args) + kwargs = {key: _arg_to_tensordict(value) for key, value in kwargs.items()} + res = func(*args, **kwargs) + if isinstance(res, TensorDictBase): + if attr.endswith("_"): + # in-place operation, return the current object + return self + elif attr in _CLEAR_METADATA: + # this is an attribute where copying the metadata makes no sense, e.g. + # .all or .any, so we replace all values with None + return type(self)._from_tensordict( + res, {k: None for k in self._non_tensordict} + ) + # create a new tensorclass from res and copy the metadata from self + return type(self)._from_tensordict(res, dict(self._non_tensordict)) + return res + + if not is_compiling(): + wrapped_func = functools.wraps(func)(wrapped_func) + + return wrapped_func + + +def _update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + non_blocking: bool = False, + update_batch_size: bool = False, + ignore_lock: bool = False, + is_leaf: Callable[[Type], bool] | None = None, +): + if is_leaf is None: + is_leaf = _is_leaf_nontensor + if isinstance(input_dict_or_td, dict): + input_dict_or_td = self.from_dict(input_dict_or_td, auto_batch_size=False) + + if is_tensorclass(input_dict_or_td): + non_tensordict = { + k: v + for k, v in input_dict_or_td.__dict__["_non_tensordict"].items() + if v is not None + } + self._tensordict.update( + input_dict_or_td.__dict__["_tensordict"], + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + is_leaf=is_leaf, + ) + self._non_tensordict.update(non_tensordict) + return self + + self._tensordict.update( + input_dict_or_td, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + is_leaf=is_leaf, + ) + # We also need to remove things from non_tensordict + if self._non_tensordict: + keys = set(self._tensordict.keys()) + ntd = {k: val for k, val in self._non_tensordict.items() if k not in keys} + self._non_tensordict.clear() + self._non_tensordict.update(ntd) + return self + + +def _update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + non_blocking: bool = False, +): + if isinstance(input_dict_or_td, dict): + input_dict_or_td = self.from_dict(input_dict_or_td, batch_size=self.batch_size) + + if is_tensorclass(input_dict_or_td): + non_tensordict = { + k: v for k, v in input_dict_or_td._non_tensordict.items() if v is not None + } + self._tensordict.update_(input_dict_or_td._tensordict) + self._non_tensordict.update(non_tensordict) + return self + + self._tensordict.update_( + input_dict_or_td, + clone=clone, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + ) + return self + + +def _update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + index: IndexType, + clone: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + non_blocking: bool = False, +): + if isinstance(input_dict_or_td, dict): + input_dict_or_td = self.from_dict(input_dict_or_td, batch_size=self.batch_size) + + if is_tensorclass(input_dict_or_td): + non_tensordict = { + k: v for k, v in input_dict_or_td._non_tensordict.items() if v is not None + } + self._tensordict.update(input_dict_or_td._tensordict) + self._non_tensordict.update(non_tensordict) + return self + + self._tensordict.update_at_( + input_dict_or_td, + index=index, + clone=clone, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + ) + return self + + +def _wrap_classmethod(td_cls, cls, func): + @functools.wraps(func) + def wrapped_func(*args, **kwargs): + res = func.__get__(td_cls)(*args, **kwargs) + # res = func(*args, **kwargs) + if isinstance(res, TensorDictBase): + # create a new tensorclass from res and copy the metadata from self + return cls._from_tensordict(res) + return res + + return wrapped_func + + +def _getitem(self, item: NestedKey) -> Tensor | TensorCollection | Any: + """Retrieve the class object at the given index. Indexing will happen for nested tensors as well. + + Args: + item (int or any other valid index type): index of the object to retrieve + + Returns: + Tensor class object at the given index + + """ + if isinstance(item, str) or ( + isinstance(item, tuple) and all(isinstance(_item, str) for _item in item) + ): + raise ValueError(f"Invalid indexing arguments: {item}.") + # tensor_res = super(type(self), self).__getattribute__("_tensordict")[item] + tensor_res = self.__dict__["_tensordict"][item] + return _from_tensordict_with_copy(self, tensor_res) # device=res.device) + + +def _setitem(self, item: NestedKey, value: Any) -> None: # noqa: D417 + """Set the value of the Tensor class object at the given index. Note that there is no strict validation on non-tensor values. + + Args: + item (int or any other valid index type): index of the object to set + value (any): value to set for the item + + """ + istuple = isinstance(item, tuple) + if istuple or isinstance(item, str): + # _unravel_key_to_tuple will return an empty tuple if the index isn't a NestedKey + idx_unravel = _unravel_key_to_tuple(item) + if idx_unravel: + raise ValueError(f"Invalid indexing arguments: {item}.") + + if istuple and len(item) == 1: + return _setitem(self, item[0], value) + if ( + ( + isinstance(item, torch.Tensor) + and item.dtype == torch.bool + and not item.shape + and item + ) + or (item is True) + or (item is None) + ) and self.batch_size == (): + return self.update(value.squeeze(0)) + + if not is_tensorclass(value) and not isinstance( + value, (TensorDictBase, numbers.Number, Tensor) + ): + raise ValueError( + f"__setitem__ only supports tensorclasses, tensordicts," + f" numeric scalars and tensors. Got {type(value)}" + ) + + if is_tensorclass(value): + if not isinstance(value, type(self)): + self_keys = set().union(self._non_tensordict, self._tensordict.keys()) + value_keys = set().union(value._non_tensordict, value._tensordict.keys()) + if self_keys != value_keys: + # if tensorclass but different class ensure that all keys are equal + raise ValueError( + "__setitem__ is only allowed for same-class or " + "compatible class (i.e. same members) assignment" + ) + + # Validating the non-tensor data before setting the item + for key, val in value._non_tensordict.items(): + # Raise a warning if non_tensor data doesn't match + if ( + key in self._non_tensordict.keys() + and val is not self._non_tensordict[key] + ): + warnings.warn( + f"Meta data at {repr(key)} may or may not be equal, " + f"this may result in undefined behaviours", + category=UserWarning, + stacklevel=2, + ) + + for key in value._tensordict.keys(): + # Making sure that the key-clashes won't happen, if the key is present + # in tensor data in value we will honor that and remove the key-value + # pair from non-tensor data + if key in self._non_tensordict.keys(): + del self._non_tensordict[key] + + self._tensordict[item] = value._tensordict + else: + # int, float etc. + self._tensordict[item] = value + + +def _repr(self) -> str: + """Return a string representation of Tensor class object.""" + fields = _td_fields(self._tensordict, sep="=") + field_str = [fields] if fields else [] + non_tensor_fields = _all_non_td_fields_as_str(self._non_tensordict) + + medatada_fields = [] + + if "batch_size" not in self.__expected_keys__: + batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ") + medatada_fields.append(batch_size_str) + elif "shape" not in self.__expected_keys__: + batch_size_str = indent(f"shape={self.shape}", 4 * " ") + medatada_fields.append(batch_size_str) + if "device" not in self.__expected_keys__: + device_str = indent(f"device={self.device}", 4 * " ") + medatada_fields.append(device_str) + + is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ") + medatada_fields.append(is_shared_str) + + if len(non_tensor_fields) > 0: + non_tensor_field_str = indent( + ",\n".join(non_tensor_fields), + 4 * " ", + ) + if field_str: + string = ",\n".join(field_str + [non_tensor_field_str, *medatada_fields]) + else: + string = ",\n".join([non_tensor_field_str, *medatada_fields]) + elif field_str: + string = ",\n".join(field_str + medatada_fields) + elif len(medatada_fields) > 0: + string = ",\n".join(medatada_fields) + else: + string = "" + return f"{type(self).__name__}({string})" + + +def _len(self) -> int: + """Returns the length of first dimension, if there is, otherwise 0.""" + return len(self._tensordict) + + +def _to_dict(self, *, retain_none: bool = True, convert_tensors: bool = False) -> dict: + td_dict = self._tensordict.to_dict( + retain_none=retain_none, convert_tensors=convert_tensors + ) + if self._non_tensordict: + if retain_none: + td_dict.update(self._non_tensordict) + else: + td_dict.update( + {k: v for k, v in self._non_tensordict.items() if v is not None} + ) + + return td_dict + + +def _from_dict( + cls, + input_dict, + *, + auto_batch_size: bool | None = None, + batch_size=None, + device=None, + batch_dims=None, +): + # we pass through a tensordict because keys could be passed as NestedKeys + # We can't assume all keys are strings, otherwise calling cls(**kwargs) + # would work ok + if issubclass(cls, NonTensorDataBase): + # Note: this won't deal with sub-tensordicts which may or may not be tensorclasses. + # We don't want to enforce them to be tensorclasses so we can't do much about it... + return cls.from_tensordict( + tensordict=TensorDict( + batch_size=batch_size, + device=device, + ), + non_tensordict=input_dict, + ) + td = TensorDict.from_dict( + input_dict, + batch_size=batch_size, + device=device, + batch_dims=batch_dims, + auto_batch_size=auto_batch_size, + ) + non_tensordict = {} + + return cls.from_tensordict(tensordict=td, non_tensordict=non_tensordict) + + +def _from_dict_instance( + self, + input_dict, + *, + auto_batch_size: bool | None = None, + batch_size=None, + device=None, + batch_dims=None, +): + if batch_dims is not None and batch_size is not None: + raise ValueError("Cannot pass both batch_size and batch_dims to `from_dict`.") + from tensordict import TensorDict + + batch_size_set = torch.Size(()) if batch_size is None else batch_size + # TODO: this is a bit slow and will be a bottleneck every time td[idx] = dict(subtd) + # is called when there are non tensor data in it + if not _is_tensor_collection(type(input_dict)): + input_tdict = TensorDict.from_dict(input_dict, auto_batch_size=auto_batch_size) + else: + input_tdict = input_dict + trsf_dict = {} + for key, value in list(input_tdict.items()): + # cur_value = getattr(self, key, None) + cur_value = self.get(key) + if _is_tensor_collection(type(cur_value)): + trsf_dict[key] = cur_value.from_dict_instance( + value, batch_size=[], device=device, batch_dims=None + ) + elif not isinstance(cur_value, torch.Tensor) and is_non_tensor(value): + trsf_dict[key] = value.data + elif cur_value is not None and not isinstance(cur_value, torch.Tensor): + # This is slightly unsafe but will work with bool, float and int + try: + trsf_dict[key] = type(cur_value)(value) + except Exception: + trsf_dict[key] = input_dict[key] + else: + trsf_dict[key] = value + out = type(self)( + **trsf_dict, + batch_size=batch_size_set, + device=device, + ) + # check that + if batch_size is None: + if auto_batch_size is None and batch_dims is None: + auto_batch_size = False + elif auto_batch_size is None: + auto_batch_size = True + if auto_batch_size: + out.auto_batch_size_() + return out + + +def _to_tensordict(self, *, retain_none: bool | None = None) -> TensorDict: + """Convert the tensorclass into a regular TensorDict. + + Makes a copy of all entries. Memmap and shared memory tensors are converted to + regular tensors. + + Args: + retain_none (bool): if ``True``, the ``None`` values will be written in the + tensordict. Otherwise they will be discrarded. Default: ``True``. + + Returns: + A new TensorDict object containing the same values as the tensorclass. + + """ + td = self._tensordict.to_tensordict(retain_none=retain_none) + for key, val in self._non_tensordict.items(): + if val is None: + if retain_none is None: + retain_none = False + if retain_none: + pass + else: + continue + td.set_non_tensor(key, val) + return td + + +def _device(self) -> torch.device: + """Retrieves the device type of tensor class.""" + return self._tensordict.device + + +def _device_setter(self, value: DeviceType) -> None: + raise RuntimeError( + "device cannot be set using tensorclass.device = device, " + "because device cannot be updated in-place. To update device, use " + "tensorclass.to(new_device), which will return a new tensorclass " + "on the new device." + ) + + +def _set( + self, key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False +): + """Sets a new key-value pair. + + Args: + key (str, tuple of str): name of the key to be set. + If tuple of str it is equivalent to chained calls of getattr + followed by a final setattr. + value (Any): value to be stored in the tensorclass + inplace (bool, optional): if ``True``, set will tentatively try to + update the value in-place. If ``False`` or if the key isn't present, + the value will be simply written at its destination. + + Returns: + self + + """ + if isinstance(key, str): + cls = type(self) + __dict__ = self.__dict__ + if __dict__["_tensordict"].is_locked: + raise RuntimeError(_LOCK_ERROR) + # if key in ("batch_size", "names", "device"): + # # handled by setattr + # return + expected_keys = cls.__expected_keys__ + if key not in expected_keys: + raise AttributeError( + f"Cannot set the attribute '{key}', expected attributes are {expected_keys}." + ) + + self_is_non_tensor = self._is_non_tensor + value_type = type(value) + + def set_tensor( + key=key, + value=value, + inplace=inplace, + non_blocking=non_blocking, + non_tensor=False, + ): + if self_is_non_tensor: + while is_non_tensor(value): + value = value.data + self._non_tensordict[key] = value + return self + if non_tensor: + value = NonTensorData( + value, batch_size=self.batch_size, device=self.device + ) + if key in self._non_tensordict: + del self._non_tensordict[key] + # Avoiding key clash, honoring the user input to assign tensor type data to the key + self._tensordict.set(key, value, inplace=inplace, non_blocking=non_blocking) + return self + + def _is_castable(datatype): + return issubclass(datatype, (int, float, np.ndarray)) + + if cls._autocast: + type_hints = cls._type_hints + if type_hints is not None: + target_cls = type_hints.get(key, _AnyType) + else: + warnings.warn("type_hints are none, cannot perform auto-casting") + target_cls = _AnyType + + if isinstance(value, dict): + if _is_tensor_collection(target_cls): + cast_val = target_cls.from_dict(value, auto_batch_size=False) + self._tensordict.set( + key, cast_val, inplace=inplace, non_blocking=non_blocking + ) + return self + elif type_hints is None: + warnings.warn(type(self)._set_dict_warn_msg) + elif value is not None and issubclass( + target_cls, tuple(tensordict_lib.base._ACCEPTED_CLASSES) + ): + try: + if not issubclass(value_type, target_cls): + if issubclass(target_cls, torch.Tensor): + # first convert to tensor to make sure that the dtype is preserved + value = torch.as_tensor(value) + cast_val = _cast_funcs[target_cls](value) + else: + cast_val = value + except TypeError: + raise TypeError( + f"Failed to cast the value {key} to the type annotation {target_cls}." + ) + return set_tensor(value=cast_val) + elif value is not None and target_cls is not _AnyType: + cast_val = _cast_funcs[target_cls](value) + return set_tensor(value=cast_val, non_tensor=True) + elif target_cls is _AnyType and _is_castable(value_type): + return set_tensor() + non_tensor = not ( + isinstance(value, _ACCEPTED_CLASSES) + or _is_tensor_collection(value_type) + ) + elif ( + issubclass(value_type, torch.Tensor) + or _is_tensor_collection(value_type) + or ( + not cls._nocast + and issubclass(value_type, (int, float, bool, np.ndarray)) + ) + ): + return set_tensor() + elif issubclass(value_type, list) and list_to_stack(): + # set() will take care of casting to non tensor + non_tensor = False + else: + non_tensor = True + + if self_is_non_tensor or value is None: + # Avoiding key clash, honoring the user input to assign non-tensor data to the key + if not self_is_non_tensor and key in self._tensordict.keys(): + if inplace: + raise RuntimeError( + f"Cannot update an existing entry of type {type(self._tensordict.get(key))} with a value of type {value_type}." + ) + self._tensordict.del_(key) + self._non_tensordict[key] = value + else: + if inplace: + if key in self._tensordict.keys(): + raise RuntimeError( + f"Cannot update an existing entry of type {type(self._tensordict.get(key))} with a value of type {value_type}." + ) + return set_tensor(value=value, non_tensor=non_tensor) + return self + + if isinstance(key, tuple) and len(key): + key = _unravel_key_to_tuple(key) + if len(key) > 1: + return self.set(key[0], getattr(self, key[0]).set(key[1:], value)) + out = self.set(key[0], value) + return out + raise ValueError( + f"Supported type for key are str and tuple, got {key} of type {type(key)}" + ) + + +def _set_str( + self, + key: NestedKey, + value: str, + *, + inplace: bool, + validated: bool, + ignore_lock: bool = False, + non_blocking: bool = False, +): + if is_non_tensor(self): + if key != "data": + raise KeyError(f"only 'data' keys are supported for {type(self).__name__}.") + while isinstance(value, (NonTensorData, NonTensorStack)): + value = value.data + self._non_tensordict[key] = value + return self + else: + if key in self._non_tensordict: + del self._non_tensordict[key] + self._tensordict._set_str( + key, + value, + inplace=inplace, + validated=validated, + ignore_lock=ignore_lock, + non_blocking=non_blocking, + ) + return self + + +def _set_at_str( + self, + key: NestedKey, + value: str, + idx, + *, + validated: bool, + non_blocking: bool = False, +): + if is_non_tensor(self): + if key != "data": + raise KeyError(f"only 'data' keys are supported for {type(self).__name__}.") + while isinstance(value, (NonTensorData, NonTensorStack)): + value = value.data + self._non_tensordict[key] = value + return self + else: + if key in self._non_tensordict: + del self._non_tensordict[key] + self._tensordict._set_at_str( + key, value, idx, validated=validated, non_blocking=non_blocking + ) + return self + + +def _delattr(self, key): + del self._tensordict[key] + + +def _del_(self, key): + key = _unravel_key_to_tuple(key) + if len(key) > 1: + td = self.get(key[0]) + td.del_(key[1:]) + return + if key[0] in self._tensordict.keys(): + self._tensordict.del_(key[0]) + # self.set(key[0], None) + elif key[0] in self._non_tensordict.keys(): + self._non_tensordict[key[0]] = None + else: + raise KeyError(f"Key {key} could not be found in tensorclass {self}.") + return + + +def _set_at_( + self, key: NestedKey, value: Any, idx: IndexType, non_blocking: bool = False +): + if key in self._non_tensordict: + del self._non_tensordict[key] + return self._tensordict.set_at_(key, value, idx, non_blocking=non_blocking) + + +def _get(self, key: NestedKey, *args, **kwargs): + """Gets the value stored with the input key. + + Args: + key (str, tuple of str): key to be queried. If tuple of str it is + equivalent to chained calls of getattr. + default: default value if the key is not found in the tensorclass. + + Returns: + value stored with the input key + + """ + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + # Find what the default is + if args: + default = args[0] + if len(args) > 1: + raise TypeError("Only one arg is allowed in TD.get.") + elif "default" in kwargs: + raise TypeError("'default' arg was passed twice.") + elif "default" in kwargs: + default = kwargs.pop("default") + if args: + raise TypeError("'default' arg was passed twice.") + elif _GET_DEFAULTS_TO_NONE: + default = None + else: + default = NO_DEFAULT + + try: + if len(key) > 1: + return _getattr(self, key[0], **kwargs).get( + key[1:], default=default, **kwargs + ) + if kwargs: + return _getattr(self, key[0], **kwargs) + return getattr(self, key[0]) + except (AttributeError, KeyError): + if default is NO_DEFAULT: + raise + return default + + +def _get_at(self, key: NestedKey, *args, **kwargs): + key = _unravel_key_to_tuple(key) + if not key: + raise KeyError(_GENERIC_NESTED_ERR.format(key)) + + try: + if len(args): + index = args[0] + args = args[1:] + else: + index = kwargs.pop("index") + except KeyError: + raise TypeError("index argument missing from get_at") + + # Find what the default is + if args: + default = args[0] + if len(args) > 1 or kwargs: + raise TypeError("only one (keyword) argument is allowed.") + elif kwargs: + default = kwargs.pop("default") + if args or kwargs: + raise TypeError("only one (keyword) argument is allowed.") + elif _GET_DEFAULTS_TO_NONE: + default = None + else: + default = NO_DEFAULT + + try: + return self.get(key, NO_DEFAULT)[index] + except (AttributeError, KeyError): + if default is NO_DEFAULT: + raise + return default + + +def _data(self): + # We allow data to be a field of the class too + if "data" in self.__dataclass_fields__: + data = self._tensordict.get("data") + if data is None: + data = self._non_tensordict.get("data") + return data + return self._from_tensordict(self._tensordict.data, self._non_tensordict) + + +def _data_setter(self, new_data): + if self._is_non_tensor: + self._non_tensor["data"] = new_data + return + if "data" in self.__dataclass_fields__: + self.set("data", new_data) + return + raise AttributeError("property 'data' is read-only.") + + +def _grad(self): + grad = self._tensordict._grad + if grad is None: + return None + return self._from_tensordict(self._tensordict.grad, self._non_tensordict) + + +def _names_setter(self, names: str) -> None: # noqa: D417 + """Set the value of ``tensorclass.names``. + + Args: + names (sequence of str) + + """ + self._tensordict.names = names + + +def _state_dict( + self, destination=None, prefix="", keep_vars=False, flatten=False +) -> dict[str, Any]: + """Returns a state_dict dictionary that can be used to save and load data from a tensorclass.""" + state_dict = { + "_tensordict": super(type(self), self) + .__getattribute__("_tensordict") + .state_dict( + destination=destination, prefix=prefix, keep_vars=keep_vars, flatten=flatten + ) + } + state_dict["_non_tensordict"] = dict(self._non_tensordict) + return state_dict + + +def _load_state_dict( + self, state_dict: dict[str, Any], strict=True, assign=False, from_flatten=False +): + """Loads a state_dict attemptedly in-place on the destination tensorclass.""" + for key, item in state_dict.items(): + # keys will never be nested which facilitates everything, but let's + # double check in case someone does something nasty + if not isinstance(key, str): + raise TypeError("Only str keys are allowed when calling load_state_dict.") + if key == "_non_tensordict": + for sub_key, sub_item in item.items(): + # sub_item is the state dict of a tensorclass + if isinstance(sub_item, dict) and "_non_tensordict" in sub_item: + raise RuntimeError( + "Loading a saved tensorclass on a uninitialized tensorclass is not allowed" + ) + else: + # check that sub_key is part of the tensorclass + if sub_key not in type(self).__dataclass_fields__: + raise KeyError( + f"Key '{sub_key}' wasn't expected in the state-dict." + ) + super(type(self), self).__getattribute__("_non_tensordict")[ + sub_key + ] = sub_item + elif key == "_tensordict": + for sub_key in item.keys(): + if sub_key not in type(self).__dataclass_fields__ and sub_key not in ( + "__batch_size", + "__device", + ): + raise KeyError( + f"Key '{sub_key}' wasn't expected in the state-dict." + ) + super(type(self), self).__getattribute__("_tensordict").load_state_dict( + item, strict=strict, assign=assign, from_flatten=from_flatten + ) + else: + raise KeyError(f"Key '{key}' wasn't expected in the state-dict.") + + return self + + +def _eq(self, other: object) -> bool: + """Compares the Tensor class object to another object for equality. However, the equality check for non-tensor data is not performed. + + Args: + other: object to compare to this object. Can be a tensorclass, a + tensordict or any compatible type (int, float or tensor), in + which case the equality check will be propagated to the leaves. + + Returns: + False if the objects are of different class types, Tensorclass of boolean + values for tensor attributes and None for non-tensor attributes + + Examples: + >>> @tensorclass + ... class MyClass: + ... x: Tensor + ... y: "MyClass" + ... z: str + ... + >>> c1 = MyClass( + ... x=torch.randn(3, 4), + ... y=MyClass( + ... x=torch.randn(3, 4, 1), + ... y=None, + ... z="bar", + ... batch_size=[3, 4, 1], + ... ), + ... z="foo", + ... batch_size=[3, 4], + ... ) + >>> c2 = c1.clone() + >>> print(c1 == c2) + MyClass( + x=Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.bool, is_shared=False), + y=MyClass( + x=Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False), + y=None, + z=None, + batch_size=torch.Size([3, 4, 1]), + device=None, + is_shared=False), + z=None, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> assert (c1 == c2).all() + >>> assert (c1[:2] == c2[:2]).all() + >>> assert not (c1 == c2.apply(lambda x: x+1)).all() + + """ + if not is_tensor_collection(other) and not isinstance( + other, (dict, numbers.Number, Tensor) + ): + return False + if is_tensorclass(other): + tensor = self._tensordict == other._tensordict + else: + tensor = self._tensordict == ( + other.exclude(*self._non_tensordict.keys()) + if _is_tensor_collection(type(other)) + else other + ) + if isinstance(self, NonTensorDataBase): + # Return a plain tensordict containing the eq value + return tensor + return _from_tensordict_with_none(self, tensor) + + +def _ne(self, other: object) -> bool: + """Compare the Tensor class object to another object for inequality. However, the equality check for non-tensor data is not performed. + + Args: + other: object to compare to this object + + Returns: + False if the objects are of different class types, Tensorclass of boolean values for tensor attributes and None for non-tensor attributes + + Examples: + >>> @tensorclass + ... class MyClass: + ... x: Tensor + ... y: "MyClass" + ... z: str + ... + >>> c1 = MyClass( + ... x=torch.randn(3, 4), + ... y=MyClass( + ... x=torch.randn(3, 4, 1), + ... y=None, + ... z="bar", + ... batch_size=[3, 4, 1], + ... ), + ... z="foo", + ... batch_size=[3, 4], + ... ) + >>> c2 = c1.clone() + >>> print(c1 != c2) + MyClass( + x=Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.bool, is_shared=False), + y=MyClass( + x=Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False), + y=None, + z=None, + batch_size=torch.Size([3, 4, 1]), + device=None, + is_shared=False), + z=None, + batch_size=torch.Size([3, 4]), + device=None, + is_shared=False) + >>> c2 = c2.apply(lambda x: x+1) + >>> assert (c1 != c2).all() + + """ + if not is_tensor_collection(other) and not isinstance( + other, (dict, numbers.Number, Tensor) + ): + return True + if is_tensorclass(other): + tensor = self._tensordict != other._tensordict + else: + tensor = self._tensordict != ( + other.exclude(*self._non_tensordict.keys()) + if _is_tensor_collection(type(other)) + else other + ) + if isinstance(self, NonTensorDataBase): + # Return a plain tensordict containing the neq value + return tensor + return _from_tensordict_with_none(self, tensor) + + +def _or(self, other: object) -> bool: + """Compares the Tensor class object to another object for logical OR. However, the logical OR check for non-tensor data is not performed. + + Args: + other: object to compare to this object. Can be a tensorclass, a + tensordict or any compatible type (int, float or tensor), in + which case the equality check will be propagated to the leaves. + + Returns: + False if the objects are of different class types, Tensorclass of boolean + values for tensor attributes and None for non-tensor attributes + + """ + if not is_tensor_collection(other) and not isinstance( + other, (dict, numbers.Number, Tensor) + ): + return False + if is_tensorclass(other): + tensor = self._tensordict | other._tensordict + else: + tensor = self._tensordict | ( + other.exclude(*self._non_tensordict.keys()) + if _is_tensor_collection(type(other)) + else other + ) + if isinstance(self, NonTensorDataBase): + # Return a plain tensordict containing the or value + return tensor + return _from_tensordict_with_none(self, tensor) + + +def _xor(self, other: object) -> bool: + """Compares the Tensor class object to another object for exclusive OR. However, the exclusive OR check for non-tensor data is not performed. + + Args: + other: object to compare to this object. Can be a tensorclass, a + tensordict or any compatible type (int, float or tensor), in + which case the equality check will be propagated to the leaves. + + Returns: + False if the objects are of different class types, Tensorclass of boolean + values for tensor attributes and None for non-tensor attributes + + """ + if not is_tensor_collection(other) and not isinstance( + other, (dict, numbers.Number, Tensor) + ): + return False + if is_tensorclass(other): + tensor = self._tensordict ^ other._tensordict + else: + tensor = self._tensordict ^ ( + other.exclude(*self._non_tensordict.keys()) + if _is_tensor_collection(type(other)) + else other + ) + if isinstance(self, NonTensorDataBase): + # Return a plain tensordict containing the xor value + return tensor + return _from_tensordict_with_none(self, tensor) + + +def _non_tensor_items(self, include_nested=False): + if include_nested: + return self.non_tensor_items() + self._tensordict.non_tensor_items( + include_nested=True + ) + elif is_tensorclass(self): + return list(self._non_tensordict.items()) + else: + return self._tensordict.non_tensor_items() + + +def _bool(self): + raise RuntimeError("Converting a tensorclass to boolean value is not permitted") + + +def _all_non_td_fields_as_str(src_dict) -> list: + """Returns a list of string representation of non-tensor key-value pairs. + + Args: + src_dict (dict): non_tensor_dict + + Returns: + result (list): list of strings with key-value representation + + """ + result = [] + for key, val in src_dict.items(): + if not is_tensor_collection(val): + result.append(f"{key}={repr(val)}") + + return result + + +def _unbind(self, dim: int): + """Returns a tuple of indexed tensorclass instances unbound along the indicated dimension. + + Resulting tensorclass instances will share the storage of the initial tensorclass instance. + + """ + # TODO: dynamo doesn't like copy, using dict instead + return tuple( + type(self)._from_tensordict(td, non_tensordict=dict(self._non_tensordict)) + for td in self._tensordict.unbind(dim) + ) + + +################ +# Custom classes +# -------------- + +NONTENSOR_HANDLED_FUNCTIONS = [] + +_MP_MANAGER = None + + +def _mp_manager(): + global _MP_MANAGER + if _MP_MANAGER is None: + _MP_MANAGER = Manager() + return _MP_MANAGER + + +def _patch_tc(cls): + cls.__setattr__ = _setattr + cls.__getattr__ = _getattr + cls.__getitem__ = _getitem + cls.__setitem__ = _setitem + cls.__len__ = _len + cls.__repr__ = _repr + cls.__eq__ = _eq + cls.__ne__ = _ne + cls.__or__ = _or + cls.__xor__ = _xor + cls.__bool__ = _bool + + cls.device = property(_device, _device_setter) + # cls.data = property(_data, _data_setter) + cls.grad = property(_grad) + + cls._from_tensordict = classmethod(_from_tensordict) + cls.from_tensordict = _from_tensordict + cls._new_unsafe = classmethod(_new_unsafe) + cls._load_memmap = classmethod(_load_memmap) + cls.from_dict = classmethod(_from_dict) + + cls.set = _set + cls.get = _get + cls.update = _update + cls.update_ = _update_ + cls.update_at_ = _update_at_ + cls.to_tensordict = _to_tensordict + cls.to_dict = _to_dict + cls.non_tensor_items = _non_tensor_items + cls.unbind = _unbind + cls._unbind = _unbind + cls.state_dict = _state_dict + cls.load_state_dict = _load_state_dict + cls._memmap_ = _memmap_ + cls.share_memory_ = _share_memory_ + cls.load_memmap = TensorDictBase.load_memmap + cls.load = TensorDictBase.load + cls.from_dict_instance = _from_dict_instance + + # # Methods from lists + for method_name in _METHOD_FROM_TD: + setattr(cls, method_name, getattr(TensorDict, method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD: + setattr(cls, method_name, _wrap_td_method(method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD_FORCE: + setattr(cls, method_name, _wrap_td_method(method_name)) + for method_name in _FALLBACK_METHOD_FROM_TD_NOWRAP: + is_property = isinstance(getattr(TensorDictBase, method_name, None), property) + # if is_property: + # print("method_name", method_name) + # continue + setattr( + cls, + method_name, + _wrap_td_method(method_name, no_wrap=True, is_property=is_property), + ) + for method_name in _FALLBACK_METHOD_FROM_TD_COPY: + setattr( + cls, + method_name, + _wrap_td_method(method_name, copy_non_tensor=True), + ) + return cls + # _set_methods(TensorClass) + + +class _TensorClassMeta(abc.ABCMeta): + def __new__( + mcs, + name, + bases, + namespace, + autocast=None, + nocast=None, + frozen=None, + tensor_only=None, + shadow=None, + **kwargs, + ): + # Create the class using the ABCMeta's __new__ method + cls = super().__new__(mcs, name, bases, namespace, **kwargs) + + # Apply the dataclass decorator to the class + if frozen is None and hasattr(cls, "_frozen"): + frozen = cls._frozen + if nocast is None and hasattr(cls, "_nocast"): + nocast = cls._nocast + if autocast is None and hasattr(cls, "_autocast"): + autocast = cls._autocast + if tensor_only is None and hasattr(cls, "_tensor_only"): + tensor_only = cls._tensor_only + if shadow is None and hasattr(cls, "_shadow"): + shadow = cls._shadow + if name == "TensorClass": + # if "tensordict.tensorclass" in namespace.get( + # "__module__", "" + # ): + cls = _patch_tc(cls) + # ideally we could make it a dataclass but that won't work because: + # - we cannot have a frozen dataclass inherit from a non frozen one + # - if we freeze it we cannot set new methods + # cls = dataclass(cls, init=False) + + delattr(cls, "batch_size") + delattr(cls, "device") + delattr(cls, "names") + else: + cls = tensorclass( + frozen=bool(frozen), + nocast=bool(nocast), + autocast=bool(autocast), + tensor_only=bool(tensor_only), + shadow=bool(shadow), + )(cls) + + return cls + + def __getitem__(cls, item: IndexType) -> Self: + if not isinstance(item, tuple): + item = (item,) + name = "_".join(item) # type: ignore + cls_name = f"TensorClass_{name}" + bases = (cls,) + class_dict = {} + # Copy the __init__ method from the original class + result = _TensorClassMeta( # type: ignore + cls_name, + bases, + class_dict, + **{_item: True for _item in item}, + ) + # Note: We must destroy any property that is set by the tensorclass decorator. + # The result is a base class, so we don't need them, and they will be populated later. + # If they are present, the dataclass decorator applied within tensorclass will look for a default value + # for these guys (when shadow=True) and it will actually find them (since they're properties of the base + # class). This is bad because then we'll be using the property as default value - not what we want. + if "device" in result.__dict__: + delattr(result, "device") + if "batch_size" in result.__dict__: + delattr(result, "batch_size") + if "names" in result.__dict__: + delattr(result, "names") + return result + + +class TensorClass(TensorCollection, metaclass=_TensorClassMeta): + """TensorClass is the inheritance-based version of the @tensorclass decorator. + + TensorClass allows you to code dataclasses that are better type-checked and more pythonic than those built with + the @tensorclass decorator. + + Examples: + >>> from typing import Any + >>> import torch + >>> from tensordict import TensorClass + >>> class Foo(TensorClass): + ... tensor: torch.Tensor + ... non_tensor: Any + ... nested: Any = None + >>> foo = Foo(tensor=torch.randn(3), non_tensor="a string!", nested=None, batch_size=[3]) + >>> print(foo) + Foo( + non_tensor=NonTensorData(data=a string!, batch_size=torch.Size([3]), device=None), + tensor=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + nested=None, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + Keyword Args: + batch_size (torch.Size, optional): The batch size of the TensorDict. Defaults to ``None``. + device (torch.device, optional): The device on which the TensorDict will be created. Defaults to ``None``. + frozen (bool, optional): If ``True``, the resulting class or instance will be immutable. Defaults to ``False``. + autocast (bool, optional): If ``True``, enables automatic type casting for the resulting class or instance. Defaults to ``False``. + nocast (bool, optional): If ``True``, disables any type casting for the resulting class or instance. Defaults to ``False``. + tensor_only (bool, optional): if ``True``, it is expected that all items in tensorclass will be + tensor instances (tensor-compatible, since non-tensor data is converted to tensors if possible). + This can bring significant speed-ups at the cost of flexible interactions with non-tensor data. + Defaults to ``False``. + shadow (bool, optional): Disables the validation of field names against TensorDict's reserved attributes. + Use with caution, as this may cause unintended consequences. Defaults to False. + + You can pass boolean keyword arguments (`"autocast"`, `"nocast"`, `"frozen"`, `"tensor_only"`, `"shadow"`) in two ways: using + brackets or keyword arguments. + + Examples: + >>> class Foo(TensorClass["autocast"]): + ... integer: int + >>> Foo(integer=torch.ones(())).integer + 1 + >>> class Foo(TensorClass, autocast=True): # equivalent + ... integer: int + >>> Foo(integer=torch.ones(())).integer + 1 + >>> class Foo(TensorClass["nocast"]): + ... integer: int + >>> Foo(integer=1).integer + 1 + >>> class Foo(TensorClass["nocast", "frozen"]): # multiple keywords can be used + ... integer: int + >>> Foo(integer=1).integer + 1 + >>> class Foo(TensorClass, nocast=True): # equivalent + ... integer: int + >>> Foo(integer=1).integer + 1 + >>> class Foo(TensorClass): + ... integer: int + >>> Foo(integer=1).integer + tensor(1) + + .. warning:: TensorClass itself is not decorated as a tensorclass, but subclasses will be. + This is because we cannot anticipate if the frozen argument will be set, and if it is, it may + conflict with the parent class (a subclass cannot be frozen if the parent class isn't). + + """ + + _autocast = False + _nocast = False + _frozen = False + _tensor_only = False + + _is_tensorclass = True + _is_non_tensor = False + + # Provide typing-friendly indexing support. Static type checkers look for + # __class_getitem__ on the class rather than a metaclass __getitem__. + # Delegate to the metaclass implementation that builds the configured subclass. + def __class_getitem__(cls, item: IndexType) -> Self: # type: ignore[override] + return _TensorClassMeta.__getitem__(cls, item) + + +def _check_equal(a, b): + # A util to check that two non-tensor data match + # We're replacing this by an identity match, not a value check (which will be faster and easier to handle). + try: + if isinstance(a, _ACCEPTED_CLASSES) or isinstance(b, _ACCEPTED_CLASSES): + iseq = (a == b).all() and a.shape == b.shape + elif isinstance(a, np.ndarray) or isinstance(b, np.ndarray): + iseq = (a == b).all() and a.shape == b.shape + else: + iseq = bool(a == b) + except Exception: + iseq = False + return iseq + + +class NonTensorDataBase(TensorClass): + """A base class to carry non-tensor data. + + There are two main `NonTensorDataBase` subclasses: :class:`~tensordict.NonTensorData` which behaves + mostly as a regular tensordict when shape operations are applied, and :class:`~tensordict.MetaData` + which is more specific. + + The main difference between the two classes is their behavior during expansion or stacking. The + :class:`~tensordict.MetaData` class will keep a single copy of the data for the entire tensordict. + As the name suggests, the intended usage is to carry data that provides additional information about + the batch of data stored in a `TensorDict`. On the other hand, the :class:`~tensordict.NontensorData` + class will carry data in a batch-size compliant manner: the batch-size of the tensorclass is indicative + of different batch elements within it. + + """ + + # Used to carry non-tensor data in a tensordict. + # The advantage of storing this in a tensorclass is that we don't need + # to patch tensordict with additional checks that will encur unwanted overhead + # and all the overhead falls back on this class. + data: Any + _metadata: dict | None = None + + _is_non_tensor: bool = True + + def __repr__(self): + data_str = str(self.data) + if len(data_str) > 200: + data_str = data_str[:20] + " ... " + data_str[-20:] + repr_str = f"{type(self).__name__}(data={data_str}" + if "batch_size" not in self.__expected_keys__: + repr_str += f", batch_size={self.batch_size}" + elif "shape" not in self.__expected_keys__: + repr_str += f", shape={self.shape}" + if "device" not in self.__expected_keys__: + repr_str += f", device={self.device}" + return repr_str + ")" + + def __post_init__(self): + _tensordict = self.__dict__["_tensordict"] + _non_tensordict = self.__dict__["_non_tensordict"] + data = _non_tensordict.get("data", NO_DEFAULT) + if data is NO_DEFAULT: + data = _tensordict._get_str("data", default=NO_DEFAULT) + data_inner = getattr(data, "data", None) + if data_inner is None: + # Support for stacks + data_inner = data.tolist() + del _tensordict["data"] + _non_tensordict["data"] = data_inner + + # TODO: this will probably fail with dynamo at some point, + it's terrible. + # Make sure it's patched properly at init time + old_eq = type(self).__eq__ + if old_eq is _eq: + global NONTENSOR_HANDLED_FUNCTIONS + NONTENSOR_HANDLED_FUNCTIONS.extend(TD_HANDLED_FUNCTIONS) + + # Patch only the first time a class is created + + @functools.wraps(_eq) + def __eq__(self, other): + if isinstance(other, NonTensorDataBase): + eqval = self.data == other.data + if isinstance(eqval, torch.Tensor): + return eqval + if isinstance(eqval, np.ndarray): + return torch.as_tensor(eqval, device=self.device) + return torch.full( + self.batch_size, + bool(eqval), + device=self.device, + ) + # # Handle comparison with scalar values (like 0, 1, etc.) + # # For non-tensor data, we should return a boolean tensor + # if isinstance(other, (int, float, bool)) or (isinstance(other, torch.Tensor) and other.numel() == 1): + # eqval = self.data == other + # if isinstance(eqval, torch.Tensor): + # return eqval + # if isinstance(eqval, np.ndarray): + # return torch.as_tensor(eqval, device=self.device) + # return torch.full( + # self.batch_size, + # bool(eqval), + # device=self.device, + # ) + return old_eq(self, other) + + type(self).__eq__ = __eq__ + + _ne = type(self).__ne__ + + @functools.wraps(_ne) + def __ne__(self, other): + if isinstance(other, NonTensorDataBase): + neqval = self.data != other.data + if isinstance(neqval, torch.Tensor): + return neqval + if isinstance(neqval, np.ndarray): + return torch.as_tensor(neqval, device=self.device) + return torch.full( + self.batch_size, + bool(neqval), + device=self.device, + ) + # # Handle comparison with scalar values (like 0, 1, etc.) + # # For non-tensor data, we should return a boolean tensor + # if isinstance(other, (int, float, bool)) or (isinstance(other, torch.Tensor) and other.numel() == 1): + # neqval = self.data != other + # if isinstance(neqval, torch.Tensor): + # return neqval + # if isinstance(neqval, np.ndarray): + # return torch.as_tensor(neqval, device=self.device) + # return torch.full( + # self.batch_size, + # bool(neqval), + # device=self.device, + # ) + return _ne(self, other) + + type(self).__ne__ = __ne__ + + _xor = type(self).__xor__ + + @functools.wraps(_xor) + def __xor__(self, other): + if isinstance(other, NonTensorDataBase): + xorval = self.data ^ other.data + if isinstance(xorval, torch.Tensor): + return xorval + if isinstance(xorval, np.ndarray): + return torch.as_tensor(xorval, device=self.device) + return torch.full( + self.batch_size, + bool(xorval), + device=self.device, + ) + return _xor(self, other) + + type(self).__xor__ = __xor__ + + _or = type(self).__or__ + + @functools.wraps(_or) + def __or__(self, other): + if isinstance(other, NonTensorDataBase): + orval = self.data | other.data # yuppie! + if isinstance(orval, torch.Tensor): + return orval + if isinstance(orval, np.ndarray): + return torch.as_tensor(orval, device=self.device) + return torch.full( + self.batch_size, + bool(orval), + device=self.device, + ) + return _or(self, other) + + type(self).__or__ = __or__ + + def __call__(self, *args, **kwargs): + """Calling a NonTensorDataBase falls back to a call of its data.""" + return self.data(*args, **kwargs) + + def __getitem__(self, idx: IndexType) -> Self | Tensor | TensorCollection | Any: + if isinstance(self.data, list): + new_data = [self.data[i] for i in torch.as_tensor(idx).tolist()] + else: + new_data = self.data[idx] + new_batch_size = ( + torch.Size(torch.Size(self.batch_size)[idx]) + if self.batch_size + else torch.Size([]) + ) + return NonTensorData(new_data, new_batch_size, self.device) + + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> T: + return self._update( + input_dict_or_td=input_dict_or_td, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + + def _update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + break_on_memmap: bool | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> T: + if isinstance(input_dict_or_td, NonTensorDataBase): + data = input_dict_or_td.data + if inplace and self._tensordict._is_shared: + _update_shared_nontensor(self._non_tensordict["data"], data) + return self + elif inplace and self._is_memmap: + _is_memmaped_from_above = self._is_memmaped_from_above() + if break_on_memmap is None: + global _BREAK_ON_MEMMAP + break_on_memmap = _BREAK_ON_MEMMAP + if _is_memmaped_from_above and break_on_memmap: + raise RuntimeError( + "Cannot update a leaf NonTensorDataBase from a memmaped parent NonTensorStack. " + "To update this leaf node, please update the NonTensorStack with the proper index." + ) + share_non_tensor = self._metadata["_share_non_tensor"] + if share_non_tensor: + _update_shared_nontensor(self._non_tensordict["data"], data) + else: + self._non_tensordict["data"] = data + # Force json update by setting is memmap to False + if not _is_memmaped_from_above and "memmap_prefix" in self._metadata: + self._tensordict._is_memmap = False + self._memmap_( + prefix=self._metadata["memmap_prefix"], + copy_existing=False, + executor=None, + futures=None, + inplace=True, + like=False, + share_non_tensor=share_non_tensor, + robust_key=True, + ) + return self + elif not inplace and self.is_locked: + raise RuntimeError(_LOCK_ERROR) + if clone: + data = deepcopy(data) + self.data = data + elif isinstance(input_dict_or_td, NonTensorStack): + raise ValueError( + "Cannot update a NonTensorDataBase object with a NonTensorStack. Call `non_tensor_data.maybe_to_stack()` " + "before calling update()." + ) + elif not input_dict_or_td.is_empty(): + raise RuntimeError(f"Unexpected type {type(input_dict_or_td)}") + return self + + def __getattr__(self, item): + if item == "data": + if self.is_memmap: + # unwrap the shared object + return _from_shared_nontensor(self._non_tensor["data"]) + return self._non_tensor["data"] + return _getattr(self, item) + + def maybe_to_stack(self): + """Converts the NonTensorDataBase object to a NonTensorStack object if it has a non-empty batch-size.""" + datalist = self.data + if not self.batch_size: + return self + for i in reversed(self.batch_size): + datalist = [datalist] * i + return NonTensorStack._from_list(datalist, device=self.device, ndim=self.ndim) + + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> T: + return self._update_( + input_dict_or_td=input_dict_or_td, + clone=clone, + keys_to_update=keys_to_update, + ) + + def _update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + break_on_memmap: bool | None = None, + ) -> T: + + if isinstance(input_dict_or_td, NonTensorStack): + raise RuntimeError( + "Cannot update a NonTensorDataBase with a NonTensorStack object." + ) + if not isinstance(input_dict_or_td, NonTensorDataBase): + raise RuntimeError( + "NonTensorDataBase.copy_ / update_ requires the source to be a NonTensorDataBase object." + ) + return self._update( + input_dict_or_td, + inplace=True, + clone=clone, + keys_to_update=keys_to_update, + break_on_memmap=break_on_memmap, + ) + + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorCollection, + index: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + ) -> T: + if index != () and index != slice(None): + raise RuntimeError("Cannot update a part of a NonTensorDataBase.") + return self.update_( + input_dict_or_td=input_dict_or_td, clone=clone, non_blocking=non_blocking + ) + + def empty(self, recurse=False, *, device=NO_DEFAULT, batch_size=None, names=None): + if batch_size is not None and names is None: + names = None + else: + names = self._maybe_names() + return type(self)( + data=self.data, + batch_size=self.batch_size if batch_size is None else batch_size, + names=names, + device=self.device if device is NO_DEFAULT else device, + ) + + def is_empty(self) -> bool: + return False + + def _apply_nest(self, *args, out=None, **kwargs): + # kwargs["filter_empty"] = False + if out is not None: + return out + return self.empty( + batch_size=kwargs.get("batch_size"), + device=kwargs.get("device", NO_DEFAULT), + names=kwargs.get("names"), + ) + + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + ) -> dict[str, Any]: + # override to_dict to return just the data + return self.data + + def to_tensordict(self, *, retain_none: bool | None = None): + return self + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + # A modified version of __torch_function__ to account for the different behaviour + # of stack, which should return lazy stacks of data of data does not match. + if func not in _TD_PASS_THROUGH or not all( + issubclass(t, (Tensor, cls)) for t in types + ): + return NotImplemented + + escape_conversion = func in (torch.stack,) + + if kwargs is None: + kwargs = {} + + # get the output type from the arguments / keyword arguments + if len(args) > 0: + tensorclass_instance = args[0] + else: + tensorclass_instance = kwargs.get("input", kwargs["tensors"]) + if isinstance(tensorclass_instance, (tuple, list)): + tensorclass_instance = tensorclass_instance[0] + if not escape_conversion: + args = tuple(_arg_to_tensordict(arg) for arg in args) + kwargs = {key: _arg_to_tensordict(value) for key, value in kwargs.items()} + + result = TD_HANDLED_FUNCTIONS[func](*args, **kwargs) + if isinstance(result, (list, tuple)): + return type(result)( + _from_tensordict_with_copy(tensorclass_instance, tensordict_result) + for tensordict_result in result + ) + if not escape_conversion: + return _from_tensordict_with_copy(tensorclass_instance, result) + return result + + def _fast_apply(self, *args, **kwargs): + kwargs["filter_empty"] = False + return _wrap_method( + self, "_fast_apply", self._tensordict._fast_apply, nowarn=True + )(*args, **kwargs) + + def _multithread_rebuild(self, *args, **kwargs): + kwargs["filter_empty"] = False + return _wrap_method( + self, + "_multithread_rebuild", + self._tensordict._multithread_rebuild, + nowarn=True, + )(*args, **kwargs) + + def tolist( + self, + *, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ): + """Converts the data in a list if the batch-size is non-empty. + + If the batch-size is empty, returns the data. + + Keyword Args: + convert_tensors (bool, "numpy"): if ``True``, tensors will be converted to lists when creating the dictionary. + If "numpy", tensors will be converted to numpy arrays. + Otherwise, they will remain as tensors. Default: ``False``. + tolist_first (bool, optional): if ``True``, the tensordict will be converted to a list first when + it has batch dimensions. Default: ``False``. + as_linked_list (bool, optional): if ``True``, the list will be converted to a :class:`tensordict.utils.LinkedList` + which will automatically update the tensordict when the list is modified. Default: ``False``. + """ + if not self.batch_size: + return self.data + result = [ + ntd.tolist( + convert_tensors=convert_tensors, + tolist_first=tolist_first, + as_linked_list=as_linked_list, + ) + for ntd in self.unbind(0) + ] + if as_linked_list: + return LinkedList(result, td=self) + return result + + def copy_( + self, src: NonTensorDataBase | NonTensorStack, non_blocking: bool = False + ): + return self.update_(src, non_blocking=non_blocking) + + def clone(self, recurse: bool = True): + if recurse: + return type(self)( + data=deepcopy(self.data), + batch_size=self.batch_size, + device=self.device, + names=self.names if self._has_names() else None, + ) + return type(self)( + data=self.data, + batch_size=self.batch_size, + device=self.device, + names=self.names if self._has_names() else None, + ) + + def share_memory_(self): + if self._tensordict._is_shared: + return self + with self.unlock_(): + self.__dict__["_non_tensordict"]["data"] = _share_memory_nontensor( + self.data, manager=_mp_manager() + ) + self._tensordict.share_memory_() + return self + + def _memmap_( + self, + *, + prefix: str | None = None, + copy_existing: bool = False, + executor=None, + futures=None, + inplace=True, + like=False, + memmaped: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key, + ): + # For efficiency, we can avoid doing this saving + # if the data is already there. + if self._tensordict._is_memmap and str( + getattr(self._tensordict, "_memmap_prefix", None) + ) == str(prefix): + return self + + _metadata = {} + if prefix is not None: + _metadata = dict(self._metadata) if self._metadata is not None else {} + _metadata["memmap_prefix"] = prefix + _metadata["memmaped"] = memmaped + + out = _memmap_( + self, + prefix=prefix, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + memmaped=memmaped, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + _metadata["_share_non_tensor"] = share_non_tensor + out._non_tensordict["_metadata"] = _metadata + if share_non_tensor: + out._non_tensordict["data"] = _share_memory_nontensor( + out.data, manager=_mp_manager() + ) + return out + + def _is_memmaped_from_above(self): + _metadata = self._metadata + if _metadata is None: + return False + return _metadata.get("memmaped", False) + + _load_memmap = classmethod(_load_memmap) + + +class NonTensorData(NonTensorDataBase): + """A carrier for non-tensordict data. + + This class can be used whenever non-tensor data needs to be carried at + any level of a tensordict instance. + + :class:`~tensordict.tensorclass.NonTensorData` instances can be created + explicitly or using :meth:`~tensordict.TensorDictBase.set_non_tensor`. + + This class is serializable using :meth:`tensordict.TensorDictBase.memmap` + and related methods, and can be loaded through :meth:`~tensordict.TensorDictBase.load_memmap`. + If the content of the object is JSON-serializable, it will be serializsed in + the `meta.json` file in the directory pointed by the parent key of the `NoneTensorData` + object. If it isn't, serialization will fall back on pickle. This implies + that we assume that the content of this class is either json-serializable or + pickable, and it is the user responsibility to make sure that one of these + holds. We try to avoid pickling/unpickling objects for performance and security + reasons (as pickle can execute arbitrary code during loading). + + .. note:: + If the data passed to :class:`NonTensorData` is a :class:`NonTensorData` + itself, the data from the nested object will be gathered. + + >>> non_tensor = NonTensorData("a string!") + >>> non_tensor = NonTensorData(non_tensor) + >>> assert non_tensor.data == "a string!" + + .. note:: + To faciliate ``NonTensorData`` integration in tensordict, the + :meth:`~tensordict.TensorDictBase.__getitem__` and :meth:`~tensordict.TensorDictBase.__setitem__` + are overloaded to set non-tensor data appropriately (unlike :meth:`~tensordict.TensorDictBase.set` + and :meth:`~tensordict.TensorDictBase.get` which are reserved for tensor-like + objects): + + >>> td = TensorDict({"a": torch.zeros(3)}, batch_size=[3]) + >>> td["a"] # gets a tensor + >>> td["b"] = "a string!" + >>> assert td["b"] == "a string!" + >>> # indexing preserves the meta-data + >>> assert td[0]["b"] == "a string!" + >>> td.get("b") # returns the NonTensorData + + One can uses lists to set multiple `NonTensorData` at the same time (if :class:`~tensordict.set_list_to_stack` + is set to `True`): + + >>> from tensordict import TensorDict, set_list_to_stack + >>> set_list_to_stack(True).set() + >>> td = TensorDict(batch_size=(3,)) + >>> td["foo"] = ["a", "b", "c"] + >>> print(td) + TensorDict( + fields={ + foo: NonTensorStack( + ['a', 'b', 'c'], + batch_size=torch.Size([3]), + device=None)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + .. note:: + Unlike other tensorclass classes, :class:`NonTensorData` supports + comparisons of two non-tensor data through :meth:`~.__eq__`, :meth:`~.__ne__`, + :meth:`~.__xor__` or :meth:`~.__or__`. These operations return a tensor + of shape `batch_size`. For compatibility with ` == `, + comparison with non-:class:`NonTensorData` will always return an empty + :class:`NonTensorData`. + + >>> a = NonTensorData(True) + >>> b = NonTensorData(True) + >>> assert a == b + >>> assert not (a != b) + >>> assert not (a ^ b) + >>> assert a | b + >>> # The output is a tensor of shape batch-size + >>> a = NonTensorData(True, batch_size=[3]) + >>> b = NonTensorData(True, batch_size=[3]) + >>> print(a == b) + tensor([True, True, True]) + + .. note:: + Stacking :class:`NonTensorData` instances results + in a :class:`~tensordict.NonTensorStack` instance. + The data is not copied during stacking / expansion etc., so that + the memory footprint of these operations is negligeable. + If you're willing to keep a single non-tensor copy during these operations, + the :class:`~tensordict.MetaData` class can be used instead. + + >>> data = torch.stack([NonTensorData(1) for _ in range(10)]) + >>> data + NonTensorStack( + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + batch_size=torch.Size([10]), + device=None) + + .. note:: + Non-tensor data can be filtered out from a tensordict using + :meth:`~tensordict.TensorDictBase.filter_non_tensor`. + + Examples: + >>> # create an instance explicitly + >>> non_tensor = NonTensorData("a string!", batch_size=[]) # batch-size can be anything + >>> data = TensorDict({}, batch_size=[3]) + >>> data.set_non_tensor(("nested", "key"), "a string!") + >>> assert isinstance(data.get(("nested", "key")), NonTensorData) + >>> assert data.get_non_tensor(("nested", "key")) == "a string!" + >>> # serialization + >>> class MyPickableClass: + ... value = 10 + >>> data.set_non_tensor("pickable", MyPickableClass()) + >>> import tempfile + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... data.memmap(tmpdir) + ... loaded = TensorDict.load_memmap(tmpdir) + ... # print directory path + ... print_directory_tree(tmpdir) + Directory size: 511.00 B + tmp2cso9og_/ + pickable/ + _tensordict/ + meta.json + other.pickle + meta.json + nested/ + key/ + _tensordict/ + meta.json + meta.json + meta.json + meta.json + >>> assert loaded.get_non_tensor("pickable").value == 10 + + .. note:: + __Preallocation__ is also possible with ``NonTensorData``. + This class can handle conversion from ``NonTensorData`` to + ``NonTensorStack`` where appropriate, as the following example + demonstrates: + + >>> td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10]) + >>> print(td) + TensorDict( + fields={ + val: NonTensorData( + data=0, + _metadata=None, + _is_non_tensor=True, + batch_size=torch.Size([10]), + device=None, + is_shared=False)}, + batch_size=torch.Size([10]), + device=None, + is_shared=False) + >>> print(td["val"]) + 0 + >>> newdata = TensorDict({"val": NonTensorData(data=1, batch_size=[5])}, [5]) + >>> td[1::2] = newdata + >>> print(td) + TensorDict( + fields={ + val: NonTensorStack( + [0, 1, 0, 1, 0, 1, 0, 1, 0, 1], + batch_size=torch.Size([10]), + device=None)}, + batch_size=torch.Size([10]), + device=None, + is_shared=False) + >>> print(td["val"]) # the stack is automatically converted to a list + [0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + + If the value is unique, the ``NonTensorData`` container is kept and + retrieving the value only returns this value. If a ``NonTensorStack`` + is used, ``__getitem__`` will return the list of values instead. + This makes the two operations not exactly interchangeable. The reason + for this inconsistency is that a single ``NonTensorData`` with a non-empty + batch-size is intended to be used as a metadata carrier for bigger + tensordicts, whereas ``NonTensorStack`` usage is aimed at allocating + one metadata atom to each corresponding batch element. + + .. note:: + ``NonTensorData`` can be shared between processes. In fact, both + :meth:`~tensordict.TensorDict.memmap_` (and the likes) and + :meth:`~tensordict.TensorDict.share_memory_` will produce sharable + instances. + + Valid methods to write data are :meth:`~tensordict.TensorDictBase.update` + with the `inplace=True` flag and :meth:`~tensordict.TensorDictBase.update_` + or :meth:`~tensordict.TensorDictBase.update_at_`. + + >>> if __name__ == "__main__": + ... td = TensorDict({"val": NonTensorData(data=0, batch_size=[])}, []) + ... td.share_memory_() + ... td.update_(TensorDict({"val": NonTensorData(data=1, batch_size=[])}, [])) # works + ... td.update(TensorDict({"val": NonTensorData(data=1, batch_size=[])}, []), inplace=True) # works + ... td["val"] = 1 # breaks + + A shared ``NonTensorData`` is writable whenever its content is a ``str``, + ``int``, ``float``, ``bool``, ``dict`` or ``list`` instance. Other types + (e.g., dataclasses) will not raise an exception during the call to + ``memmap_`` or ``share_memory_`` but they will cause the code to break + when the data is overwritten. + + >>> @dataclass + ... class MyClass: + ... string: str + ... + >>> if __name__ == "__main__": + ... td = TensorDict({"val": MyClass("a string!")}, []) + ... td.share_memory_() # works and can be shared between processes + ... td.update_(TensorDict({"val": MyClass("another string!")}, [])) # breaks! + + :class:`~tensordict.tensorclass.TensorStack` instances are also sharable + in a similar way. Crucially, preallocation must be properly handled for + this to work. + + >>> td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10]) + >>> newdata = TensorDict({"val": NonTensorData(data=1, batch_size=[5])}, [5]) + >>> td[1::2] = newdata + >>> # If TD is properly preallocated, we can share it and change its content + >>> td.share_memory_() + >>> newdata = TensorDict({"val": NonTensorData(data=2, batch_size=[5])}, [5]) + >>> td[1::2] = newdata # Works! + >>> # In contrast, not preallocating the tensordict properly will break when assigning values + >>> td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10]) + >>> td.share_memory_() + >>> newdata = TensorDict({"val": NonTensorData(data=2, batch_size=[5])}, [5]) + >>> td[1::2] = newdata # breaks! + + Writable memmapped-``NonTensorData`` instances will update the underlying + metadata if required. This involves writing in a JSON file, which can + introduce some overhead. We advise against this usage whenever one seeks + performance and long-lasting data sharing isn't required (``share_memory_`` + should be preferred in these cases). + + >>> if __name__ == "__main__": + ... td = TensorDict({"val": NonTensorData(data=0, batch_size=[])}, []) + ... td.memmap_(dest_folder) + ... td.update_(TensorDict({"val": NonTensorData(data=1, batch_size=[])}, [])) + ... # The underlying metadata on disk is updated during calls to update_ + ... td_load = TensorDict.load_memmap(dest_folder) + ... assert (td == td_load).all() + + ``NonTensorData`` can store callables. If called, it will fallback on the `__call__` of `.data`: + + >>> td0 = TensorDict({"a": 0, "b": 0}) + >>> td1 = TensorDict({"a": 1, "b": 1}) + >>> td_func = TensorDict({"a": lambda x, y: x-y, "b": lambda x, y: x+y}) + >>> td = td0.apply(lambda x, y, func: func(x, y), td1, td_func) + >>> assert td["a"] == -1 + >>> assert td["b"] == 1 + + """ + + _load_memmap = classmethod(_load_memmap) + _from_dict = classmethod(_from_dict) + _from_tensordict = classmethod(_from_tensordict) + __repr__ = NonTensorDataBase.__repr__ + + def expand(self, *args, **kwargs) -> T: + # tensordict_dims = self.batch_dims + shape = _get_shape_from_args(*args, **kwargs) + + # Replicate self until we have the appropriate batch size + out = self + for i, s in enumerate(reversed(shape)): + j = -i - 1 + if i < self.ndim and self.batch_size[j] == s: + continue + elif i < self.ndim and self.batch_size[j] == 1: + out = torch.cat([out.copy() for _ in range(s)], j) + else: + out = torch.stack([out.copy() for _ in range(s)]) + return out + + def unsqueeze(self, dim: int): + return torch.stack([self], dim) + + @classmethod + def _stack_non_tensor( + cls, + list_of_non_tensor: list[NonTensorDataBase | NonTensorStack], + dim: int = 0, + raise_if_non_unique=False, + ): + # checks have been performed previously, so we're sure the list is non-empty + first = list_of_non_tensor[0] + + ids = set() + firstdata = NO_DEFAULT + return_stack = not capture_non_tensor_stack() + if return_stack: + return NonTensorStack(*list_of_non_tensor, stack_dim=dim) + for data in list_of_non_tensor: + if not isinstance(data, cls): + if raise_if_non_unique: + cls._stack_non_tensor(data, raise_if_non_unique=raise_if_non_unique) + else: + return_stack = True + break + if firstdata is NO_DEFAULT: + firstdata = data.data + ids.add(id(data.data)) + if len(ids) > 1: + if _check_equal(data.data, firstdata): + continue + if raise_if_non_unique: + raise ValueError( + "More than one unique value has been found in the stack." + ) + return_stack = True + break + else: + return_stack = not capture_non_tensor_stack() + if not return_stack: + batch_size = list(first.batch_size) + batch_size.insert(dim, len(list_of_non_tensor)) + return NonTensorData( + data=first.data, + batch_size=batch_size, + names=first._maybe_names(), + device=first.device, + ) + + return NonTensorStack(*list_of_non_tensor, stack_dim=dim) + + +def _reconstruct_typed_metadata(item, data, state): + """Reconstruct a typed MetaData instance during unpickling.""" + from tensordict.tensorclass import MetaData + + instance = MetaData[item](data) + instance.__dict__.update(state) + return instance + + +class _MetaDataMeta(_TensorClassMeta): + def __new__( + mcs, + name, + bases, + namespace, + datatype=None, + **kwargs, + ): + # Create the class using the parent's __new__ method + cls = super().__new__(mcs, name, bases, namespace, **kwargs) + if datatype is not None: + cls._datatype = datatype + # Initialize cache for typed classes + if not hasattr(cls, "_typed_class_cache"): + cls._typed_class_cache = {} + return cls + + def __getitem__(cls, item): + """Create a typed version of MetaData that validates the data type.""" + if cls.__name__ != "MetaData": + # Only allow type specification on the base MetaData class + raise TypeError(f"Cannot specify type for {cls.__name__}") + + # Check cache first + if item in cls._typed_class_cache: # type: ignore + return cls._typed_class_cache[item] # type: ignore + + # Create a new class that validates the data type + type_name = getattr(item, "__name__", str(item)) + class_name = f"MetaData[{type_name}]" + + class TypedMetaData(cls): + _expected_type = item + # Define all the classes defined by MetaData + _load_memmap = classmethod(_load_memmap) + _from_dict = classmethod(_from_dict) + _from_tensordict = classmethod(_from_tensordict) + __repr__ = NonTensorDataBase.__repr__ + + def __post_init__(self): + super().__post_init__() + # Validate the data type + if not isinstance(self.data, item): + expected_name = getattr(item, "__name__", str(item)) + actual_name = type(self.data).__name__ + raise TypeError( + f"Expected data of type {expected_name}, got {actual_name}" + ) + + TypedMetaData.__name__ = class_name + TypedMetaData.__qualname__ = class_name + + # Add pickle support for the dynamically created class + def __reduce__(self): + # Return a callable that can reconstruct the object + return (_reconstruct_typed_metadata, (item, self.data, self.__dict__)) + + TypedMetaData.__reduce__ = __reduce__ + + # Cache the class + cls._typed_class_cache[item] = TypedMetaData # type: ignore + + return TypedMetaData + + +class MetaData(NonTensorDataBase, metaclass=_MetaDataMeta): + """A non-tensor, metadata carrier class for `TensorDict`. + + This class mainly behaves as :class:`~tensordict.NonTensorData`, except for indexing, + stacking, squeezing/unsqueezing and similar operations. + + During __stacking__, `MetaData` will check if the content of the various items match + in identity (i.e., using `is` and not `==`). If so, a single `MetaData` instance will be + returned with the shape adapted to the stack operations. If not, a :class:`~tensordict.NonTensorStack` + instance will be returned. + + Similarly, :func:`~torch.unsqueeze` will return a `MetaData` instance and not a stack (as it does for + :class:`~tensordict.NonTensorData`). + + """ + + # Remove the __class_getitem__ method since the metaclass handles it + + _load_memmap = classmethod(_load_memmap) + _from_dict = classmethod(_from_dict) + _from_tensordict = classmethod(_from_tensordict) + __repr__ = NonTensorDataBase.__repr__ + + @classmethod + def _stack_non_tensor( + cls, + list_of_non_tensor: list[NonTensorDataBase | NonTensorStack], + dim: int = 0, + raise_if_non_unique=False, + ): + # checks have been performed previously, so we're sure the list is non-empty + first = list_of_non_tensor[0] + + ids = set() + firstdata = NO_DEFAULT + return_stack = False + for data in list_of_non_tensor: + if not isinstance(data, cls): + if raise_if_non_unique: + cls._stack_non_tensor(data, raise_if_non_unique=raise_if_non_unique) + else: + return_stack = True + break + if firstdata is NO_DEFAULT: + firstdata = data.data + ids.add(id(data.data)) + if len(ids) > 1: + if raise_if_non_unique: + raise ValueError( + "More than one unique value has been found in the stack." + ) + return_stack = True + break + if not return_stack: + batch_size = list(first.batch_size) + batch_size.insert(dim, len(list_of_non_tensor)) + return cls( + data=first.data, + batch_size=batch_size, + names=first._maybe_names(), + device=first.device, + ) + + return NonTensorStack(*list_of_non_tensor, stack_dim=dim) + + +# For __setitem__ and _update_at_ we don't pass a kwarg but use a global variable instead +_BREAK_ON_MEMMAP = True + + +class NonTensorStack(LazyStackedTensorDict): + """A thin wrapper around LazyStackedTensorDict to make stack on non-tensor data easily recognizable. + + A ``NonTensorStack`` is returned whenever :func:`~torch.stack` is called on + a list of :class:`~tensordict.NonTensorData` or ``NonTensorStack``. + + Examples: + >>> from tensordict import NonTensorData + >>> import torch + >>> data = torch.stack([ + ... torch.stack([NonTensorData(data=(i, j), batch_size=[]) for i in range(2)]) + ... for j in range(3)]) + >>> print(data) + NonTensorStack( + [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(0, 2), (1, ..., + batch_size=torch.Size([3, 2]), + device=None) + + To obtain the values stored in a ``NonTensorStack``, call :class:`~.tolist`. + + """ + + _is_non_tensor: bool = True + + def __init__(self, *args, **kwargs): + args = [ + arg if is_tensor_collection(arg) else NonTensorData(arg) for arg in args + ] + super().__init__(*args, **kwargs) + if not all(is_non_tensor(item) for item in self.tensordicts): + raise RuntimeError("All tensordicts must be non-tensors.") + + def tolist( + self, + *, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ): + """Extracts the content of a :class:`tensordict.tensorclass.NonTensorStack` in a nested list. + + Keyword Args: + convert_tensors (bool): if ``True``, tensors will be converted to lists. + Otherwise, they will remain as tensors. Default: ``False``. + tolist_first (bool, optional): if ``True``, the tensordict will be converted to a list first when + it has batch dimensions. Default: ``True``. + as_linked_list (bool, optional): if ``True``, the list will be converted to a :class:`tensordict.utils.LinkedList` + which will automatically update the tensordict when the list is modified. Default: ``False``. + + Examples: + >>> from tensordict import NonTensorData + >>> import torch + >>> data = torch.stack([ + ... torch.stack([NonTensorData(data=(i, j), batch_size=[]) for i in range(2)]) + ... for j in range(3)]) + >>> data.tolist() + [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(0, 2), (1, 2)]] + + """ + iterator = self.tensordicts if self.stack_dim == 0 else self.unbind(0) + result = [ + td.tolist( + convert_tensors=convert_tensors, + tolist_first=tolist_first, + as_linked_list=as_linked_list, + ) + for td in iterator + ] + if as_linked_list: + return LinkedList(result, td=self) + return result + + def maybe_to_stack(self): + """Placeholder for interchangeability between stack and non-stack of non-tensors.""" + return type(self)( + *[ntd.maybe_to_stack() for ntd in self.tensordicts], + stack_dim=self.stack_dim, + ) + + @classmethod + def from_list(cls, non_tensors: List[Any]): + # Use local function because refers to cls + def _maybe_from_list(nontensor): + if isinstance(nontensor, list): + return cls.from_list(nontensor) + if is_non_tensor(nontensor): + return nontensor + return NonTensorData(nontensor) + + return cls(*[_maybe_from_list(nontensor) for nontensor in non_tensors]) + + def is_empty(self) -> bool: + return False + + _stack_non_tensor = NonTensorData._stack_non_tensor + + @classmethod + def from_nontensordata(cls, non_tensor: NonTensorData): + data = non_tensor.data + prev = NonTensorData(data, batch_size=[], device=non_tensor.device) + for dim in reversed(non_tensor.shape): + prev = cls(*[prev.clone(False) for _ in range(dim)], stack_dim=0) + return prev + + def __repr__(self): + selfrepr = str(self.tolist()) + if len(selfrepr) > 50: + selfrepr = f"{selfrepr[:50]}..." + selfrepr = indent(selfrepr, prefix=4 * " ") + batch_size = indent(f"batch_size={self.batch_size}", prefix=4 * " ") + device = indent(f"device={self.device}", prefix=4 * " ") + return f"NonTensorStack(\n{selfrepr}," f"\n{batch_size}," f"\n{device})" + + @classmethod + def lazy_stack( + cls, + items: Sequence[TensorCollection], + dim: int = 0, + *, + device: DeviceType | None = None, + out: T | None = None, + stack_dim_name: str | None = None, + **kwargs, + ) -> T: + result = super().lazy_stack( + items=items, + dim=dim, + out=out, + stack_dim_name=stack_dim_name, + device=device, + **kwargs, + ) + if not isinstance(result, cls): + raise RuntimeError( + f"Unexpected result type: {type(result)} - expected one of {cls}." + ) + return result + + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ) -> dict[str, Any]: + return self.tolist( + convert_tensors=convert_tensors, + tolist_first=tolist_first, + as_linked_list=as_linked_list, + ) + + def to_tensordict(self, *, retain_none: bool | None = None): + return self + + def _memmap_( + self, + *, + prefix: str | None = None, + copy_existing: bool = False, + executor=None, + futures=None, + inplace=True, + like=False, + memmaped: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key, + ) -> T: + + memmaped_leaves = memmaped + if not memmaped and prefix is not None: + memmaped_leaves = True + + def save_metadata(prefix=prefix, self=self): + data = self.tolist() + device = str(self.device) if self.device is not None else None + if not prefix.exists(): + os.makedirs(prefix, exist_ok=True) + jsondict = { + "_type": str(type(self)), + "stack_dim": self.stack_dim, + "device": device, + } + if _is_json_serializable(data): + jsondict["data"] = data + else: + jsondict["data"] = "pickle.pkl" + with open(prefix / "pickle.pkl", "wb") as f: + pickle.dump(data, f) + with open(prefix / "meta.json", "wb") as f: + from tensordict.utils import json_dumps + + json_str = json_dumps(jsondict, separators=(",", ":")) + # Ensure we write bytes to the binary file + if isinstance(json_str, str): + f.write(json_str.encode("utf-8")) + else: + f.write(json_str) + + if executor is None: + save_metadata() + else: + futures.append(executor.submit(save_metadata)) + # The leaves are all non-tensor or non-tensor stacks, and we already saved this on disk + # The only thing remaining to do is share the data between processes + results = [] + for i, td in enumerate(self.tensordicts): + td: NonTensorData + results.append( + td._memmap_( + prefix=(prefix / str(i)) if prefix is not None else None, + copy_existing=copy_existing, + executor=executor, + futures=futures, + inplace=inplace, + like=like, + # tell the nested stack / nontensor that + # no memmapping should be executed + memmaped=memmaped_leaves, + share_non_tensor=share_non_tensor, + existsok=existsok, + robust_key=robust_key, + ) + ) + if not inplace: + results = self.lazy_stack(results, dim=self.stack_dim) + else: + results = self + if not memmaped and prefix is not None: + results.__dict__["_path_to_memmap"] = prefix + return results + + @classmethod + def _load_memmap( + cls, prefix: str, metadata: dict, *, out=None, robust_key, **kwargs + ) -> LazyStackedTensorDict: + data = metadata.get("data") + if data is not None: + if isinstance(data, str): + with open(prefix / data, "rb") as file: + data = pickle.load(file) + device = metadata["device"] + if device is not None: + device = torch.device(device) + return cls._from_list(data, device=device) + return super()._load_memmap( + prefix=prefix, metadata=metadata, robust_key=robust_key, **kwargs + ) + + @classmethod + def _from_list(cls, datalist: List, device: torch.device, ndim: int | None = None): + if ( + all(isinstance(item, list) for item in datalist) + and all(len(item) == len(datalist[0]) for item in datalist) + and (ndim is None or ndim > 1) + ): + ndim = ndim - 1 if ndim is not None else None + return NonTensorStack( + *(cls._from_list(item, device=device, ndim=ndim) for item in datalist), + stack_dim=0, + ) + return NonTensorStack( + *( + NonTensorData(data=item, device=device, batch_size=torch.Size([])) + for item in datalist + ), + stack_dim=0, + ) + + def densify(self, layout: torch.layout = torch.strided): + # No need to do anything with a non tensor stack + return self + + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> T: + return self._update( + input_dict_or_td=input_dict_or_td, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> T: + return self._update( + input_dict_or_td=input_dict_or_td, + clone=clone, + inplace=True, + keys_to_update=keys_to_update, + ) + + def _update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + keys_to_update: Sequence[NestedKey] | None = None, + break_on_memmap: bool | None = None, + non_blocking: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> T: + if inplace and self.is_locked and not (self._is_shared or self._is_memmap): + raise RuntimeError(_LOCK_ERROR) + + if isinstance(input_dict_or_td, NonTensorData): + datalist = input_dict_or_td.data + for d in reversed(self.batch_size): + datalist = [datalist] * d + reconstructed = self._from_list( + datalist, device=self.device, ndim=self.ndim + ) + return self.update( + reconstructed, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + + memmap = False + if self._is_memmap and hasattr(self, "_path_to_memmap"): + if break_on_memmap is None: + global _BREAK_ON_MEMMAP + break_on_memmap = _BREAK_ON_MEMMAP + if not break_on_memmap: + raise RuntimeError( + "Calling _update with break_on_memmap=False is not permitted if the stack has a path." + ) + # this is the only way break_on_memmap is False + break_on_memmap = False + # remove memmap + if self._path_to_memmap.exists(): + shutil.rmtree(self._path_to_memmap) + memmap = True + if is_tensorclass(input_dict_or_td): + input_dict_or_td = input_dict_or_td._tensordict + + # update content + if isinstance(input_dict_or_td, NonTensorStack): + for leaf_dest, leaf_src in _zip_strict( + self.tensordicts, input_dict_or_td.unbind(self.stack_dim) + ): + leaf_dest._update( + leaf_src, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + break_on_memmap=break_on_memmap, + is_leaf=is_leaf, + update_batch_size=update_batch_size, + ignore_lock=ignore_lock, + ) + if memmap: + self._memmap_( + prefix=self._path_to_memmap, inplace=True, robust_key=True + ) + else: + raise NotImplementedError( + f"The data type {type(input_dict_or_td)} is not supported within {type(self).__name__}.update" + ) + return self + + def __setitem__(self, index: IndexType, value: Any): + memmap = False + if self._is_memmap and hasattr(self, "_path_to_memmap"): + global _BREAK_ON_MEMMAP + _BREAK_ON_MEMMAP = False + memmap = True + try: + if not is_tensor_collection(value): + if isinstance(value, list): + value = NonTensorStack(*value) + else: + value = NonTensorData(value) + super().__setitem__(index, value) + if memmap: + self._memmap_( + prefix=self._path_to_memmap, inplace=True, robust_key=True + ) + finally: + _BREAK_ON_MEMMAP = True + + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorCollection, + index: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + ) -> T: + memmap = False + if self._is_memmap and hasattr(self, "_path_to_memmap"): + global _BREAK_ON_MEMMAP + _BREAK_ON_MEMMAP = False + memmap = True + try: + super().update_at_( + input_dict_or_td, index, clone=clone, non_blocking=non_blocking + ) + if memmap: + self._memmap_( + prefix=self._path_to_memmap, inplace=True, robust_key=True + ) + finally: + _BREAK_ON_MEMMAP = True + return self + + @property + def data(self) -> Self: + """Attempts to return the unique value in the stack. + + Raises a ValueError if there is more than one unique value. + """ + try: + with set_capture_non_tensor_stack(True): + nt = NonTensorData._stack_non_tensor( + self.tensordicts, raise_if_non_unique=True + ) + if not isinstance(nt, NonTensorData): + raise ValueError + return nt.data + except ValueError: + raise AttributeError( + "Cannot get the non-unique data of a NonTensorStack. Use .tolist() instead." + ) + + +_register_tensor_class(NonTensorStack) + + +def _share_memory_nontensor(data, manager: Manager): + if isinstance(data, int): + return mp.Value(ctypes.c_int, data) + if isinstance(data, float): + return mp.Value(ctypes.c_double, data) + if isinstance(data, bool): + return mp.Value(ctypes.c_bool, data) + if isinstance(data, bytes): + return mp.Value(ctypes.c_byte, data) + if isinstance(data, dict): + result = manager.dict() + result.update(data) + return result + if isinstance(data, str): + result = mp.Array(ctypes.c_char, 100) + data = data.encode("utf-8") + result[: len(data)] = data + return result + if isinstance(data, list): + result = manager.list() + result.extend(data) + return result + # In all other cases, we just return the tensor. It's ok because the content + # will be passed to the remote process using regular serialization. We will + # lock the update in _update_shared_nontensor though. + return data + + +def _from_shared_nontensor(nontensor): + if isinstance(nontensor, multiprocessing.managers.ListProxy): + return list(nontensor) + if isinstance(nontensor, multiprocessing.managers.DictProxy): + return dict(nontensor) + if isinstance(nontensor, multiprocessing.sharedctypes.Synchronized): + return nontensor.value + if isinstance(nontensor, multiprocessing.sharedctypes.SynchronizedArray): + byte_list = [] + for byte in nontensor: + if byte == b"\x00": + break + byte_list.append(byte) + return b"".join(byte_list).decode("utf-8") + return nontensor + + +def _update_shared_nontensor(nontensor, val): + if isinstance(nontensor, multiprocessing.managers.ListProxy): + nontensor[:] = [] + nontensor.extend(val) + elif isinstance(nontensor, multiprocessing.managers.DictProxy): + nontensor.clear() + nontensor.update(val) + elif isinstance(nontensor, multiprocessing.sharedctypes.Synchronized): + nontensor.value = val + elif isinstance(nontensor, multiprocessing.sharedctypes.SynchronizedArray): + val = val.encode("utf-8") + for i, byte in enumerate(nontensor): + if i < len(val): + v = val[i] + nontensor[i] = v + elif byte == b"\x00": + break + else: + nontensor[i] = b"\x00" + # nontensor[0] = val.encode("utf-8") + else: + raise NotImplementedError( + f"Updating {type(nontensor).__name__} within a shared/memmaped structure is not supported." + ) diff --git a/lib/python3.12/site-packages/tensordict/tensorclass.pyi b/lib/python3.12/site-packages/tensordict/tensorclass.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a044079c69ae9fe5918a42dbc0b96b774a72e39b --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/tensorclass.pyi @@ -0,0 +1,1550 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import abc +import dataclasses +import enum +import types +from abc import abstractmethod +from collections.abc import MutableMapping +from pathlib import Path +from typing import ( + Any, + Callable, + dataclass_transform, + Generator, + Iterator, + Literal, + OrderedDict, + overload, + Sequence, + Tuple, + Type, + TYPE_CHECKING, + TypeVar, +) + +import numpy as np + +import torch +import torch.distributed as dist +from _typeshed import Incomplete +from tensordict import TensorDictBase +from tensordict._contextlib import LAST_OP_MAPS as LAST_OP_MAPS +from tensordict._nestedkey import NestedKey +from tensordict._tensorcollection import TensorCollection +from tensordict.memmap import MemoryMappedTensor as MemoryMappedTensor +from tensordict.utils import ( + Buffer as Buffer, + cache as cache, + convert_ellipsis_to_idx as convert_ellipsis_to_idx, + DeviceType as DeviceType, + erase_cache as erase_cache, + implement_for as implement_for, + IndexType as IndexType, + infer_size_impl as infer_size_impl, + int_generator as int_generator, + is_namedtuple as is_namedtuple, + is_namedtuple_class as is_namedtuple_class, + lazy_legacy as lazy_legacy, + lock_blocked as lock_blocked, + prod as prod, + set_lazy_legacy as set_lazy_legacy, + strtobool as strtobool, + TensorDictFuture as TensorDictFuture, + unravel_key as unravel_key, + unravel_key_list as unravel_key_list, +) +from torch import multiprocessing as mp, nn, Tensor + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + +class _NoDefault(enum.IntEnum): + ZERO = 0 + +NO_DEFAULT: Incomplete + +class _BEST_ATTEMPT_INPLACE: + def __bool__(self) -> bool: ... + +BEST_ATTEMPT_INPLACE: Incomplete +from typing import TypeAlias + +CompatibleType: TypeAlias = Tensor + +T = TypeVar("T", bound="TensorDictBase") +# Use Any for methods that don't have T in their parameters but return T +T_Any = TypeVar("T_Any", bound="TensorDictBase") + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + +class TensorClass: + _autocast: bool = False + _nocast: bool = False + _frozen: bool = False + + @overload + def __class_getitem__(cls, item: Literal["autocast"]) -> Type["TensorClass"]: ... + @overload + def __class_getitem__(cls, item: Literal["nocast"]) -> Type["TensorClass"]: ... + @overload + def __class_getitem__(cls, item: Literal["frozen"]) -> Type["TensorClass"]: ... + @overload + def __class_getitem__(cls, item: Literal["tensor_only"]) -> Type["TensorClass"]: ... + @overload + def __class_getitem__(cls, item: Literal["shadow"]) -> Type["TensorClass"]: ... + @overload + def __class_getitem__( + cls, + item: tuple[ + Literal["autocast", "nocast", "frozen", "tensor_only", "shadow"], ... + ], + ) -> Type["TensorClass"]: ... + def __init__( + self, + *args, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: DeviceType | None = None, + names: Sequence[str] | None = None, + non_blocking: bool | None = None, + lock: bool = False, + **kwargs, + ) -> None: ... + @property + def is_meta(self) -> bool: ... + def __bool__(self) -> bool: ... + def __ne__(self, other: object) -> Self: ... + def __xor__(self, other: TensorCollection | float) -> Self: ... + def __or__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __eq__(self, other: object) -> Self: ... + def __ge__(self, other: object) -> Self: ... + def __gt__(self, other: object) -> Self: ... + def __le__(self, other: object) -> Self: ... + def __lt__(self, other: object) -> Self: ... + def __deepcopy__(self, memodict={}) -> Self: ... + def __iter__(self) -> Generator: ... + def __len__(self) -> int: ... + def __contains__(self, key: NestedKey) -> bool: ... + def __getitem__( + self, index: IndexType + ) -> Self | Tensor | TensorCollection | Any: ... + __getitems__ = __getitem__ + + def __setitem__(self, index: IndexType, value: Any) -> None: ... + def __delitem__(self, key: NestedKey) -> Any: ... + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: ... + def all(self, dim: int | None = None) -> bool | TensorDictBase: ... + def any(self, dim: int | None = None) -> bool | TensorDictBase: ... + def isfinite(self) -> Any: ... + def isnan(self) -> Any: ... + def isneginf(self) -> Any: ... + def isposinf(self) -> Any: ... + def isreal(self) -> Any: ... + @overload + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + ) -> Self: ... + @overload + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + def amin( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + @overload + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + def min( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + @overload + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + ) -> Self: ... + @overload + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + ) -> Self | torch.Tensor: ... + def amax( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + return_indices: bool = True, + ) -> Self: ... + @overload + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + def max( + self, + dim: int | NO_DEFAULT = ..., + keepdim: bool = False, + *, + reduce: bool | None = None, + return_indices: bool = True, + ) -> Self | torch.Tensor: ... + @overload + def cummin(self, dim: int, *, return_indices: bool = True) -> Self: ... + @overload + def cummin( + self, dim: int, *, reduce: bool, return_indices: bool = True + ) -> Self | torch.Tensor: ... + def cummin( + self, dim: int, *, reduce: bool | None = None, return_indices: bool = True + ) -> Self | torch.Tensor: ... + @overload + def cummax(self, dim: int, *, return_indices: bool = True) -> Self: ... + @overload + def cummax( + self, dim: int, *, reduce: bool, return_indices: bool = True + ) -> Self | torch.Tensor: ... + def cummax( + self, dim: int, *, reduce: bool | None = None, return_indices: bool = True + ) -> Self | torch.Tensor: ... + @overload + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def mean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def nanmean( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def prod( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def sum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + ) -> Self: ... + @overload + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool, + ) -> Self | torch.Tensor: ... + def nansum( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + dtype: torch.dtype | None = None, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + ) -> Self: ... + @overload + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + def std( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + ) -> Self: ... + @overload + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool, + ) -> Self | torch.Tensor: ... + def var( + self, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + correction: int = 1, + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + ) -> Self: ... + @overload + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + reduce: bool, + ) -> Self | torch.Tensor: ... + def quantile( + self, + q: float | torch.Tensor, + dim: int | tuple[int] = ..., + keepdim: bool = ..., + *, + interpolation: str = "linear", + reduce: bool | None = None, + ) -> Self | torch.Tensor: ... + def auto_batch_size_(self, batch_dims: int | None = None) -> Self: ... + def auto_device_(self) -> Self: ... + @classmethod + def from_dataclass( + cls, dataclass, *, auto_batch_size: bool = False, as_tensorclass: bool = False + ) -> Self: ... + @classmethod + def from_any(cls, obj, *, auto_batch_size: bool = False) -> Self: ... + @classmethod + def from_dict( + cls, + input_dict, + *, + batch_size: Sequence[int] | torch.Size | int | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: list[str] | None = None, + ) -> Self: ... + def from_dict_instance( + self, + input_dict, + batch_size: Incomplete | None = None, + device: Incomplete | None = None, + batch_dims: Incomplete | None = None, + names: list[str] | None = None, + ) -> Self: ... + @classmethod + def from_pytree( + cls, + pytree, + *, + batch_size: torch.Size | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + ) -> Self: ... + def to_pytree(self) -> Self: ... + @classmethod + def from_h5(cls, filename, mode: str = "r") -> Self: ... + @classmethod + def from_module( + cls, + module, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + ) -> Self: ... + @classmethod + def from_modules( + cls, + *modules, + as_module: bool = False, + lock: bool = True, + use_state_dict: bool = False, + lazy_stack: bool = False, + expand_identical: bool = False, + ) -> Self: ... + def to_module( + self, + module: nn.Module, + *, + inplace: bool | None = None, + return_swap: bool = True, + swap_dest: Incomplete | None = None, + use_state_dict: bool = False, + non_blocking: bool = False, + memo: Incomplete | None = None, + ) -> Self: ... + @property + def shape(self) -> torch.Size: ... + @shape.setter + def shape(self, value) -> torch.Size: ... + @property + def batch_size(self) -> torch.Size: ... + @batch_size.setter + def batch_size(self, new_size: Sequence[int] | torch.Size) -> None: ... + def size(self, dim: int | None = None) -> torch.Size | int: ... + @property + def data(self) -> Self: ... + @property + def grad(self) -> Self: ... + @grad.setter + def grad(self, grad) -> None: ... + def data_ptr(self, *, storage: bool = False) -> int: ... + def zero_grad(self, set_to_none: bool = True) -> Self: ... + @property + def dtype(self) -> torch.dtype: ... + @property + def batch_dims(self) -> int: ... + @batch_dims.setter + def batch_dims(self, value: int) -> None: ... + def ndimension(self) -> int: ... + @property + def ndim(self) -> int: ... + def dim(self) -> int: ... + def numel(self) -> int: ... + @property + def depth(self) -> int: ... + @overload + def expand(self, *shape: int) -> Self: ... + @overload + def expand(self, shape: torch.Size | Sequence[int]) -> Self: ... + def expand_as(self, other: TensorCollection | torch.Tensor) -> Self: ... + def new_zeros( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ) -> Self: ... + def new_ones( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ) -> Self: ... + def new_empty( + self, + *size: torch.Size, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ) -> Self: ... + def new_full( + self, + size: torch.Size, + fill_value, + *, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + layout: torch.layout = ..., + pin_memory: bool | None = None, + empty_lazy: bool = False, + ) -> Self: ... + def new_tensor( + self, + data: torch.Tensor | TensorDictBase, + *, + dtype: torch.dtype = None, + device: DeviceType = ..., + requires_grad: bool = False, + pin_memory: bool | None = None, + ) -> Self: ... + def unbind(self, dim: int) -> tuple[T, ...]: ... + def chunk(self, chunks: int, dim: int = 0) -> tuple[Self, ...]: ... + def unsqueeze(self, dim: int) -> Self: ... + def squeeze(self, dim: int | None = None) -> Self: ... + @overload + def reshape(self, *shape: int) -> Self: ... + @overload + def reshape(self, shape: Sequence[int] | torch.Size) -> Self: ... + def repeat_interleave( + self, + repeats: torch.Tensor | int, + dim: int | None = None, + *, + output_size: int | None = None, + ) -> Self: ... + def repeat(self, *repeats: int) -> Self: ... + def cat_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: ... + def stack_tensors( + self, + *keys: NestedKey, + out_key: NestedKey, + dim: int = 0, + keep_entries: bool = False, + ) -> Self: ... + def cat_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | list[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: ... + def stack_from_tensordict( + self, + dim: int = 0, + *, + sorted: bool | list[NestedKey] | None = None, + out: torch.Tensor | None = None, + ) -> torch.Tensor: ... + @classmethod + def stack(cls, input, dim: int = 0, *, out: Incomplete | None = None) -> Self: ... + @classmethod + def cat(cls, input, dim: int = 0, *, out: Incomplete | None = None) -> Self: ... + @classmethod + def lazy_stack( + cls, input, dim: int = 0, *, out: Incomplete | None = None, **kwargs + ) -> Self: ... + @classmethod + def maybe_dense_stack( + cls, input, dim: int = 0, *, out: Incomplete | None = None, **kwargs + ) -> Self: ... + def split(self, split_size: int | list[int], dim: int = 0) -> list[Self]: ... + def gather(self, dim: int, index: Tensor, out: T | None = None) -> Self: ... + @overload + def view(self, *shape: int) -> Self: ... + @overload + def view(self, dtype) -> Self: ... + @overload + def view(self, shape: torch.Size) -> Self: ... + def view( + self, + *shape: int, + size: Sequence[int] | torch.Size | None = None, + batch_size: torch.Size | None = None, + ) -> Self: ... + def transpose(self, dim0, dim1) -> Self: ... + def swapaxes(self, axis0: int, axis1: int) -> Self: ... + def swapdims(self, dim0: int, dim1: int) -> Self: ... + def flip(self, dims: int | tuple[int, ...]) -> Self: ... + def fliplr(self) -> Self: ... + def flipud(self) -> Self: ... + def roll( + self, shifts: int | tuple[int, ...], dims: int | tuple[int, ...] | None = None + ) -> Self: ... + def rot90(self, k: int = 1, dims: tuple[int, int] = (0, 1)) -> Self: ... + def narrow(self, dim: int, start: int, length: int) -> Self: ... + def tile(self, dims: tuple[int, ...]) -> Self: ... + def broadcast_to(self, shape: tuple[int, ...]) -> Self: ... + def atleast_1d(self) -> Self: ... + def atleast_2d(self) -> Self: ... + def atleast_3d(self) -> Self: ... + def movedim( + self, source: int | tuple[int, ...], destination: int | tuple[int, ...] + ) -> Self: ... + def moveaxis( + self, source: int | tuple[int, ...], destination: int | tuple[int, ...] + ) -> Self: ... + @overload + def permute(self, *dims: int) -> Self: ... + @overload + def permute(self, dims: Sequence[int]) -> Self: ... + @property + def names(self) -> list[str]: ... + @names.setter + def names(self, value: Sequence[str]) -> None: ... + def refine_names(self, *names: str) -> Self: ... + def rename(self, *names: str, **rename_map: str) -> Self: ... + def rename_(self, *names: str, **rename_map: str) -> Self: ... + @property + def device(self) -> torch.device | None: ... + @device.setter + def device(self, value: DeviceType) -> torch.device | None: ... + def clear(self) -> Self: ... + def clear_refs_for_compile_(self) -> Self: ... + @classmethod + def fromkeys(cls, keys: list[NestedKey], value: Any = 0) -> Self: ... + def popitem(self) -> tuple[NestedKey, CompatibleType]: ... + def clear_device_(self) -> Self: ... + def param_count(self, *, count_duplicates: bool = True) -> int: ... + def bytes(self, *, count_duplicates: bool = True) -> int: ... + def pin_memory( + self, num_threads: int | None = None, inplace: bool = False + ) -> Self: ... + def pin_memory_(self, num_threads: int | str = 0) -> Self: ... + def cpu(self, **kwargs) -> Self: ... + def cuda(self, device: int | None = None, **kwargs) -> Self: ... + @property + def is_cuda(self) -> bool: ... + @property + def is_cpu(self) -> bool: ... + def state_dict( + self, + destination: Incomplete | None = None, + prefix: str = "", + keep_vars: bool = False, + flatten: bool = False, + ) -> OrderedDict[str, Any]: ... + def load_state_dict( + self, + state_dict: OrderedDict[str, Any], + strict: bool = True, + assign: bool = False, + from_flatten: bool = False, + ) -> Self: ... + def is_shared(self) -> bool: ... + def is_memmap(self) -> bool: ... + def share_memory_(self) -> Self: ... + def densify(self, layout: torch.layout = ...) -> Self: ... + @property + def saved_path(self) -> Self: ... + def consolidate( + self, + filename: Path | str | None = None, + *, + num_threads: int = 0, + device: torch.device | None = None, + non_blocking: bool = False, + inplace: bool = False, + return_early: bool = False, + use_buffer: bool = False, + share_memory: bool = False, + pin_memory: bool = False, + metadata: bool = False, + ) -> Self: ... + @classmethod + def from_consolidated(cls, filename) -> Self: ... + def is_consolidated(self) -> bool: ... + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: ... + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def make_memmap_from_tensor( + self, + key: NestedKey, + tensor: torch.Tensor, + *, + copy_data: bool = True, + robust_key: bool | None = None, + ) -> MemoryMappedTensor: ... + def save( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + def dumps( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + def memmap( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + existsok: bool = True, + robust_key: bool | None = None, + ) -> Self: ... + def memmap_like( + self, + prefix: str | None = None, + copy_existing: bool = False, + *, + existsok: bool = True, + num_threads: int = 0, + return_early: bool = False, + share_non_tensor: bool = False, + robust_key: bool | None = None, + ) -> Self: ... + @classmethod + def load(cls, prefix: str | Path, *args, **kwargs) -> Self: ... + def load_(self, prefix: str | Path, *args, **kwargs) -> Self: ... + @classmethod + def load_memmap( + cls, + prefix: str | Path, + device: torch.device | None = None, + non_blocking: bool = False, + *, + out: TensorCollection | None = None, + robust_key: bool | None = None, + ) -> Self: ... + def load_memmap_( + self, prefix: str | Path, robust_key: bool | None = None + ) -> Self: ... + def memmap_refresh_(self) -> Self: ... + def entry_class(self, key: NestedKey) -> type: ... + def set( + self, + key: NestedKey, + item: CompatibleType, + inplace: bool = False, + *, + non_blocking: bool = False, + **kwargs: Any, + ) -> Self: ... + def set_non_tensor(self, key: NestedKey, value: Any) -> Self: ... + def get_non_tensor(self, key: NestedKey, default=...) -> CompatibleType: ... + def filter_non_tensor_data(self) -> Self: ... + def filter_empty_(self) -> Self: ... + def set_at_( + self, + key: NestedKey, + value: CompatibleType, + index: IndexType, + *, + non_blocking: bool = False, + ) -> Self: ... + def set_( + self, key: NestedKey, item: CompatibleType, *, non_blocking: bool = False + ) -> Self: ... + @overload + def get(self, key: NestedKey) -> CompatibleType: ... + @overload + def get(self, key: NestedKey, default: CompatibleType | Any) -> CompatibleType: ... + @overload + def get(self, key: NestedKey, *args, **kwargs) -> CompatibleType: ... + @overload + def get( + self, + key: NestedKey, + *, + as_list: bool = False, + as_padded_tensor: bool = False, + as_nested_tensor: bool = False, + padding_side: str = "right", + layout: torch.layout | None = None, + padding_value: float | int | bool = 0.0, + **kwargs, + ) -> CompatibleType: ... + @overload + def get_at(self, key: NestedKey, index: IndexType) -> CompatibleType: ... + @overload + def get_at( + self, key: NestedKey, index: IndexType, default: CompatibleType | Any + ) -> CompatibleType: ... + def get_at( + self, + key: NestedKey, + *args, + **kwargs, + ) -> CompatibleType: ... + def get_item_shape(self, key: NestedKey) -> torch.Size: ... + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> Self: ... + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: ... + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + idx: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> Self: ... + def replace(self, *args, **kwargs) -> Self: ... + def create_nested(self, key: NestedKey) -> Self: ... + def copy_(self, tensordict: T, non_blocking: bool = False) -> Self: ... + def copy_at_( + self, tensordict: T, idx: IndexType, non_blocking: bool = False + ) -> Self: ... + def is_empty(self) -> bool: ... + def setdefault( + self, key: NestedKey, default: CompatibleType, inplace: bool = False + ) -> CompatibleType: ... + def items( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Incomplete | None = None, + *, + sort: bool = False, + ) -> Iterator[tuple[str, CompatibleType]]: ... + def non_tensor_items( + self, include_nested: bool = False + ) -> Iterator[tuple[str, Any]]: ... + def values( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Incomplete | None = None, + *, + sort: bool = False, + ) -> Iterator[CompatibleType]: ... + def keys( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[type], bool] | None = None, + *, + sort: bool = False, + ) -> Iterator[NestedKey]: ... + def pop(self, key: NestedKey, default: Any = ...) -> CompatibleType: ... + @property + def sorted_keys(self) -> list[NestedKey]: ... + def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Self: ... + def unflatten(self, dim: int, unflattened_size: torch.Size) -> Self: ... + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> Self: ... + def del_(self, key: NestedKey) -> Self: ... + def gather_and_stack( + self, dst: int, group: "dist.ProcessGroup" | None = None + ) -> Self | None: ... + def send( + self, + dst: int, + *, + group: dist.ProcessGroup | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> None: ... + def recv( + self, + src: int, + *, + group: dist.ProcessGroup | None = None, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> int: ... + @classmethod + def from_remote_init( + cls: T, + src: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ) -> T: ... + def init_remote( + self, + dst: int, + group: "ProcessGroup" | None = None, # noqa: F821 + device: torch.device | None = None, + ) -> Self: ... + def isend( + self, + dst: int, + *, + group: "dist.ProcessGroup" | None = None, # noqa: F821 + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> int: ... + def irecv( + self, + src: int, + *, + group: dist.ProcessGroup | None = None, + return_premature: bool = False, + init_tag: int = 0, + pseudo_rand: bool = False, + ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None: ... + def reduce( + self, + dst, + op: Incomplete | None = None, + async_op: bool = False, + return_premature: bool = False, + group: Incomplete | None = None, + ) -> Self: ... + def apply_(self, fn: Callable, *others, **kwargs) -> Self: ... + def apply( + self, + fn: Callable, + *others: T, + batch_size: Sequence[int] | None = None, + device: torch.device | None = ..., + names: Sequence[str] | None = ..., + inplace: bool = False, + default: Any = ..., + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorCollection | None = None, + **constructor_kwargs, + ) -> Self | None: ... + def named_apply( + self, + fn: Callable, + *others: T, + nested_keys: bool = False, + batch_size: Sequence[int] | None = None, + device: torch.device | None = ..., + names: Sequence[str] | None = ..., + inplace: bool = False, + default: Any = ..., + filter_empty: bool | None = None, + propagate_lock: bool = False, + call_on_nested: bool = False, + out: TensorCollection | None = None, + **constructor_kwargs, + ) -> Self | None: ... + def map( + self, + fn: Callable[[TensorDictBase], TensorDictBase | None], + dim: int = 0, + num_workers: int | None = None, + *, + out: TensorCollection | None = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = False, + pbar: bool = False, + mp_start_method: str | None = None, + ) -> Self: ... + def map_iter( + self, + fn: Callable[[TensorDictBase], TensorDictBase | None], + dim: int = 0, + num_workers: int | None = None, + *, + shuffle: bool = False, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool | None = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + index_with_generator: bool = True, + pbar: bool = False, + mp_start_method: str | None = None, + ) -> Self: ... + def record_stream(self, stream: torch.cuda.Stream) -> Self: ... + def __add__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __iadd__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __abs__(self) -> Self: ... + def __truediv__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __itruediv__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __mod__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __mul__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __imul__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __sub__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __isub__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __pow__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def __ipow__(self, other: TensorCollection | torch.Tensor) -> Self: ... + def abs(self) -> Self: ... + def abs_(self) -> Self: ... + def acos(self) -> Self: ... + def acos_(self) -> Self: ... + def exp(self) -> Self: ... + def exp_(self) -> Self: ... + def neg(self) -> Self: ... + def neg_(self) -> Self: ... + def reciprocal(self) -> Self: ... + def reciprocal_(self) -> Self: ... + def sigmoid(self) -> Self: ... + def sigmoid_(self) -> Self: ... + def sign(self) -> Self: ... + def sign_(self) -> Self: ... + def sin(self) -> Self: ... + def sin_(self) -> Self: ... + def sinh(self) -> Self: ... + def sinh_(self) -> Self: ... + def tan(self) -> Self: ... + def tan_(self) -> Self: ... + def tanh(self) -> Self: ... + def tanh_(self) -> Self: ... + def trunc(self) -> Self: ... + def trunc_(self) -> Self: ... + def lgamma(self) -> Self: ... + def lgamma_(self) -> Self: ... + def frac(self) -> Self: ... + def frac_(self) -> Self: ... + def expm1(self) -> Self: ... + def expm1_(self) -> Self: ... + def log(self) -> Self: ... + def log_(self) -> Self: ... + def log10(self) -> Self: ... + def log10_(self) -> Self: ... + def log1p(self) -> Self: ... + def log1p_(self) -> Self: ... + def log2(self) -> Self: ... + def log2_(self) -> Self: ... + def ceil(self) -> Self: ... + def ceil_(self) -> Self: ... + def floor(self) -> Self: ... + def floor_(self) -> Self: ... + def round(self) -> Self: ... + def round_(self) -> Self: ... + def erf(self) -> Self: ... + def erf_(self) -> Self: ... + def erfc(self) -> Self: ... + def erfc_(self) -> Self: ... + def asin(self) -> Self: ... + def asin_(self) -> Self: ... + def atan(self) -> Self: ... + def atan_(self) -> Self: ... + def cos(self) -> Self: ... + def cos_(self) -> Self: ... + def cosh(self) -> Self: ... + def cosh_(self) -> Self: ... + def add( + self, + other: TensorCollection | torch.Tensor, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: ... + def add_( + self, other: TensorCollection | float, *, alpha: float | None = None + ) -> Self: ... + def lerp( + self, + end: TensorCollection | torch.Tensor, + weight: TensorCollection | torch.Tensor | float, + ) -> Self: ... + def lerp_( + self, end: TensorCollection | float, weight: TensorCollection | float + ) -> Self: ... + def addcdiv( + self, + other1: TensorCollection | torch.Tensor, + other2: TensorCollection | torch.Tensor, + value: float | None = 1, + ) -> Self: ... + def addcdiv_(self, other1, other2, *, value: float | None = 1) -> Self: ... + def addcmul(self, other1, other2, *, value: float | None = 1) -> Self: ... + def addcmul_(self, other1, other2, *, value: float | None = 1) -> Self: ... + def sub( + self, + other: TensorCollection | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: ... + def rsub( + self, + other: TensorCollection | float, + *, + alpha: float | None = None, + default: str | CompatibleType | None = None, + ) -> Self: ... + def sub_( + self, other: TensorCollection | float, alpha: float | None = None + ) -> Self: ... + def mod(self, other: TensorCollection | torch.Tensor) -> Self: ... + def mul_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def mul( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def maximum_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def maximum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def minimum_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def minimum( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def clamp( + self, + min: TensorCollection | torch.Tensor = None, + max: TensorCollection | torch.Tensor = None, + *, + out=None, + ) -> Self: ... + def logsumexp(self, dim=None, keepdim=False, *, out=None) -> Self: ... + def clamp_max_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def clamp_max( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def clamp_min_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def clamp_min( + self, + other: TensorCollection | torch.Tensor, + default: str | CompatibleType | None = None, + ) -> Self: ... + def pow_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def pow( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def div_(self, other: TensorCollection | torch.Tensor) -> Self: ... + def div( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def sqrt_(self) -> Self: ... + def sqrt(self) -> Self: ... + def __enter__(self) -> Any: ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: types.TracebackType | None, + ) -> None: ... + def select( + self, *keys: NestedKey, inplace: bool = False, strict: bool = True + ) -> Self: ... + def exclude(self, *keys: NestedKey, inplace: bool = False) -> Self: ... + def to_tensordict(self, *, retain_none: bool | None = None) -> Self: ... + def clone(self, recurse: bool = True, **kwargs) -> Self: ... + def copy(self) -> Self: ... + def to_padded_tensor( + self, padding: float = 0.0, mask_key: NestedKey | None = None + ) -> Self: ... + def as_tensor(self) -> Self: ... + def to_lazystack(self, dim: int = 0) -> Self: ... + def to_dict( + self, + *, + retain_none: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + ) -> dict[str, Any]: ... + def to_mds( + self, + *, + out: str | tuple[str, str], + columns: dict[str, str] | None = None, + writer: "MDSWriter" | None = None, + ) -> None: ... + @classmethod + def from_list( + cls, + input, + *, + auto_batch_size: bool | None = None, + batch_size: torch.Size | None = None, + device: torch.device | None = None, + batch_dims: int | None = None, + names: List[str] | None = None, + lazy: bool | None = None, + ) -> Self: ... + def tolist( + self, + *, + convert_nodes: bool = True, + convert_tensors: bool | Literal["numpy"] = False, + tolist_first: bool = False, + as_linked_list: bool = False, + ) -> list[Any]: ... + def numpy(self) -> np.ndarray | dict[str, Any]: ... + def to_namedtuple(self, dest_cls: type | None = None) -> Any: ... + @classmethod + def _from_tensordict( + cls, + tensordict: TensorCollection, + non_tensordict: dict | None = None, + safe: bool = True, + ) -> Self: ... + @classmethod + def from_tensordict( + cls, + tensordict: TensorCollection, + non_tensordict: dict | None = None, + safe: bool = True, + ) -> Self: ... + def _from_tensordict_with_copy(self, tensordict: TensorCollection) -> Self: ... + def _from_tensordict_with_none(self, tensordict: TensorCollection) -> Self: ... + @classmethod + def from_namedtuple(cls, named_tuple, *, auto_batch_size: bool = False) -> Self: ... + def from_tuple( + cls, + obj, + *, + auto_batch_size: bool = False, + batch_dims: int | None = None, + device: torch.device | None = None, + batch_size: torch.Size | None = None, + ) -> Self: ... + def logical_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + def bitwise_and( + self, + other: TensorCollection | torch.Tensor, + *, + default: str | CompatibleType | None = None, + ) -> Self: ... + @classmethod + def from_struct_array( + cls, struct_array: np.ndarray, device: torch.device | None = None + ) -> Self: ... + def to_struct_array(self) -> np.ndarray: ... + def to_h5(self, filename: str, **kwargs) -> Any: ... + def empty( + self, + recurse: bool = False, + *, + batch_size: Incomplete | None = None, + device=..., + names: Incomplete | None = None, + ) -> Self: ... + def zero_(self) -> Self: ... + def fill_(self, key: NestedKey, value: float | bool) -> Self: ... + def masked_fill_(self, mask: Tensor, value: float | bool) -> Self: ... + def masked_fill(self, mask: Tensor, value: float | bool) -> Self: ... + def where( + self, + condition, + other, + *, + out: Incomplete | None = None, + pad: Incomplete | None = None, + update_batch_size: bool = False, + ) -> Self: ... + def masked_select(self, mask: Tensor) -> Self: ... + def is_contiguous(self) -> bool: ... + def contiguous(self) -> Self: ... + def flatten_keys( + self, + separator: str = ".", + inplace: bool = False, + is_leaf: Callable[[type], bool] | None = None, + ) -> Self: ... + def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> Self: ... + def tensor_split( + self, + indices_or_sections: int | list[int] | tuple[int, ...] | torch.Tensor, + dim: int = 0, + ) -> tuple[Self, ...]: ... + def split_keys( + self, + *key_sets, + inplace: bool = False, + strict: bool = True, + reproduce_struct: bool = False, + ) -> tuple[Self, ...]: ... + def separates( + self, + *keys: NestedKey, + default: Any = NO_DEFAULT, + strict: bool = True, + filter_empty: bool = True, + ) -> Self: ... + def norm( + self, + *, + out=None, + dtype: torch.dtype | None = None, + ) -> Self: ... + def softmax(self, dim: int, dtype: torch.dtype | None = None) -> Self: ... + @property + def is_locked(self) -> bool: ... + @is_locked.setter + def is_locked(self, value: bool) -> None: ... + def lock_(self) -> Self: ... + def unlock_(self) -> Self: ... + @overload + def to( + self, + device: int | torch.device | None = ..., + dtype: torch.dtype | None = ..., + non_blocking: bool = ..., + inplace: bool = False, + ) -> Self: ... + @overload + def to(self, dtype: torch.dtype, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, tensor: Tensor, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, *, other: T, non_blocking: bool = ...) -> Self: ... + @overload + def to(self, *, batch_size: torch.Size) -> Self: ... + def to(self, *args, **kwargs) -> Self: ... + def is_floating_point(self) -> bool: ... + def double(self) -> Self: ... + def float(self) -> Self: ... + def int(self) -> Self: ... + def bool(self) -> Self: ... + def half(self) -> Self: ... + def type(self, dst_type: torch.dtype) -> Self: ... + @property + def requires_grad(self) -> bool: ... + def requires_grad_(self, requires_grad: bool = True) -> Self: ... + def detach_(self) -> Self: ... + def detach(self) -> Self: ... + def bfloat16(self) -> Self: ... + def complex128(self) -> Self: ... + def complex32(self) -> Self: ... + def complex64(self) -> Self: ... + def float16(self) -> Self: ... + def float32(self) -> Self: ... + def float64(self) -> Self: ... + def int16(self) -> Self: ... + def int32(self) -> Self: ... + def int64(self) -> Self: ... + def int8(self) -> Self: ... + def qint32(self) -> Self: ... + def qint8(self) -> Self: ... + def quint4x2(self) -> Self: ... + def quint8(self) -> Self: ... + def uint16(self) -> Self: ... + def uint32(self) -> Self: ... + def uint64(self) -> Self: ... + def uint8(self) -> Self: ... + +class NonTensorDataBase(TensorClass): ... +class NonTensorData(NonTensorDataBase): ... +class MetaData(NonTensorDataBase): ... +class NonTensorStack(TensorDictBase): ... + +@dataclass_transform() +def tensorclass( + cls: T = None, + /, + *, + autocast: bool = False, + frozen: bool = False, + nocast: bool = False, + shadow: bool = False, + tensor_only: bool = False, +) -> T: ... +def is_non_tensor(obj) -> bool: ... +def from_dataclass( + obj: Any, + *, + dest_cls: Type | None = None, + auto_batch_size: bool = False, + batch_dims: int | None = None, + batch_size: torch.Size | None = None, + frozen: bool = False, + autocast: bool = False, + nocast: bool = False, + inplace: bool = False, + shadow: bool = False, + tensor_only: bool = False, + device: torch.device | None = None, +) -> Any: ... diff --git a/lib/python3.12/site-packages/tensordict/tensordict.py b/lib/python3.12/site-packages/tensordict/tensordict.py new file mode 100644 index 0000000000000000000000000000000000000000..4b043f41687ce70dfbed80083d8f81ae164ea757 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/tensordict.py @@ -0,0 +1,36 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from tensordict._lazy import LazyStackedTensorDict # noqa: F401 +from tensordict._td import TensorDict # noqa: F401 +from tensordict.base import ( # noqa: F401 + is_tensor_collection, + NO_DEFAULT, + TensorDictBase, +) +from tensordict.functional import ( # noqa: F401 + dense_stack_tds, + make_tensordict, + merge_tensordicts, + pad, + pad_sequence, +) +from tensordict.memmap import MemoryMappedTensor # noqa: F401 +from tensordict.utils import ( # noqa: F401 + assert_allclose_td, + cache, + convert_ellipsis_to_idx, + erase_cache, + expand_as_right, + expand_right, + implement_for, + infer_size_impl, + int_generator, + is_nested_key, + is_seq_of_nested_key, + is_tensorclass, + lock_blocked, + NestedKey, +) diff --git a/lib/python3.12/site-packages/tensordict/utils.py b/lib/python3.12/site-packages/tensordict/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..db5bbfbf53e10ee06c42d3482cfd7d411b2e4463 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/utils.py @@ -0,0 +1,3241 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import collections +import concurrent.futures +import functools +import importlib.util +import itertools +import logging + +import math +import os +import re +import sys +import threading +import time +import warnings +import weakref +from collections import defaultdict +from collections.abc import KeysView +from contextlib import nullcontext +from functools import wraps +from numbers import Number +from textwrap import indent +from typing import ( + Any, + Callable, + Iterator, + List, + Sequence, + Tuple, + TYPE_CHECKING, + TypeVar, + Union, +) + +import numpy as np +import torch +from pyvers import get_backend, implement_for, register_backend, set_backend + +from tensordict._C import ( # noqa: F401 # @manual=//pytorch/tensordict:_C + _unravel_key_to_tuple as _unravel_key_to_tuple_cpp, + unravel_key as unravel_key_cpp, + unravel_key_list as unravel_key_list_cpp, + unravel_keys as unravel_keys_cpp, +) + +from tensordict._nestedkey import NestedKey + +from torch import Tensor +from torch._C import _disabled_torch_function_impl +from torch.nn.parameter import ( + UninitializedBuffer, + UninitializedParameter, + UninitializedTensorMixin, +) +from torch.utils._contextlib import _DecoratorContextManager +from torch.utils.data._utils.worker import _generate_state + +try: + from functorch import dim as ftdim + + _has_funcdim = True +except ImportError: + _has_funcdim = False +try: + from torch.compiler import assume_constant_result, is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import assume_constant_result, is_compiling + +if TYPE_CHECKING: + from tensordict.base import TensorDictBase + from tensordict.tensorclass import NonTensorStack + +try: + from dataclasses import GenericAlias +except ImportError: + # python < 3.9 + class GenericAlias: + """Placeholder.""" + + ... + + +try: + try: + from torch._C._functorch import ( # @manual=fbcode//caffe2:torch + get_unwrapped, + is_batchedtensor, + ) + except ImportError: + from functorch._C import ( # @manual=fbcode//caffe2/functorch:_C # noqa + get_unwrapped, + is_batchedtensor, + ) +except ImportError: + pass + + +# Utility function to wrap C++ functorch functions for torch.compile support +def _wrap_functorch_function(func): + """Wrap a functorch C++ function to make it compatible with torch.compile. + + PyTorch's Dynamo compiler cannot trace C++ functions from torch._C._functorch. + This wrapper uses torch.compiler.allow_in_graph to allow compilation through + these functions, with a fallback for older PyTorch versions. + + Args: + func: The C++ function to wrap + + Returns: + The wrapped function, or the original if wrapping is not available + """ + try: + from torch.compiler import allow_in_graph + + return allow_in_graph(func) + except (ImportError, AttributeError): + # Fallback for older PyTorch versions without allow_in_graph + return func + + +def _import_and_wrap_functorch(*names, wrap_names=None): + """Import functorch C++ functions and optionally wrap them for torch.compile. + + This utility handles the common pattern of importing functions from + torch._C._functorch and wrapping those that need to be compatible with + torch.compile. + + Args: + *names: Names of functions to import from torch._C._functorch + wrap_names: Optional list/set of function names to wrap with allow_in_graph. + If None, wraps all functions that start with '_' (private functions). + + Returns: + Tuple of imported (and optionally wrapped) functions in the same order as names + + Example: + >>> _add_batch_dim, _remove_batch_dim, is_batchedtensor = _import_and_wrap_functorch( + ... '_add_batch_dim', '_remove_batch_dim', 'is_batchedtensor' + ... ) + """ + from torch._C._functorch import ( # @manual=fbcode//caffe2:torch + _add_batch_dim as _functorch_add_batch_dim, + _remove_batch_dim as _functorch_remove_batch_dim, + get_unwrapped as _functorch_get_unwrapped, + is_batchedtensor as _functorch_is_batchedtensor, + ) + + # Map of available functions + available_funcs = { + "_add_batch_dim": _functorch_add_batch_dim, + "_remove_batch_dim": _functorch_remove_batch_dim, + "get_unwrapped": _functorch_get_unwrapped, + "is_batchedtensor": _functorch_is_batchedtensor, + } + + # Determine which functions to wrap + if wrap_names is None: + # By default, wrap all private functions (starting with '_') + wrap_names = {name for name in names if name.startswith("_")} + else: + wrap_names = set(wrap_names) + + # Get and optionally wrap the requested functions + result = [] + for name in names: + if name not in available_funcs: + raise ValueError(f"Function '{name}' not available in torch._C._functorch") + func = available_funcs[name] + if name in wrap_names: + func = _wrap_functorch_function(func) + result.append(func) + + return tuple(result) if len(result) > 1 else result[0] + + +# Import and wrap _add_batch_dim for compilation support +_add_batch_dim_c = None +try: + _add_batch_dim_c = _import_and_wrap_functorch("_add_batch_dim") +except ImportError: + pass + + +if not _has_funcdim: + + class _ftdim_mock: + class Dim: + pass + + class Tensor: + pass + + def dims(self, *args, **kwargs): + raise ImportError("functorch.dim not found") + + ftdim = _ftdim_mock # noqa: F811 + +T = TypeVar("T", bound="TensorDictBase") + +_PIN_MEM_TIMEOUT = 10 +_TORCH_DTYPES = ( + torch.bfloat16, + torch.bool, + torch.complex128, + torch.complex32, + torch.complex64, + torch.float16, + torch.float32, + torch.float64, + torch.int16, + torch.int32, + torch.int64, + torch.int8, + torch.qint32, + torch.qint8, + torch.quint4x2, + torch.quint8, + torch.uint8, +) +if hasattr(torch, "uint16"): + _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint16,) +if hasattr(torch, "uint32"): + _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint32,) +if hasattr(torch, "uint64"): + _TORCH_DTYPES = _TORCH_DTYPES + (torch.uint64,) +_STR_DTYPE_TO_DTYPE = {str(dtype): dtype for dtype in _TORCH_DTYPES} +_STRDTYPE2DTYPE = _STR_DTYPE_TO_DTYPE +_DTYPE_TO_STR_DTYPE = { + dtype: str_dtype for str_dtype, dtype in _STR_DTYPE_TO_DTYPE.items() +} +_DTYPE2STRDTYPE = _STR_DTYPE_TO_DTYPE + +IndexType = Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]] +DeviceType = Union[torch.device, str, int] + + +_KEY_ERROR = 'key "{}" not found in {} with ' "keys {}" +_LOCK_ERROR = ( + "Cannot modify locked TensorDict. For in-place modification, consider " + "using the `set_()` method and make sure the key is present." +) + + +LOGGING_LEVEL = os.environ.get("TD_LOGGING_LEVEL", "DEBUG") +logger = logging.getLogger("tensordict") +logger.setLevel(getattr(logging, LOGGING_LEVEL)) +# Disable propagation to the root logger +logger.propagate = False +# Remove all attached handlers +while logger.hasHandlers(): + logger.removeHandler(logger.handlers[0]) +console_handler = logging.StreamHandler() +console_handler.setLevel(logging.INFO) +formatter = logging.Formatter("%(asctime)s [%(name)s][%(levelname)s] %(message)s") +console_handler.setFormatter(formatter) +logger.addHandler(console_handler) + + +def strtobool(val): + """Convert a string representation of truth to true (1) or false (0). + + True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values + are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if + 'val' is anything else. + """ + val = val.lower() + if val in ("y", "yes", "t", "true", "on", "1"): + return 1 + elif val in ("n", "no", "f", "false", "off", "0"): + return 0 + else: + raise ValueError(f"invalid truth value {val!r}") + + +def _sub_index(tensor: Tensor, idx: IndexType) -> Tensor: + """Allows indexing of tensors with nested tuples. + + >>> sub_tensor1 = tensor[tuple1][tuple2] + >>> sub_tensor2 = _sub_index(tensor, (tuple1, tuple2)) + >>> assert torch.allclose(sub_tensor1, sub_tensor2) + + Args: + tensor (Tensor): tensor to be indexed. + idx (tuple of indices): indices sequence to be used. + + """ + if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple): + idx0 = idx[0] + idx1 = idx[1:] + return _sub_index(_sub_index(tensor, idx0), idx1) + return tensor[idx] + + +def convert_ellipsis_to_idx( + idx: tuple[int | Ellipsis] | Ellipsis, batch_size: list[int] +) -> tuple[int, ...]: + """Given an index containing an ellipsis or just an ellipsis, converts any ellipsis to slice(None). + + Example: + >>> idx = (..., 0) + >>> batch_size = [1,2,3] + >>> new_index = convert_ellipsis_to_idx(idx, batch_size) + >>> print(new_index) + (slice(None, None, None), slice(None, None, None), 0) + + Args: + idx (tuple, Ellipsis): Input index + batch_size (list): Shape of tensor to be indexed + + Returns: + new_index (tuple): Output index + """ + istuple = isinstance(idx, tuple) + if (not istuple and idx is not Ellipsis) or ( + istuple and all(_idx is not Ellipsis for _idx in idx) + ): + return idx + new_index = () + num_dims = len(batch_size) + + if idx is Ellipsis: + idx = (...,) + + num_ellipsis = sum(_idx is Ellipsis for _idx in idx) + if num_dims < (len(idx) - num_ellipsis - sum(item is None for item in idx)): + raise RuntimeError("Not enough dimensions in TensorDict for index provided.") + + start_pos, after_ellipsis_length = None, 0 + for i, item in enumerate(idx): + if item is Ellipsis: + if start_pos is not None: + raise RuntimeError("An index can only have one ellipsis at most.") + else: + start_pos = i + if item is not Ellipsis and start_pos is not None: + after_ellipsis_length += 1 + if item is None: + # unsqueeze + num_dims += 1 + + before_ellipsis_length = start_pos + if start_pos is None: + return idx + else: + ellipsis_length = num_dims - after_ellipsis_length - before_ellipsis_length + + new_index += idx[:start_pos] + + ellipsis_start = start_pos + ellipsis_end = start_pos + ellipsis_length + new_index += (slice(None),) * (ellipsis_end - ellipsis_start) + + new_index += idx[start_pos + 1 : start_pos + 1 + after_ellipsis_length] + + if len(new_index) != num_dims: + raise RuntimeError( + f"The new index {new_index} is incompatible with the dimensions of the batch size {num_dims}." + ) + + return new_index + + +def _copy(self: list[int]) -> list[int]: + return list(self) + + +def infer_size_impl(shape: list[int], numel: int) -> list[int]: + """Infers the shape of an expanded tensor whose number of elements is indicated by :obj:`numel`. + + Copied from pytorch for compatibility issues (See #386). + See https://github.com/pytorch/pytorch/blob/35d4fa444b67cbcbe34a862782ddf2d92f5b1ce7/torch/jit/_shape_functions.py + for the original copy. + + """ + newsize = 1 + infer_dim: int | None = None + for dim in range(len(shape)): + if shape[dim] == -1: + if infer_dim is not None: + raise AssertionError("only one dimension can be inferred") + infer_dim = dim + elif shape[dim] >= 0: + newsize *= shape[dim] + else: + raise AssertionError("invalid shape dimensions") + if not ( + numel == newsize + or (infer_dim is not None and newsize > 0 and numel % newsize == 0) + ): + raise AssertionError("invalid shape") + out = _copy(shape) + if infer_dim is not None: + out[infer_dim] = numel // newsize + return out + + +def _unwrap_value(value: Tensor) -> Tensor: + # batch_dims = value.ndimension() + if not isinstance(value, Tensor): + out = value + elif is_batchedtensor(value): + out = get_unwrapped(value) + else: + out = value + return out + # batch_dims = out.ndimension() - batch_dims + # batch_size = out.shape[:batch_dims] + # return out, batch_size + + +if hasattr(math, "prod"): # Python 3.8+ + + def prod(sequence): + """General prod function, that generalised usage across math and np. + + Created for multiple python versions compatibility. + + """ + return math.prod(sequence) + +else: + + def prod(sequence): + """General prod function, that generalised usage across math and np. + + Created for multiple python versions compatibility. + + """ + return int(np.prod(sequence)) + + +def expand_as_right( + tensor: torch.Tensor | TensorDictBase, + dest: torch.Tensor | TensorDictBase, +) -> torch.Tensor | TensorDictBase: + """Expand a tensor on the right to match another tensor shape. + + Args: + tensor: tensor to be expanded + dest: tensor providing the target shape + + Returns: + a tensor with shape matching the dest input tensor shape. + + Examples: + >>> tensor = torch.zeros(3,4) + >>> dest = torch.zeros(3,4,5) + >>> print(expand_as_right(tensor, dest).shape) + torch.Size([3,4,5]) + + """ + if dest.ndimension() < tensor.ndimension(): + raise RuntimeError( + "expand_as_right requires the destination tensor to have less " + f"dimensions than the input tensor, got" + f" tensor.ndimension()={tensor.ndimension()} and " + f"dest.ndimension()={dest.ndimension()}" + ) + if any( + tensor.shape[i] != dest.shape[i] and tensor.shape[i] != 1 + for i in range(tensor.ndimension()) + ): + raise RuntimeError( + f"tensor shape is incompatible with dest shape, " + f"got: tensor.shape={tensor.shape}, dest={dest.shape}" + ) + for _ in range(dest.ndimension() - tensor.ndimension()): + tensor = tensor.unsqueeze(-1) + return tensor.expand(dest.shape) + + +def expand_right(tensor: Tensor, shape: Sequence[int]) -> Tensor: + """Expand a tensor on the right to match a desired shape. + + Args: + tensor: tensor to be expanded + shape: target shape + + Returns: + a tensor with shape matching the target shape. + + Examples: + >>> tensor = torch.zeros(3,4) + >>> shape = (3,4,5) + >>> print(expand_right(tensor, shape).shape) + torch.Size([3,4,5]) + + """ + tensor_expand = tensor + while tensor_expand.ndimension() < len(shape): + tensor_expand = tensor_expand.unsqueeze(-1) + tensor_expand = tensor_expand.expand(shape) + return tensor_expand + + +def _populate_np_dtypes(): + d = {} + for dtype in _TORCH_DTYPES: + dtype_str = str(dtype).split(".")[-1] + try: + d[np.dtype(dtype_str)] = dtype + except TypeError: + continue + return d + + +NUMPY_TO_TORCH_DTYPE_DICT = _populate_np_dtypes() + +TORCH_TO_NUMPY_DTYPE_DICT = { + value: key for key, value in NUMPY_TO_TORCH_DTYPE_DICT.items() +} + + +def is_nested_key(key: NestedKey) -> bool: + """Returns True if key is a NestedKey.""" + if isinstance(key, str): + return True + if key and isinstance(key, (list, tuple)): + return all(isinstance(subkey, str) for subkey in key) + return False + + +def is_seq_of_nested_key(seq: Sequence[NestedKey]) -> bool: + """Returns True if seq is a Sequence[NestedKey].""" + if seq and isinstance(seq, Sequence): + return all(is_nested_key(k) for k in seq) + elif isinstance(seq, Sequence): + # we allow empty inputs + return True + return False + + +def _ndimension(tensor: Tensor) -> int: + if isinstance(tensor, Tensor): + return tensor.ndimension() + else: + return tensor.ndimension() + + +def _shape(tensor: Tensor, nested_shape=False) -> torch.Size: + if isinstance(tensor, UninitializedTensorMixin): + return torch.Size([*getattr(tensor, "batch_size", ()), -1]) + elif not isinstance(tensor, Tensor): + return tensor.shape + if tensor.is_nested: + if nested_shape: + return tensor._nested_tensor_size() + shape = [] + for i in range(tensor.ndim): + try: + shape.append(tensor.size(i)) + except RuntimeError: + shape.append(-1) + return torch.Size(shape) + return tensor.shape + + +def _device(tensor: Tensor) -> torch.device: + if isinstance(tensor, Tensor): + return tensor.device + else: + return tensor.device + + +def _is_shared(tensor: Tensor) -> bool: + if isinstance(tensor, Tensor): + if torch._C._functorch.is_batchedtensor(tensor): + return None + return tensor.is_shared() + if isinstance(tensor, ftdim.Tensor): + return None + else: + return tensor.is_shared() + + +def _is_meta(tensor: Tensor) -> bool: + if isinstance(tensor, Tensor): + return tensor.is_meta + else: + return tensor.is_meta + + +def _dtype(tensor: Tensor) -> torch.dtype: + if isinstance(tensor, Tensor): + return tensor.dtype + else: + return tensor.dtype + + +def _get_item(tensor: Tensor, index: IndexType) -> Tensor: + try: + return tensor[index] + except IndexError as err: + # try to map list index to tensor, and assess type. If bool, we + # likely have a nested list of booleans which is not supported by pytorch + if _is_lis_of_list_of_bools(index): + index = torch.tensor(index, device=tensor.device) + if index.dtype is torch.bool: + raise RuntimeError( + "Indexing a tensor with a nested list of boolean values is " + "not supported by PyTorch.", + ) + return tensor[index] + raise err + + +def _set_item( + tensor: Tensor, index: IndexType, value: Tensor, *, validated, non_blocking +) -> Tensor: + # the tensor must be validated + if not validated: + raise RuntimeError + if isinstance(tensor, Tensor): + tensor[index] = value + return tensor + from tensordict.tensorclass import NonTensorData, NonTensorStack + + if is_non_tensor(tensor): + if ( + isinstance(value, NonTensorData) + and isinstance(tensor, NonTensorData) + and tensor.data == value.data + ): + return tensor + elif isinstance(tensor, NonTensorData): + tensor = NonTensorStack.from_nontensordata(tensor) + if tensor.stack_dim != 0: + tensor = NonTensorStack(*tensor.unbind(0), stack_dim=0) + tensor[index] = value + return tensor + else: + tensor[index] = value + return tensor + + +def _requires_grad(tensor: Tensor) -> bool: + if isinstance(tensor, Tensor): + return tensor.requires_grad + else: + return tensor.requires_grad + + +class timeit: + """A dirty but easy to use decorator for profiling code.""" + + _REG = {} + + def __init__(self, name) -> None: + self.name = name + + def __call__(self, fn): + @wraps(fn) + def decorated_fn(*args, **kwargs): + with self: + out = fn(*args, **kwargs) + return out + + return decorated_fn + + def __enter__(self): + self.t0 = time.time() + + def __exit__(self, exc_type, exc_val, exc_tb): + t = time.time() - self.t0 + val = self._REG.setdefault(self.name, [0.0, 0.0, 0]) + + count = val[2] + N = count + 1 + val[0] = val[0] * (count / N) + t / N + val[1] += t + val[2] = N + + @staticmethod + def print(prefix=None): # noqa: T202 + keys = list(timeit._REG) + keys.sort() + for name in keys: + strings = [] + if prefix: + strings.append(prefix) + strings.append( + f"{name} took {timeit._REG[name][0] * 1000:4.4} msec (total = {timeit._REG[name][1]} sec)" + ) + logger.info(" -- ".join(strings)) + + @staticmethod + def erase(): + for k in timeit._REG: + timeit._REG[k] = [0.0, 0.0, 0] + + +def int_generator(seed): + """A pseudo-random chain generator. + + To be used to produce deterministic integer sequences + + Examples: + >>> for _ in range(2): + ... init_int = 10 + ... for _ in range(10): + ... init_int = int_generator(init_int) + ... print(init_int, end=", ") + ... print("") + 6756, 1717, 4410, 9740, 9611, 9716, 5397, 7745, 4521, 7523, + 6756, 1717, 4410, 9740, 9611, 9716, 5397, 7745, 4521, 7523, + """ + max_seed_val = 10_000 + rng = np.random.default_rng(seed) + seed = int.from_bytes(rng.bytes(8), "big") + return seed % max_seed_val + + +def _is_lis_of_list_of_bools(index, first_level=True): + # determines if an index is a list of list of bools. + # this is aimed at catching a deprecation feature where list of list + # of bools are valid indices + if first_level: + if not isinstance(index, list): + return False + if not len(index): + return False + if isinstance(index[0], list): + return _is_lis_of_list_of_bools(index[0], False) + return False + # then we know it is a list of lists + if isinstance(index[0], bool): + return True + if isinstance(index[0], list): + return _is_lis_of_list_of_bools(index[0], False) + return False + + +def is_tensorclass(obj: type | Any) -> bool: + """Returns True if obj is either a tensorclass or an instance of a tensorclass.""" + cls = obj if isinstance(obj, type) else type(obj) + return _is_tensorclass(cls) + + +_TENSORCLASS_MEMO = {} + + +def _is_tensorclass(cls: type) -> bool: + out = _TENSORCLASS_MEMO.get(cls) + if out is None: + out = getattr(cls, "_is_tensorclass", False) + if not is_compiling(): + _TENSORCLASS_MEMO[cls] = out + return out + + +def _unfold_sequence(seq): + for item in seq: + if isinstance(item, (list, tuple)): + yield tuple(_unfold_sequence(item)) + else: + if isinstance(item, (str, int, slice)) or item is Ellipsis: + yield item + else: + yield id(item) + + +def _make_cache_key(args, kwargs): + """Creates a key for the cache such that memory footprint is minimized.""" + # Fast path for the common args + if not args and not kwargs: + return ((), ()) + elif not kwargs and len(args) == 1 and type(args[0]) is str: + return (args, ()) + else: + return ( + tuple(_unfold_sequence(args)), + tuple(_unfold_sequence(sorted(kwargs.items()))), + ) + + +def cache(fun): + """A cache for TensorDictBase subclasses. + + This decorator will cache the values returned by a method as long as the + input arguments match. + Leaves (tensors and such) are not cached. + The cache is stored within the tensordict such that it can be erased at any + point in time. + + Examples: + >>> import timeit + >>> from tensordict import TensorDict + >>> class SomeOtherTd(TensorDict): + ... @cache + ... def all_keys(self): + ... return set(self.keys(include_nested=True)) + >>> td = SomeOtherTd({("a", "b", "c", "d", "e", "f", "g"): 1.0}, []) + >>> td.lock_() + >>> print(timeit.timeit("set(td.keys(True))", globals={'td': td})) + 11.057 + >>> print(timeit.timeit("set(td.all_keys())", globals={'td': td})) + 0.88 + """ + + @wraps(fun) + def newfun(_self: "TensorDictBase", *args, **kwargs): + if not _self.is_locked or is_compiling(): + return fun(_self, *args, **kwargs) + cache = _self._cache + if cache is None: + cache = _self._cache = defaultdict(dict) + cache = cache[fun.__name__] + key = _make_cache_key(args, kwargs) + if key not in cache: + out = fun(_self, *args, **kwargs) + if not isinstance(out, Tensor): + # we don't cache tensors to avoid filling the mem and / or + # stacking them from their origin + cache[key] = out + else: + out = cache[key] + return out + + return newfun + + +def erase_cache(fun): + """A decorator to erase the cache at each call.""" + + @wraps(fun) + def new_fun(self, *args, **kwargs): + self._erase_cache() + return fun(self, *args, **kwargs) + + return new_fun + + +_NON_STR_KEY_TUPLE_ERR = "Nested membership checks with tuples of strings is only supported when setting `include_nested=True`." +_NON_STR_KEY_ERR = "TensorDict keys are always strings. Membership checks are only supported for strings or non-empty tuples of strings (for nested TensorDicts)" +_GENERIC_NESTED_ERR = "Only NestedKeys are supported. Got key {}." + + +class _StringKeys(KeysView): + """A key view where contains is restricted to strings. + + Saving the keys as an attribute is 25% faster than just subclassing KeysView. + + """ + + def __init__(self, keys): + self.keys = keys + + def __getitem__(self, key: str) -> Any: + return self.keys.__getitem__(key) + + def __iter__(self): + yield from self.keys + + def __repr__(self): + return f"{type(self).__name__}({self.keys})" + + def __len__(self): + return len(self.keys) + + def __contains__(self, item): + if not isinstance(item, str): + try: + unravel_item = _unravel_key_to_tuple(item) + if not unravel_item: # catch errors during unravel + raise TypeError + except Exception: + raise TypeError(_NON_STR_KEY_ERR) + if len(unravel_item) > 1: + raise TypeError(_NON_STR_KEY_TUPLE_ERR) + else: + item = unravel_item[0] + return self.keys.__contains__(item) + + +_StringOnlyDict = dict + + +def lock_blocked(func): + """Checks that the tensordict is unlocked before executing a function.""" + + @wraps(func) + def new_func(self, *args, **kwargs): + if ( + not kwargs.get("ignore_lock", False) + and self.is_locked + and not kwargs.get("inplace") + ): + raise RuntimeError(_LOCK_ERROR) + return func(self, *args, **kwargs) + + return new_func + + +def _strong_ref(self): + return lambda self: self + + +def _as_context_manager(attr=None): + """Converts a method to a decorator. + + Examples: + >>> from tensordict import TensorDict + >>> data = TensorDict() + >>> with data.lock_(): # lock_ is decorated + ... assert data.is_locked + >>> assert not data.is_locked + """ + + def __call__(func): + if attr is not None: + + @wraps(func) + def func_as_decorator(_self, *args, **kwargs): + _attr_pre = getattr(_self, attr) + out = func(_self, *args, **kwargs) + _attr_post = getattr(_self, attr) + if out is not None: + if _attr_post is not _attr_pre: + ref = weakref.ref(_self) + out_lo = out + if is_tensorclass(out_lo): + # We write in the tensordict but the ref is still to self (the tensorclass object) + # we do this because we don't want to call the __setattr__ of the tensorclass + out_lo = out_lo._tensordict + out_lo._last_op = ( + func.__name__, + ( + args, + kwargs, + ref, + ), + ) + else: + out._last_op = None + return out + + else: + + @wraps(func) + def func_as_decorator(_self, *args, **kwargs): + out = func(_self, *args, **kwargs) + if out is not None: + ref = weakref.ref(_self) + out_lo = out + if is_tensorclass(out_lo): + # We write in the tensordict but the ref is still to self (the tensorclass object) + # we do this because we don't want to call the __setattr__ of the tensorclass + out_lo = out_lo._tensordict + + out_lo._last_op = (func.__name__, (args, kwargs, ref)) + return out + + return func_as_decorator + + return __call__ + + +def _find_smallest_uint(N): + if not hasattr(torch, "uint32"): + # Fallback + return torch.int64 + if N < 0: + raise ValueError("N must be a non-negative integer") + + int8_max = 127 + int16_max = 32767 + int32_max = 2147483647 + int64_max = 9223372036854775807 + if N <= int8_max: + return torch.int8 + elif N <= int16_max: + return torch.int16 + elif N <= int32_max: + return torch.int32 + elif N <= int64_max: + return torch.int64 + else: + return "uint is too large to be represented by uint64" + + +def _split_tensordict( + td, + chunksize, + num_chunks, + num_workers, + dim, + use_generator=False, + to_tensordict=False, + shuffle=False, +): + if shuffle and not use_generator: + raise RuntimeError( + "Shuffling is not permitted unless use_generator is set to ``True`` for efficiency purposes." + ) + if chunksize is None and num_chunks is None: + num_chunks = num_workers + if chunksize is not None and num_chunks is not None: + raise ValueError( + "Either chunksize or num_chunks must be provided, but not both." + ) + if num_chunks is not None: + num_chunks = min(td.shape[dim], num_chunks) + if use_generator: + + def next_index(td=td, dim=dim, num_chunks=num_chunks): + idx_start = 0 + n = td.shape[dim] + chunksize = -(n // -num_chunks) + idx_end = chunksize + while idx_start < n: + yield slice(idx_start, idx_end) + idx_start = idx_end + idx_end += chunksize + + else: + return td.chunk(num_chunks, dim=dim) + else: + if chunksize == 0: + if use_generator: + + def next_index(td=td, dim=dim): + yield from range(td.shape[dim]) + + else: + return td.unbind(dim=dim) + else: + if use_generator: + + def next_index(td=td, dim=dim, chunksize=chunksize): + idx_start = 0 + idx_end = chunksize + n = td.shape[dim] + while idx_start < n: + yield slice(idx_start, idx_end) + idx_start = idx_end + idx_end += chunksize + + else: + chunksize = min(td.shape[dim], chunksize) + return td.split(chunksize, dim=dim) + # end up here only when use_generator = True + if shuffle: + + def next_index_shuffle(next_index=next_index): + n = td.shape[dim] + device = td.device + rp = torch.randperm(n, dtype=_find_smallest_uint(n), device=device) + for idx in next_index(): + yield rp[idx].long() + + next_index = next_index_shuffle + + def _split_generator(): + base = (slice(None),) * dim + for idx in next_index(): + out = td[base + (idx,)] + if to_tensordict: + out = out.to_tensordict() + yield out + + return _split_generator() + + +def _parse_to(*args, **kwargs): + batch_size = kwargs.pop("batch_size", None) + non_blocking_pin = kwargs.pop("non_blocking_pin", False) + num_threads = kwargs.pop("num_threads", None) + other = kwargs.pop("other", None) + inplace = kwargs.pop("inplace", False) + if not is_compiling(): + device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to( + *args, **kwargs + ) + else: + non_blocking = kwargs.get("non_blocking", False) + convert_to_format = kwargs.get("convert_to_format") + if len(args) > 0: + device = torch.device(args[0]) + if len(args) > 1: + dtype = args[1] + else: + dtype = kwargs.get("dtype") + else: + device = kwargs.get("device") + dtype = kwargs.get("dtype") + if device is not None: + device = torch.device(device) + + if device and device.type == "cuda" and device.index is None: + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + if other is not None: + if device is not None and device != other.device: + raise ValueError("other and device cannot be both passed") + device = other.device + dtypes = {val.dtype for val in other.values(True, True)} + if len(dtypes) > 1 or len(dtypes) == 0: + dtype = None + elif len(dtypes) == 1: + dtype = list(dtypes)[0] + return ( + device, + dtype, + non_blocking, + convert_to_format, + batch_size, + non_blocking_pin, + num_threads, + inplace, + ) + + +class _ErrorInteceptor: + """Context manager for catching errors and modifying message. + + Intended for use with stacking / concatenation operations applied to TensorDicts. + + """ + + DEFAULT_EXC_MSG = "Expected all tensors to be on the same device" + + def __init__( + self, + key: NestedKey, + prefix: str, + exc_msg: str | None = None, + exc_type: type[Exception] | None = None, + ) -> None: + self.exc_type = exc_type if exc_type is not None else RuntimeError + self.exc_msg = exc_msg if exc_msg is not None else self.DEFAULT_EXC_MSG + self.prefix = prefix + self.key = key + + def _add_key_to_error_msg(self, msg: str) -> str: + if msg.startswith(self.prefix): + return f'{self.prefix} "{self.key}" /{msg[len(self.prefix):]}' + return f'{self.prefix} "{self.key}". {msg}' + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_value, _): + if exc_type is self.exc_type and ( + self.exc_msg is None or self.exc_msg in str(exc_value) + ): + exc_value.args = (self._add_key_to_error_msg(str(exc_value)),) + + +def _nested_keys_to_dict(keys: Iterator[NestedKey]) -> dict[str, Any]: + nested_keys = {} + for key in keys: + if isinstance(key, str): + nested_keys.setdefault(key, {}) + else: + d = nested_keys + for subkey in key: + d = d.setdefault(subkey, {}) + return nested_keys + + +def _dict_to_nested_keys( + nested_keys: dict[NestedKey, NestedKey], prefix: tuple[str, ...] = () +) -> tuple[str, ...]: + for key, subkeys in nested_keys.items(): + if subkeys: + yield from _dict_to_nested_keys(subkeys, prefix=(*prefix, key)) + elif prefix: + yield (*prefix, key) + else: + yield key + + +def _default_hook(td: T, key: tuple[str, ...]) -> None: + """Used to populate a tensordict. + + For example, ``td.set(("a", "b"))`` may require to create ``"a"``. + + """ + out = td.get(key[0]) + if out is None: + td._create_nested_str(key[0]) + out = td._get_str(key[0], None) + return out + + +def _get_leaf_tensordict( + tensordict: T, key: tuple[str, ...], hook: Callable = None +) -> tuple[TensorDictBase, str]: + # utility function for traversing nested tensordicts + # hook should return the default value for tensordict.get(key) + while len(key) > 1: + if hook is not None: + tensordict = hook(tensordict, key) + else: + tensordict = tensordict.get(key[0]) + if tensordict is None: + raise KeyError(f"No sub-tensordict with key {key[0]}.") + key = key[1:] + return tensordict, key[0] + + +def assert_close( + actual: T, + expected: T, + rtol: float | None = None, + atol: float | None = None, + equal_nan: bool = True, + intersection: bool = False, + msg: str = "", + prefix: NestedKey = (), +) -> bool: + """Asserts that two tensordicts, `actual` and `expected`, are element-wise equal within a tolerance for all entries. + + This function checks if the elements of the `actual` tensor are close to the corresponding elements + of the `expected` tensordict, within a relative tolerance (`rtol`) and an absolute tolerance (`atol`). + + It is similar to the :func:`~torch.testing.assert_close` function in PyTorch, but with tensordicts inputs. + + Args: + actual (T): The tensordict containing actual values. + expected (T): The tensordict containing expected values. + rtol (float | None, optional): The relative tolerance parameter. Default is None. + atol (float | None, optional): The absolute tolerance parameter. Default is None. + equal_nan (bool, optional): If True, ``NaNs`` will be considered equal to ``NaNs``. Default is ``True``. + intersection (bool, optional): If True, only the intersection of the two tensordicts will be compared. + Default is ``False``. + msg (str, optional): An optional message to include in the assertion error if the check fails. + prefix (NestedKey, optional): a prefix to add to the key for error messages. + + Returns: + bool: True if the tensors are close within the specified tolerances, raise an exception otherwise. + + Raises: + AssertionError: If the tensordicts are not close within the specified tolerances. + + """ + from tensordict.base import _is_tensor_collection + + if not _is_tensor_collection(type(actual)) or not _is_tensor_collection( + type(expected) + ): + raise TypeError( + f"assert_allclose inputs must be of TensorDict type, got {type(actual)} and {type(expected)}" + ) + + from tensordict._lazy import LazyStackedTensorDict + + if is_tensorclass(actual): + actual = actual._tensordict + if is_tensorclass(expected): + expected = expected._tensordict + + if isinstance(actual, LazyStackedTensorDict) and isinstance( + expected, LazyStackedTensorDict + ): + if expected.stack_dim != actual.stack_dim: + # turn expected in actual stack dim + expected = expected.to_lazystack(actual.stack_dim) + + for sub_actual, sub_expected in _zip_strict( + actual.tensordicts, expected.tensordicts + ): + assert_close( + sub_actual, + sub_expected, + rtol=rtol, + atol=atol, + msg=msg, + intersection=intersection, + equal_nan=equal_nan, + prefix=prefix, + ) + return True + + try: + set1 = set(actual.keys()) + set2 = set(expected.keys()) + except ValueError: + # Persistent tensordicts do not work with is_leaf + def istensor(cls): + return issubclass(cls, torch.Tensor) + + set1 = set(actual.keys(is_leaf=istensor)) + set2 = set(expected.keys(is_leaf=istensor)) + if not intersection and ( + not (len(set1.difference(set2)) == 0 and len(set2) == len(set1)) + ): + _mismatch_keys(set1, set2) + elif intersection and set1 != set2: + actual = actual.select(*set2, strict=False) + expected = expected.select(*set1, strict=False) + + keys = sorted(actual.keys(), key=str) + for key in keys: + input1 = actual.get(key) + input2 = expected.get(key) + if _is_tensor_collection(type(input1)): + if is_non_tensor(input1): + # We skip non-tensor data + continue + assert_close( + input1, + input2, + rtol=rtol, + atol=atol, + msg=msg, + intersection=intersection, + equal_nan=equal_nan, + prefix=prefix + (key,), + ) + continue + elif not isinstance(input1, torch.Tensor): + continue + try: + if input1.is_nested: + input1v = input1.values() + input2v = input2.values() + mse = (input1v.to(torch.float) - input2v.to(torch.float)).pow(2).sum() + input1o = input1.offsets() + input2o = input2.offsets() + mse = ( + mse + + (input1o.to(torch.float) - input2o.to(torch.float)).pow(2).sum() + ) + else: + mse = (input1.to(torch.float) - input2.to(torch.float)).pow(2).sum() + except Exception as err: + raise RuntimeError( + f"Failed to compare key {prefix + (key,)}. Scroll up for more details." + ) from err + mse = mse.data.div(input1.numel()).sqrt().item() + + local_msg = f"key {prefix + (key,)} does not match, got mse = {mse:4.4f}" + new_msg = ",\t".join([local_msg, msg]) if len(msg) else local_msg + if input1.is_nested: + torch.testing.assert_close( + input1v.data, + input2v.data, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + msg=new_msg, + ) + else: + torch.testing.assert_close( + input1.data, + input2.data, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + msg=new_msg, + ) + local_msg = f"key {prefix + (key,)} matches" + msg = "\t".join([local_msg, msg]) if len(msg) else local_msg + + return True + + +def _get_repr(tensor: Tensor) -> str: + s = ", ".join( + [ + f"shape={_shape(tensor)}", + f"device={_device(tensor)}", + f"dtype={_dtype(tensor)}", + f"is_shared={_is_shared(tensor)}", + ] + ) + return f"{type(tensor).__name__}({s})" + + +def _get_repr_custom(cls, shape, device, dtype, is_shared) -> str: + s = ", ".join( + [ + f"shape={shape}", + f"device={device}", + f"dtype={dtype}", + f"is_shared={is_shared}", + ] + ) + return f"{cls.__name__}({s})" + + +def _make_repr(key: NestedKey, item, tensordict: T, sep) -> str: + from tensordict.base import _is_tensor_collection + + if _is_tensor_collection(type(item)): + return sep.join([key, repr(tensordict.get(key))]) + return sep.join([key, _get_repr(item)]) + + +def _td_fields(td: T, keys=None, sep=": ") -> str: + strs = [] + if keys is None: + keys = td.keys() + for key in keys: + shape = td.get_item_shape(key) + if -1 not in shape: + item = td.get(key) + strs.append(_make_repr(key, item, td, sep=sep)) + else: + # we know td is lazy stacked and the key is a leaf + # so we can get the shape and escape the error + temp_td = td + from tensordict import ( + is_tensor_collection, + LazyStackedTensorDict, + TensorDictBase, + ) + + while isinstance( + temp_td, LazyStackedTensorDict + ): # we need to grab the heterogeneous tensor from the inner nesting level + temp_td = temp_td.tensordicts[0] + tensor = temp_td.get(key) + if is_tensor_collection(tensor): + tensor = td.get(key) + strs.append(_make_repr(key, tensor, td, sep=sep)) + continue + + if isinstance(tensor, TensorDictBase): + substr = _td_fields(tensor) + else: + is_shared = ( + tensor.is_shared() + if not isinstance(tensor, UninitializedTensorMixin) + else None + ) + substr = _get_repr_custom( + type(tensor), + shape=shape, + device=tensor.device, + dtype=tensor.dtype, + is_shared=is_shared, + ) + strs.append(sep.join([key, substr])) + + return indent( + "\n" + ",\n".join(sorted(strs)), + 4 * " ", + ) + + +def _check_keys( + list_of_tensordicts: Sequence[TensorDictBase], + strict: bool = False, + include_nested: bool = False, + leaves_only: bool = False, +) -> set[str] | list[str]: + from tensordict.base import _is_leaf_nontensor + + if not len(list_of_tensordicts): + return set() + keys = list_of_tensordicts[0].keys( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=_is_leaf_nontensor, + ) + # TODO: compile doesn't like set() over an arbitrary object + is_comp = is_compiling() + if is_comp: + keys_set = {k for k in keys} # noqa: C416 + else: + keys_set: set[str] = set(keys) + for td in list_of_tensordicts[1:]: + k = td.keys( + include_nested=include_nested, + leaves_only=leaves_only, + is_leaf=_is_leaf_nontensor, + ) + if not strict: + keys_set = keys_set.intersection(k) + else: + if is_comp: + k = {v for v in k} # noqa: C416 + else: + k = set(k) + if k != keys_set: + raise KeyError( + f"got keys {keys} and {set(td.keys())} which are incompatible" + ) + if strict: + if is_comp: + return [key for key in keys] # noqa: C416 + else: + return list(keys) + return keys_set + + +def _set_max_batch_size( + source: T, batch_dims: int | None = None, keep_compliant_size: bool = False +): + """Updates a tensordict with its maximum batch size.""" + from tensordict.base import _is_tensor_collection + + tensor_data = [ + val + for val in source.values() + if not (_pass_through(val) and not val.batch_size) + ] + + for val in tensor_data: + if _is_tensor_collection(type(val)): + if ( + batch_dims is not None + and keep_compliant_size + and val.batch_dims >= batch_dims + ): + continue + _set_max_batch_size(val, batch_dims=batch_dims) + + batch_size = [] + if not tensor_data: # when source is empty + if batch_dims: + source.batch_size = source.batch_size[:batch_dims] + return source + else: + return source + + curr_dim = 0 + tensor_shapes = [_shape(_tensor_data) for _tensor_data in tensor_data] + + while True: + if len(tensor_shapes[0]) > curr_dim: + curr_dim_size = tensor_shapes[0][curr_dim] + else: + source.batch_size = batch_size + return + for leaf, shape in _zip_strict(tensor_data[1:], tensor_shapes[1:]): + # if we have a nested empty tensordict we can modify its batch size at will + if _is_tensor_collection(type(leaf)) and leaf.is_empty(): + continue + if (len(shape) <= curr_dim) or (shape[curr_dim] != curr_dim_size): + source.batch_size = batch_size + return + if batch_dims is None or len(batch_size) < batch_dims: + batch_size.append(curr_dim_size) + curr_dim += 1 + + +def _clone_value(value, recurse: bool): + from tensordict.base import _is_tensor_collection + + if recurse: + # this is not a problem for locked tds as we will not lock it + return value.clone() + elif _is_tensor_collection(type(value)): + return value._clone(recurse=False) + else: + return value + + +def _is_number(item): + if isinstance(item, (Number, ftdim.Dim)): + return True + if isinstance(item, Tensor) and item.ndim == 0: + return True + if isinstance(item, np.ndarray) and item.ndim == 0: + return True + return False + + +def _expand_index(index, batch_size): + len_index = sum(True for idx in index if idx is not None) + if len_index > len(batch_size): + raise ValueError + if len_index < len(batch_size): + index = index + (slice(None),) * (len(batch_size) - len_index) + return index + + +def _renamed_inplace_method(fn): + def wrapper(*args, **kwargs): + raise RuntimeError( + f"{fn.__name__.rstrip('_')} has been removed, use {fn.__name__} instead" + ) + + return wrapper + + +def _broadcast_tensors(index): + # tensors and range need to be broadcast + tensors = { + i: torch.as_tensor(tensor) + for i, tensor in enumerate(index) + if isinstance(tensor, (range, list, np.ndarray, Tensor)) + } + if tensors: + shape = torch.broadcast_shapes(*[tensor.shape for tensor in tensors.values()]) + tensors = {i: tensor.expand(shape) for i, tensor in tensors.items()} + index = tuple( + idx if i not in tensors else tensors[i] for i, idx in enumerate(index) + ) + return index + + +def _reduce_index(index): + if all( + idx is Ellipsis or (isinstance(idx, slice) and idx == slice(None)) + for idx in index + ): + index = () + return index + + +def _get_shape_from_args(*args, kwarg_name="size", **kwargs): + if not args and not kwargs: + return () + if args: + if len(args) > 1 or isinstance(args[0], Number): + size = args + else: + size = args[0] + if len(kwargs): + raise TypeError( + f"Either the kwarg `{kwarg_name}`, a single shape argument or a sequence of integers can be passed. Got args={args} and kwargs={kwargs}." + ) + else: + size = kwargs.pop(kwarg_name, None) + if size is None: + raise TypeError( + f"Either the kwarg `{kwarg_name}`, a single shape argument or a sequence of integers can be passed. Got args={args} and kwargs={kwargs}." + ) + return size + + +if hasattr(torch.nn, "Buffer"): + _parent_buffer_cls = torch.nn.Buffer + + class Buffer: # noqa: D101 + ... + + class _BufferMeta: ... + +else: + + class _BufferMeta(torch._C._TensorMeta): + # Make `isinstance(t, Buffer)` return True for custom tensor instances that have the _is_buffer flag. + def __instancecheck__(self, instance): + if self is Buffer: + if isinstance(instance, torch.Tensor) and getattr( + instance, "_is_buffer", False + ): + return True + return super().__instancecheck__(instance) + + class Buffer(torch.Tensor, metaclass=_BufferMeta): + """A replicate of torch.nn.Buffer if not available (prior to torch v2.5).""" + + def __new__(cls, data=None, *, persistent=True): + if data is None: + data = torch.empty(0) + + t = data.detach().requires_grad_(data.requires_grad) + t.persistent = persistent + t._is_buffer = True + return t + + __torch_function__ = _disabled_torch_function_impl + + _parent_buffer_cls = Buffer + + +class BufferLegacy(_parent_buffer_cls): + """A buffer subclass that keeps the grad fn history.""" + + def __new__(cls, data=None, *, persistent=True): + if data is None: + data = torch.empty(0) + + t = data + t.persistent = persistent + t._is_buffer = True + return t + + +def _getitem_batch_size(batch_size, index): + """Given an input shape and an index, returns the size of the resulting indexed tensor. + + This function is aimed to be used when indexing is an + expensive operation. + Args: + shape (torch.Size): Input shape + items (index): Index of the hypothetical tensor + + Returns: + Size of the resulting object (tensor or tensordict) + + Examples: + >>> idx = (None, ..., None) + >>> torch.zeros(4, 3, 2, 1)[idx].shape + torch.Size([1, 4, 3, 2, 1, 1]) + >>> _getitem_batch_size([4, 3, 2, 1], idx) + torch.Size([1, 4, 3, 2, 1, 1]) + """ + if not isinstance(index, tuple): + if isinstance(index, int): + return batch_size[1:] + if isinstance(index, slice) and index == slice(None): + return batch_size + index = (index,) + # index = convert_ellipsis_to_idx(index, batch_size) + # broadcast shapes + shapes_dict = {} + look_for_disjoint = False + disjoint = False + bools = [] + for i, idx in enumerate(index): + boolean = False + if isinstance(idx, (range, list)): + shape = len(idx) + elif isinstance(idx, torch.Tensor): + if idx.dtype == torch.bool: + shape = torch.Size([idx.sum()]) + boolean = True + else: + shape = idx.shape + elif isinstance(idx, np.ndarray): + if idx.dtype == np.dtype("bool"): + shape = torch.Size([idx.sum()]) + boolean = True + else: + shape = idx.shape + elif isinstance(idx, slice): + look_for_disjoint = not disjoint and (len(shapes_dict) > 0) + shape = None + else: + shape = None + if shape is not None: + if look_for_disjoint: + disjoint = True + shapes_dict[i] = shape + bools.append(boolean) + bs_shape = None + if shapes_dict: + bs_shape = torch.broadcast_shapes(*shapes_dict.values()) + out = [] + count = -1 + for i, idx in enumerate(index): + if idx is True or idx is None: + out.append(1) + continue + count += 1 if not bools[i] else idx.ndim + if i in shapes_dict: + if bs_shape is not None: + if disjoint: + # the indices will be put at the beginning + out = list(bs_shape) + out + else: + # if there is a single tensor or similar, we just extend + out.extend(bs_shape) + bs_shape = None + continue + elif isinstance(idx, (int, ftdim.Dim)): + # could be spared for efficiency + continue + elif isinstance(idx, slice): + batch = batch_size[count] + if is_compiling(): + out.append(len(range(*_slice_indices(idx, batch)))) + else: + out.append(len(range(*idx.indices(batch)))) + count += 1 + if batch_size[count:]: + out.extend(batch_size[count:]) + return torch.Size(out) + + +# Lazy classes control (legacy feature) +_DEFAULT_LAZY_OP = False +_LAZY_OP = os.environ.get("LAZY_LEGACY_OP") + + +class set_lazy_legacy(_DecoratorContextManager): + """Sets the behaviour of some methods to a lazy transform. + + These methods include :meth:`~tensordict.TensorDict.view`, :meth:`~tensordict.TensorDict.permute`, + :meth:`~tensordict.TensorDict.transpose`, :meth:`~tensordict.TensorDict.squeeze` + and :meth:`~tensordict.TensorDict.unsqueeze`. + + This property is dynamic, ie. it can be changed during the code execution, but + it won't propagate to sub-processes unless it has been called before the process + has been created. + + """ + + def __init__(self, mode: bool) -> None: + super().__init__() + self.mode = mode + + def clone(self) -> set_lazy_legacy: + # override this method if your children class takes __init__ parameters + return type(self)(self.mode) + + def __enter__(self) -> None: + self.set() + + def set(self) -> None: + global _LAZY_OP + self._old_mode = _LAZY_OP + _LAZY_OP = bool(self.mode) + # we do this such that sub-processes see the same lazy op than the main one + os.environ["LAZY_LEGACY_OP"] = str(_LAZY_OP) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + global _LAZY_OP + _LAZY_OP = self._old_mode + os.environ["LAZY_LEGACY_OP"] = str(_LAZY_OP) + + +def lazy_legacy(allow_none=False): + """Returns `True` if lazy representations will be used for selected methods.""" + global _LAZY_OP + if _LAZY_OP is None and allow_none: + return None + elif _LAZY_OP is None: + return _DEFAULT_LAZY_OP + return strtobool(_LAZY_OP) if isinstance(_LAZY_OP, str) else _LAZY_OP + + +def _legacy_lazy(func): + if not func.__name__.startswith("_legacy_"): + raise NameError( + f"The function name {func.__name__} must start with _legacy_ if it's decorated with _legacy_lazy." + ) + func.LEGACY = True + return func + + +# non tensor stack control +_DEFAULT_CAPTURE_NONTENSOR_STACK = False +_CAPTURE_NONTENSOR_STACK = os.environ.get("CAPTURE_NONTENSOR_STACK") + + +class set_capture_non_tensor_stack(_DecoratorContextManager): + """A context manager or decorator to control whether identical non-tensor data should be stacked into a single NonTensorData object or a NonTensorStack. + + Args: + mode (bool): Whether to capture non-tensor stacks. If ``False``, identical + non-tensor data will be stacked into a :class:`~tensordict.NonTensorStack`. If ``True``, + a single :class:`~tensordict.NonTensorData` object will contain the unique value, but with the desired batch-size. + Defaults to ``True``. + + .. note:: Since v0.9, `capture_non_tensor_stack()` returns `False` by default. + You can set the value of :func:`~tensordict.capture_non_tensor_stack` through: + + - The ``CAPTURE_NON_TENSOR_STACK`` environment variable; + - By setting ``set_capture_non_tensor_stack(val: bool).set()`` at the beginning of your script; + - By using ``set_capture_non_tensor_stack(val: bool)`` as a context manager or a decorator. + + It is recommended to use the `set_capture_non_tensor_stack(False)` behavior. + + .. seealso:: :class:`~tensordict.capture_non_tensor_stack` + + Examples: + >>> with set_capture_non_tensor_stack(False): + ... torch.stack([NonTensorData("a"), NonTensorData("a")]) + NonTensorData("a", batch_size=[2]) + >>> @set_capture_non_tensor_stack(False) + ... def my_function(): + ... return torch.stack([NonTensorData("a"), NonTensorData("a")]) + >>> my_function() + NonTensorStack(["a", "a"], stack_dim=0) + """ + + def __init__(self, mode: bool) -> None: + super().__init__() + self.mode = mode + + def clone(self) -> set_capture_non_tensor_stack: + # override this method if your children class takes __init__ parameters + return type(self)(self.mode) + + def __enter__(self) -> None: + self.set() + + def set(self) -> None: + global _CAPTURE_NONTENSOR_STACK + self._old_mode = _CAPTURE_NONTENSOR_STACK + _CAPTURE_NONTENSOR_STACK = bool(self.mode) + # we do this such that sub-processes see the same lazy op than the main one + os.environ["CAPTURE_NONTENSOR_STACK"] = str(_CAPTURE_NONTENSOR_STACK) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + global _CAPTURE_NONTENSOR_STACK + _CAPTURE_NONTENSOR_STACK = self._old_mode + os.environ["CAPTURE_NONTENSOR_STACK"] = str(_CAPTURE_NONTENSOR_STACK) + + +def capture_non_tensor_stack(allow_none=False): + """Get the current setting for capturing non-tensor stacks. + + Args: + allow_none (bool, optional): If ``True``, returns ``None`` if no setting has been + specified. Otherwise, returns the default setting. Defaults to ``False``. + + seealso: :func:`~tensordict.set_capture_non_tensor_stack` + + Returns: + bool or None: The current setting for capturing non-tensor stacks. + + """ + global _CAPTURE_NONTENSOR_STACK + if _CAPTURE_NONTENSOR_STACK is None and allow_none: + return None + elif _CAPTURE_NONTENSOR_STACK is None: + return _DEFAULT_CAPTURE_NONTENSOR_STACK + elif ( + isinstance(_CAPTURE_NONTENSOR_STACK, str) + and _CAPTURE_NONTENSOR_STACK.lower() == "none" + ): + return _DEFAULT_CAPTURE_NONTENSOR_STACK + return ( + strtobool(_CAPTURE_NONTENSOR_STACK) + if isinstance(_CAPTURE_NONTENSOR_STACK, str) + else _CAPTURE_NONTENSOR_STACK + ) + + +# list to stack constrol +_DEFAULT_LIST_TO_STACK = "1" +_LIST_TO_STACK = os.environ.get("LIST_TO_STACK") + + +class set_list_to_stack(_DecoratorContextManager): + """Context manager and decorator to control the behavior of list handling in TensorDict. + + When enabled, lists assigned to a TensorDict will be automatically stacked along the batch dimension. + This can be useful for ensuring that lists of tensors or other elements are treated as stackable entities + within a TensorDict. + + Current Behavior: + If a list is assigned to a TensorDict without this context manager, it will be converted to a numpy array + and wrapped in a NonTensorData if it cannot be cast to a Tensor. + + Args: + mode (bool): If True, enables list-to-stack conversion. If False, disables it. + + Example: + >>> with set_list_to_stack(True): + ... td = TensorDict(a=[torch.zeros(()), torch.ones(())], batch_size=2) + ... assert (td["a"] == torch.tensor([0, 1])).all() + ... assert td[0]["a"] == 0 + ... assert td[1]["a"] == 1 + + .. seealso:: :func:`~tensordict.list_to_stack`. + + """ + + def __init__(self, mode: bool) -> None: + super().__init__() + self.mode = mode + + def clone(self) -> set_list_to_stack: + # override this method if your children class takes __init__ parameters + return type(self)(self.mode) + + def __enter__(self) -> None: + self.set() + + def set(self) -> None: + global _LIST_TO_STACK + self._old_mode = _LIST_TO_STACK + _LIST_TO_STACK = bool(self.mode) + os.environ["LIST_TO_STACK"] = str(_LIST_TO_STACK) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + global _LIST_TO_STACK + _LIST_TO_STACK = self._old_mode + os.environ["LIST_TO_STACK"] = str(_LIST_TO_STACK) + + +def list_to_stack(allow_none=False): + """Retrieves the current setting for list-to-stack conversion in TensorDict. + + This function checks the global environment variable or the context manager setting to determine + whether lists should be automatically stacked when assigned to a TensorDict. + + Current Behavior: + Returns the current setting for list-to-stack conversion. If the setting is not defined and `allow_none` + is True, it returns None. Otherwise, it returns the default setting. + + Args: + allow_none (bool): If True, allows the function to return None if the setting is not defined. + + Returns: + bool or None: The current setting for list-to-stack conversion. + + .. seealso:: :class:`~tensordict.set_list_to_stack`. + + """ + global _LIST_TO_STACK + if _LIST_TO_STACK is None and allow_none: + return None + elif _LIST_TO_STACK is None: + return _DEFAULT_LIST_TO_STACK + elif isinstance(_LIST_TO_STACK, str) and _LIST_TO_STACK.lower() == "none": + return _DEFAULT_LIST_TO_STACK + return ( + strtobool(_LIST_TO_STACK) if isinstance(_LIST_TO_STACK, str) else _LIST_TO_STACK + ) + + +def _convert_list_to_stack( + a_list: list[Any], +) -> tuple[torch.Tensor | TensorDictBase | NonTensorStack, bool]: # noqa + # First, check elements and determine if there are lists within + nontensor = True + if all(isinstance(elt, list) for elt in a_list): + a_list, nontensor = zip(*[_convert_list_to_stack(elt) for elt in a_list]) + nontensor = any(nontensor) + # FIXME: we should check that the type is unique + all_castable = all(isinstance(elt, (bool, int, float)) for elt in a_list) + if all_castable: + return torch.tensor(a_list), False + all_tensors = all(isinstance(elt, torch.Tensor) for elt in a_list) + if not nontensor or all_tensors: + # should we stack? + if all_tensors and len({x.shape for x in a_list}) == 1: + # FIXME: this may lead to some weird behaviours if we have nested lists and by chance one of them has + # things that can be stacked, and others don't. + return torch.stack(a_list), False + # TODO: check that LazyStack understands that a list is a bunch of elements to write in separate tds + return list(a_list), False + from tensordict.base import _is_tensor_collection + + if all(_is_tensor_collection(type(elt)) for elt in a_list): + return torch.stack(a_list), False + from tensordict import NonTensorStack + + return NonTensorStack(*a_list), True + + +def _recursive_unbind_list(a_list, dim): + if dim == 0: + return list(a_list) + try: + return map( + list, _zip_strict(*[_recursive_unbind_list(elt, dim - 1) for elt in a_list]) + ) + except Exception: + raise ValueError("lengths of nested lists differed.") + + +# Process initializer for map +def _proc_init(base_seed, queue, num_threads): + worker_id = queue.get(timeout=120) + seed = base_seed + worker_id + torch.manual_seed(seed) + np_seed = _generate_state(base_seed, worker_id) + np.random.seed(np_seed) + torch.set_num_threads(num_threads) + + +def _prune_selected_keys(keys_to_update, prefix): + if keys_to_update is None: + return None + return tuple( + key[1:] for key in keys_to_update if isinstance(key, tuple) and key[0] == prefix + ) + + +class TensorDictFuture: + """A custom future class for TensorDict multithreaded operations. + + Args: + futures (list of futures): a list of concurrent.futures.Future objects to wait for. + resulting_td (TensorDictBase): instance that will result from the futures + completing. + + """ + + def __init__(self, futures, resulting_td): + self.futures = futures + self.resulting_td = resulting_td + + def result(self): + """Wait and returns the resulting tensordict.""" + concurrent.futures.wait(self.futures) + return self.resulting_td + + +def _is_json_serializable(item): + if isinstance(item, dict): + for key, val in item.items(): + # Per se, int, float and bool are serializable but not recoverable + # as such + if not isinstance(key, (str,)) or not _is_json_serializable(val): + return False + else: + return True + if isinstance(item, (list, tuple, set)): + for val in item: + if not _is_json_serializable(val): + return False + else: + return True + return isinstance(item, (str, int, float, bool)) or item is None + + +def print_directory_tree(path, indent="", display_metadata=True) -> str: + """Prints the directory tree starting from the specified path. + + Args: + path (str): The path of the directory to print. + indent (str): The current indentation level for formatting. + display_metadata (bool): if ``True``, metadata of the dir will be + displayed too. + + Returns: + the string printed with the logger. + + """ + string = [] + if display_metadata: + + def get_directory_size(path="."): + total_size = 0 + + for dirpath, _, filenames in os.walk(path): + for filename in filenames: + file_path = os.path.join(dirpath, filename) + total_size += os.path.getsize(file_path) + + return total_size + + def format_size(size): + # Convert size to a human-readable format + for unit in ["B", "KB", "MB", "GB", "TB"]: + if size < 1024.0: + return f"{size:.2f} {unit}" + size /= 1024.0 + + total_size_bytes = get_directory_size(path) + formatted_size = format_size(total_size_bytes) + string.append(f"Directory size: {formatted_size}") + logger.info(string[-1]) + + if os.path.isdir(path): + string.append(indent + os.path.basename(path) + "/") + logger.info(string[-1]) + indent += " " + for item in os.listdir(path): + string.append( + print_directory_tree( + os.path.join(path, item), indent=indent, display_metadata=False + ) + ) + else: + string.append(indent + os.path.basename(path)) + logger.info(string[-1]) + return "\n".join(string) + + +def isin( + input: TensorDictBase, + reference: TensorDictBase, + key: NestedKey, + dim: int = 0, +) -> Tensor: + """Tests if each element of ``key`` in input ``dim`` is also present in the reference. + + This function returns a boolean tensor of length ``input.batch_size[dim]`` that is ``True`` for elements in + the entry ``key`` that are also present in the ``reference``. This function assumes that both ``input`` and + ``reference`` have the same batch size and contain the specified entry, otherwise an error will be raised. + + Args: + input (TensorDictBase): Input TensorDict. + reference (TensorDictBase): Target TensorDict against which to test. + key (Nestedkey): The key to test. + dim (int, optional): The dimension along which to test. Defaults to ``0``. + + Returns: + out (Tensor): A boolean tensor of length ``input.batch_size[dim]`` that is ``True`` for elements in + the ``input`` ``key`` tensor that are also present in the ``reference``. + + Examples: + >>> td = TensorDict( + ... { + ... "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [7, 8, 9]]), + ... "tensor2": torch.tensor([[10, 20], [30, 40], [40, 50], [50, 60]]), + ... }, + ... batch_size=[4], + ... ) + >>> td_ref = TensorDict( + ... { + ... "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [10, 11, 12]]), + ... "tensor2": torch.tensor([[10, 20], [30, 40], [50, 60]]), + ... }, + ... batch_size=[3], + ... ) + >>> in_reference = isin(td, td_ref, key="tensor1") + >>> expected_in_reference = torch.tensor([True, True, True, False]) + >>> torch.testing.assert_close(in_reference, expected_in_reference) + """ + # Get the data + reference_tensor = reference.get(key) + target_tensor = input.get(key) + + # Check key is present in both tensordict and reference_tensordict + if not isinstance(target_tensor, torch.Tensor): + raise KeyError(f"Key '{key}' not found in input or not a tensor.") + if not isinstance(reference_tensor, torch.Tensor): + raise KeyError(f"Key '{key}' not found in reference or not a tensor.") + + # Check that both TensorDicts have the same number of dimensions + if len(input.batch_size) != len(reference.batch_size): + raise ValueError( + "The number of dimensions in the batch size of the input and reference must be the same." + ) + + # Check dim is valid + batch_dims = input.ndim + if dim >= batch_dims or dim < -batch_dims or batch_dims == 0: + raise ValueError( + f"The specified dimension '{dim}' is invalid for an input TensorDict with batch size '{input.batch_size}'." + ) + + # Convert negative dimension to its positive equivalent + if dim < 0: + dim = batch_dims + dim + + # Find the common indices + N = reference_tensor.shape[dim] + cat_data = torch.cat([reference_tensor, target_tensor], dim=dim) + _, unique_indices = torch.unique( + cat_data, dim=dim, sorted=True, return_inverse=True + ) + out = torch.isin(unique_indices[N:], unique_indices[:N], assume_unique=True) + + return out + + +def _index_preserve_data_ptr(index): + if isinstance(index, tuple): + return all(_index_preserve_data_ptr(idx) for idx in index) + # we can't use a list comprehension here because it fails with tensor indices + if index is None or index is Ellipsis: + return True + if isinstance(index, int): + return True + if isinstance(index, slice) and (index.start == 0 or index.start is None): + return True + return False + + +def remove_duplicates( + input: TensorDictBase, + key: NestedKey, + dim: int = 0, + *, + return_indices: bool = False, +) -> TensorDictBase: + """Removes indices duplicated in `key` along the specified dimension. + + This method detects duplicate elements in the tensor associated with the specified `key` along the specified + `dim` and removes elements in the same indices in all other tensors within the TensorDict. It is expected for + `dim` to be one of the dimensions within the batch size of the input TensorDict to ensure consistency in all + tensors. Otherwise, an error will be raised. + + Args: + input (TensorDictBase): The TensorDict containing potentially duplicate elements. + key (NestedKey): The key of the tensor along which duplicate elements should be identified and removed. It + must be one of the leaf keys within the TensorDict, pointing to a tensor and not to another TensorDict. + dim (int, optional): The dimension along which duplicate elements should be identified and removed. It must be one of + the dimensions within the batch size of the input TensorDict. Defaults to ``0``. + return_indices (bool, optional): If ``True``, the indices of the unique elements in the input tensor will be + returned as well. Defaults to ``False``. + + Returns: + output (TensorDictBase): input tensordict with the indices corrsponding to duplicated elements + in tensor `key` along dimension `dim` removed. + unique_indices (torch.Tensor, optional): The indices of the first occurrences of the unique elements in the + input tensordict for the specified `key` along the specified `dim`. Only provided if return_index is True. + + Example: + >>> td = TensorDict( + ... { + ... "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [7, 8, 9]]), + ... "tensor2": torch.tensor([[10, 20], [30, 40], [40, 50], [50, 60]]), + ... } + ... batch_size=[4], + ... ) + >>> output_tensordict = remove_duplicate_elements(td, key="tensor1", dim=0) + >>> expected_output = TensorDict( + ... { + ... "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), + ... "tensor2": torch.tensor([[10, 20], [30, 40], [50, 60]]), + ... }, + ... batch_size=[3], + ... ) + >>> assert (td == expected_output).all() + """ + tensor = input.get(key) + + # Check if the key is a TensorDict + if tensor is None: + raise KeyError(f"The key '{key}' does not exist in the TensorDict.") + + # Check that the key points to a tensor + if not isinstance(tensor, torch.Tensor): + raise KeyError(f"The key '{key}' does not point to a tensor in the TensorDict.") + + # Check dim is valid + batch_dims = input.ndim + if dim >= batch_dims or dim < -batch_dims or batch_dims == 0: + raise ValueError( + f"The specified dimension '{dim}' is invalid for a TensorDict with batch size '{input.batch_size}'." + ) + + # Convert negative dimension to its positive equivalent + if dim < 0: + dim = batch_dims + dim + + # Get indices of unique elements (e.g. [0, 1, 0, 2]) + _, unique_indices, counts = torch.unique( + tensor, dim=dim, sorted=True, return_inverse=True, return_counts=True + ) + + # Find first occurrence of each index (e.g. [0, 1, 3]) + _, unique_indices_sorted = torch.sort(unique_indices, stable=True) + cum_sum = counts.cumsum(0, dtype=torch.long) + cum_sum = torch.cat( + (torch.zeros(1, device=input.device, dtype=torch.long), cum_sum[:-1]) + ) + first_indices = unique_indices_sorted[cum_sum] + + # Remove duplicate elements in the TensorDict + output = input[(slice(None),) * dim + (first_indices,)] + + if return_indices: + return output, unique_indices + + return output + + +class _CloudpickleWrapper(object): + def __init__(self, fn): + self.fn = fn + + def __getstate__(self): + import cloudpickle + + return cloudpickle.dumps(self.fn) + + def __setstate__(self, ob: bytes): + import pickle + + self.fn = pickle.loads(ob) + + def __call__(self, *args, **kwargs): + return self.fn(*args, **kwargs) + + +class _BatchedUninitializedParameter(UninitializedParameter): + batch_size: torch.Size + in_dim: int | None = None + vmap_level: int | None = None + + def materialize(self, shape, device=None, dtype=None): + UninitializedParameter.materialize( + self, (*self.batch_size, *shape), device=device, dtype=dtype + ) + + +class _BatchedUninitializedBuffer(UninitializedBuffer): + batch_size: torch.Size + in_dim: int | None = None + vmap_level: int | None = None + + def materialize(self, shape, device=None, dtype=None): + UninitializedBuffer.materialize( + self, (*self.batch_size, *shape), device=device, dtype=dtype + ) + + +class _add_batch_dim_pre_hook: + def __call__(self, mod: torch.nn.Module, args, kwargs): + for name, param in list(mod.named_parameters(recurse=False)): + if hasattr(param, "in_dim") and hasattr(param, "vmap_level"): + if _add_batch_dim_c is None: + from torch._C._functorch import ( + _add_batch_dim, + ) # @manual=//caffe2:_C + else: + _add_batch_dim = _add_batch_dim_c + + param = _add_batch_dim(param, param.in_dim, param.vmap_level) + delattr(mod, name) + setattr(mod, name, param) + for key, val in list(mod._forward_pre_hooks.items()): + if val is self: + del mod._forward_pre_hooks[key] + return + else: + raise RuntimeError("did not find pre-hook") + + +def is_non_tensor(data) -> bool: + """Checks if an item is a non-tensor.""" + return _is_non_tensor(type(data)) + + +def _pass_through(data) -> bool: + return _pass_through_cls(type(data)) + + +_NON_TENSOR_MEMO = {} + + +def _is_non_tensor(cls: type): + out = None + is_dynamo = is_compiling() + if not is_dynamo: + out = _NON_TENSOR_MEMO.get(cls) + if out is None: + out = bool(getattr(cls, "_is_non_tensor", False)) + if not is_dynamo: + _NON_TENSOR_MEMO[cls] = out + return out + + +_PASSTHROUGH_MEMO = {} + + +def _pass_through_cls(cls: type): + out = None + is_dynamo = is_compiling() + if not is_dynamo: + out = _PASSTHROUGH_MEMO.get(cls) + if out is None: + out = bool(getattr(cls, "_is_non_tensor", False)) or getattr( + cls, "_pass_through", False + ) + if not is_dynamo: + _PASSTHROUGH_MEMO[cls] = out + return out + + +class KeyDependentDefaultDict(collections.defaultdict): + """A key-dependent default dict. + + Examples: + >>> my_dict = KeyDependentDefaultDict(lambda key: "foo_" + key) + >>> print(my_dict["bar"]) + foo_bar + """ + + def __init__(self, fun): + self.fun = fun + super().__init__() + + def __missing__(self, key): + value = self.fun(key) + self[key] = value + return value + + +def is_namedtuple(obj): + """Check if obj is a namedtuple.""" + return isinstance(obj, tuple) and hasattr(obj, "_fields") + + +def is_namedtuple_class(cls): + """Check if a class is a namedtuple class.""" + base_attrs = {"_fields", "_replace", "_asdict"} + return all(hasattr(cls, attr) for attr in base_attrs) + + +def _make_dtype_promotion(func): + dtype = getattr(torch, func.__name__, None) + + @wraps(func) + def new_func(self): + if dtype is None: + raise NotImplementedError( + f"Your pytorch version {torch.__version__} does not support {dtype}." + ) + + def todtype(x): + return x.to(dtype) + + return self._fast_apply(todtype, propagate_lock=True) + + new_func.__doc__ = rf"""Casts all tensors to ``{str(dtype)}``.""" + return new_func + + +def _unravel_key_to_tuple(key): + if not is_compiling(): + return _unravel_key_to_tuple_cpp(key) + if isinstance(key, str): + return (key,) + if not isinstance(key, tuple): + return () + return tuple(subk for k in key for subk in _unravel_key_to_tuple(k)) + + +def unravel_key(key): + """Unravel a nested key. + + Examples: + >>> unravel_key("a") + "a" + >>> unravel_key(("a",)) + "a" + >>> unravel_key((("a", ("b",)))) + ("a", "b") + + """ + if not is_compiling(): + return unravel_key_cpp(key) + if isinstance(key, str): + return key + if isinstance(key, tuple): + if len(key) == 1: + return unravel_key(key[0]) + return tuple(unravel_key(_key) for _key in key) + raise ValueError("the key must be a str or a tuple of str") + + +def unravel_keys(*keys): + """Unravels a sequence of keys.""" + if not is_compiling(): + return unravel_keys_cpp(*keys) + return tuple(unravel_key(key) for key in keys) + + +def unravel_key_list(keys): + """Unravels a list of keys.""" + if not is_compiling(): + return unravel_key_list_cpp(keys) + return [unravel_key(key) for key in keys] + + +def _encode_key_for_filesystem(key: str, *, robust: bool = True) -> str: + """Encode a TensorDict key to be safe for filesystem paths. + + This function provides a bijective mapping from TensorDict keys to + filesystem-safe filenames by percent-encoding problematic characters. + + Args: + key (str): The original TensorDict key + robust (bool): If True, uses the new robust encoding that percent-encodes + problematic characters. If False, returns the key unchanged (legacy + behavior). Defaults to True. + + Returns: + str: A filesystem-safe encoded key if robust=True, otherwise the original key + + Examples: + >>> _encode_key_for_filesystem("a/b/c") + "a%2Fb%2Fc" + >>> _encode_key_for_filesystem("a/b/c", robust=False) + "a/b/c" + >>> _encode_key_for_filesystem("normal_key") + "normal_key" + """ + if not robust: + # Legacy behavior: return key unchanged + return key + + # Characters that are problematic across filesystems: + # - Unix/Linux: / (path separator), null + # - Windows: < > : " | ? * \ / (path separator), null, control chars + # - macOS: : (path separator in HFS), null + # - General: space, percent (for our encoding) + unsafe_chars = set('/<>:"|?*\\ \0%') + + # Also encode control characters (0-31) and DEL (127) + unsafe_chars.update(chr(i) for i in range(32)) + unsafe_chars.add(chr(127)) + + encoded_parts = [] + for char in key: + if char in unsafe_chars: + # Percent-encode using uppercase hex for consistency + encoded_parts.append(f"%{ord(char):02X}") + else: + encoded_parts.append(char) + + return "".join(encoded_parts) + + +def _get_robust_key_setting_with_warning(key: str, robust_key) -> bool: + """Handle the robust_key parameter with smart deprecation warning. + + Only warns when there's actually a difference between robust and legacy encoding. + + Args: + key: The TensorDict key to check + robust_key: None (auto-detect with warning), False (legacy), True (robust) + + Returns: + bool: The effective robust setting to use + """ + if robust_key is not None: + return robust_key + + # Check if robust encoding would produce a different filename + robust_encoded = _encode_key_for_filesystem(key, robust=True) + # Keep this in case we need some futher mapping + legacy_encoded = _encode_key_for_filesystem(key, robust=False) + + if robust_encoded != legacy_encoded: + # Only warn when there's actually a difference + import warnings + + warnings.warn( + f"The key '{key}' contains characters that will be handled differently " + f"in TensorDict v0.12 for better cross-platform support. " + f"To opt into the new behavior now, use `robust_key=True`. " + f"To suppress this warning and keep the current behavior, use `robust_key=False`. " + f"See https://github.com/pytorch/tensordict/issues/1440 for details.", + FutureWarning, + stacklevel=3, + ) + + # Always use legacy behavior when robust_key=None + return False + + +def _get_robust_key_setting(robust_key) -> bool: + """Handle the robust_key parameter without key-specific logic. + + Args: + robust_key: None (fallback to False), False (legacy), True (robust) + + Returns: + bool: The effective robust setting to use + """ + if robust_key is None: + return False + return robust_key + + +def _decode_key_from_filesystem(encoded_key: str) -> str: + """Decode a filesystem-safe key back to the original TensorDict key. + + This is the reverse of _encode_key_for_filesystem. + + Args: + encoded_key (str): A filesystem-safe encoded key + + Returns: + str: The original TensorDict key + + Examples: + >>> _decode_key_from_filesystem("a%2Fb%2Fc") + "a/b/c" + >>> _decode_key_from_filesystem("key%20with%20spaces") + "key with spaces" + >>> _decode_key_from_filesystem("normal_key") + "normal_key" + """ + decoded_parts = [] + i = 0 + while i < len(encoded_key): + if encoded_key[i] == "%" and i + 2 < len(encoded_key): + try: + # Decode the hex value + hex_str = encoded_key[i + 1 : i + 3] + char_code = int(hex_str, 16) + decoded_parts.append(chr(char_code)) + i += 3 + except ValueError: + # Invalid hex sequence, treat as literal % + decoded_parts.append(encoded_key[i]) + i += 1 + else: + decoded_parts.append(encoded_key[i]) + i += 1 + + return "".join(decoded_parts) + + +def _slice_indices(index: slice, len: int): + """A pure python implementation of slice.indices(len) since torch.compile doesn't recognise it.""" + step = index.step + if step is None: + step = 1 + elif step == 0: + raise ValueError("Step cannot be zero.") + + start = index.start + stop = index.stop + if start is None: + if step > 0: + start = 0 + else: + start = len - 1 + elif start < 0: + start = max(0, len + start) + + if stop is None: + if step > 0: + stop = len + else: + stop = -1 + elif stop > 0: + stop = min(len, stop) + elif step < 0 or (step > 0 and start >= 0): + stop = len + stop + return start, stop, step + + +assert_allclose_td = assert_close + + +def _prefix_last_key(key, prefix): + if isinstance(key, str): + return prefix + key + if len(key) == 1: + return (_prefix_last_key(key[0], prefix),) + return key[:-1] + (_prefix_last_key(key[-1], prefix),) + + +NESTED_TENSOR_ERR = ( + "The PyTorch version isn't compatible with " + "nested tensors. Please upgrade to a more recent " + "version." +) + +_DEVICE2STRDEVICE = KeyDependentDefaultDict(str) + + +def _lock_warn(): + warnings.warn( + "Using lock_() in a compiled graph should " + "only be done if users make sure that the code runs in eager mode. " + "torch.compile doesn't support weakrefs which are used to reference root tensordicts " + "to sub-tensordict and prevent unlocking a node when the graph is locked. " + "Such operation will fail in eager mode but won't be captured by torch.compile.", + category=UserWarning, + ) + + +_lock_warn = assume_constant_result(_lock_warn) + + +def _check_inbuild(): + if not torch._dynamo.config.inline_inbuilt_nn_modules: + raise RuntimeError( + "to_module requires torch._dynamo.config.inline_inbuilt_nn_modules to be set to True." + ) + + +_check_inbuild = assume_constant_result(_check_inbuild) + +if sys.version_info >= (3, 10): + _zip_strict = functools.partial(zip, strict=True) +else: + + def _zip_strict(*iterables): + iterables = tuple(tuple(it) for it in iterables) + lengths = {len(it) for it in iterables} + if len(lengths) > 1: + raise ValueError("lengths of iterables differ.") + + return zip(*iterables) + + +def _pin_mem(q_in, q_out): + while not q_in.empty(): + input = q_in.get(timeout=_PIN_MEM_TIMEOUT) + try: + key, val = input[0], input[1].pin_memory() + except Exception as err: + # Surface the exception + q_out.put(err) + return + q_out.put((key, val)) + + +def _infer_size_impl(shape: List[int], numel: int) -> List[int]: + # A local copy of torch.jit._shape_functions.infer_size_impl which is skipped by torch.compile + newsize = 1 + infer_dim: int | None = None + for dim in range(len(shape)): + if shape[dim] == -1: + if infer_dim is not None: + raise AssertionError("only one dimension can be inferred") + infer_dim = dim + elif shape[dim] >= 0: + newsize *= shape[dim] + else: + raise AssertionError("invalid shape dimensions") + if not ( + numel == newsize + or (infer_dim is not None and newsize > 0 and numel % newsize == 0) + ): + raise AssertionError("invalid shape") + out = _copy(shape) + if infer_dim is not None: + out[infer_dim] = numel // newsize + return out + + +def parse_tensor_dict_string(s: str): + """Parse a TensorDict repr to a TensorDict. + + .. note:: + This functions is intended to be used for debugging, to reproduce a tensordict + given its printed version, and should not be used in real applications. + + """ + from tensordict import TensorDict + + # Regular expression patterns + field_pattern = r"(\w+): Tensor\(shape=torch.Size\((\[(.*?)\])\), device=(\w+), dtype=torch.(\w+), is_shared=(\w+)\)" + nested_field_pattern = r"(\w+): TensorDict\(" + batch_size_pattern = r"batch_size=torch.Size\((\[(.*?)\])\)" + device_pattern = r"device=(\w+)(?=,|$)" + + # Find all nested TensorDicts first + nested_dict_ranges = [] + for match in re.finditer(nested_field_pattern, s): + start_idx = match.start() + depth = 1 + for i in range(start_idx + len(match.group(0)), len(s)): + if s[i] == "(": + depth += 1 + elif s[i] == ")": + depth -= 1 + if depth == 0: + end_idx = i + break + nested_dict_ranges.append((start_idx, end_idx)) + + # Find all fields in the string that are not part of a nested TensorDict + fields = {} + for match in re.finditer(field_pattern, s): + name, _, shape, device, dtype, is_shared = match.groups() + field_start = match.start() + field_end = match.end() + if any( + field_start >= start and field_end <= end + for start, end in nested_dict_ranges + ): + continue # skip if this field is inside a nested TensorDict + shape = [int(x) for x in shape.split(", ")] if shape else [] + fields[name] = torch.zeros( + tuple(shape), device=torch.device(device), dtype=getattr(torch, dtype) + ) + + # Now find nested TensorDicts and add them to the fields + for match in re.finditer(nested_field_pattern, s): + name = match.group(1) + start_idx = match.end() + depth = 1 + for i in range(start_idx, len(s)): + if s[i] == "(": + depth += 1 + elif s[i] == ")": + depth -= 1 + if depth == 0: + end_idx = i + break + content = s[start_idx:end_idx] + nested_fields = parse_tensor_dict_string(f"TensorDict({content})") + fields[name] = nested_fields + + # Parse batch size + batch_size_matches = re.findall(batch_size_pattern, s) + if batch_size_matches: + batch_size_match = batch_size_matches[-1] # Take the last match + if batch_size_match[1]: + batch_size = [int(x) for x in batch_size_match[1].split(", ")] + else: + batch_size = [] + else: + raise ValueError("Batch size not found in the string") + # Parse device + device_matches = re.findall(device_pattern, s) + if device_matches: + device = device_matches[-1] # Take the last match + if isinstance(device, str) and device.lower() == "none": + device = None + else: + device = torch.device(device) + else: + raise ValueError("Device not found in the string") + tensor_dict = TensorDict(fields, batch_size=torch.Size(batch_size), device=device) + return tensor_dict + + +def _rebuild_njt_from_njt(x, values, offsets, lengths): + from torch._subclasses.fake_tensor import FakeTensor + from torch._subclasses.functional_tensor import FunctionalTensor + from torch.nested._internal.nested_tensor import ( + _tensor_symint_registry, + NestedTensor, + ) + from torch.nested._internal.ops import extract_kwargs + + kwargs = extract_kwargs(x) + kwargs["offsets"] = offsets + if x._lengths is not None: + kwargs["lengths"] = lengths + ragged_source = x._lengths + else: + ragged_source = x._offsets + new_thing = kwargs.get("lengths", kwargs.get("offsets")) + if isinstance(new_thing, (FakeTensor, FunctionalTensor)): + from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor + + # Temporary hack until we have the union find + tgt = mb_unwrap_functional_tensor(new_thing) + src = mb_unwrap_functional_tensor(ragged_source) + tgt.nested_int_memo = src.nested_int_memo + elif new_thing is not None: + _tensor_symint_registry[new_thing] = _tensor_symint_registry[ragged_source] + + return NestedTensor( + values, + **kwargs, + ) + + +def _mismatch_keys(keys1, keys2): + def keyfunc(key): + return "".join(key) if isinstance(key, tuple) else key + + keys1 = sorted( + keys1, + key=keyfunc, + ) + keys2 = sorted( + keys2, + key=keyfunc, + ) + if set(keys1) - set(keys2): + sub1 = rf"The first TD has keys {set(keys1) - set(keys2)} that the second does not have." + else: + sub1 = None + if set(keys2) - set(keys1): + sub2 = rf"The second TD has keys {set(keys2) - set(keys1)} that the first does not have." + else: + sub2 = None + main = [r"keys in tensordicts mismatch."] + if sub1 is not None: + main.append(sub1) + if sub2 is not None: + main.append(sub2) + raise KeyError(r" ".join(main)) + + +def _is_dataclass(obj): + """Check if an object is a dataclass.""" + try: + from dataclasses import is_dataclass + + return is_dataclass(obj) + except ImportError: + # Fallback for older Python versions + cls = ( + obj + if isinstance(obj, type) and not isinstance(obj, GenericAlias) + else type(obj) + ) + return hasattr(cls, "__dataclass_fields__") + + +def _is_list_tensor_compatible(t) -> Tuple[bool, tuple | None, type | None]: + length_t = len(t) + dtypes = set() + sizes = set() + for i in t: + if isinstance(i, (float, int, torch.SymInt, Number)): + dtypes.add(type(i)) + if len(dtypes) > 1: + return False, None, None + continue + elif isinstance(i, list): + is_compat, size_i, dtype = _is_list_tensor_compatible(i) + if not is_compat: + return False, None, None + if dtype is not None: + dtypes.add(dtype) + if len(dtypes) > 1: + return False, None, None + sizes.add(size_i) + if len(sizes) > 1: + return False, None, None + continue + return False, None, None + else: + if len(dtypes): + dtype = list(dtypes)[0] + else: + dtype = None + if len(sizes): + return True, (length_t, *list(sizes)[0]), dtype + return True, (length_t,), dtype + + +class _ContextManager: + def __init__(self, default=None): + self._mode: Any | None = default + self._lock = threading.Lock() + + def get_mode(self) -> Any | None: + cm = self._lock if not is_compiling() else nullcontext() + with cm: + return self._mode + + def set_mode(self, type: Any | None) -> None: + cm = self._lock if not is_compiling() else nullcontext() + with cm: + self._mode = type + + +def _maybe_correct_neg_dim( + dim: int, shape: torch.Size | None, ndim: int | None = None +) -> int: + """Corrects neg dim to pos.""" + if ndim is None: + ndim = len(shape) + if dim < 0: + new_dim = ndim + dim + else: + new_dim = dim + if new_dim < 0 or new_dim >= ndim: + if shape is not None: + raise IndexError( + f"Incompatible dim {new_dim} for tensordict with shape {shape}." + ) + raise IndexError( + f"Incompatible dim {new_dim} for tensordict with batch dims {ndim}." + ) + return new_dim + + +# Check if the new shape is a flatten / unflatten version of the current one +def _check_is_flatten(new_shape, old_shape, return_flatten_dim=False): + if not new_shape or not old_shape: + if return_flatten_dim: + return False, (-1, -1) + return False + if new_shape.numel() != old_shape.numel(): + if return_flatten_dim: + return False, (-1, -1) + return False + # a shape is a flatten version of another if all the first sizes and/or all the last sizes match + for i, (first_new, first_old) in enumerate(zip(new_shape, old_shape)): # noqa: B007 + if first_new != first_old: + break + # 'i' must be the result of the flatten op + for j, (last_new, last_old) in enumerate( # noqa: B007 + zip(reversed(new_shape), reversed(old_shape)) + ): + if last_new != last_old: + break + # j is also the result of the flatten, so if j and i match this is the result of a flatten + if i == len(new_shape) - j - 1: + if return_flatten_dim: + j = len(old_shape) - j - 1 + return True, (i, j) + return True + if return_flatten_dim: + return False, (-1, -1) + return False + + +def _check_is_unflatten(new_shape, old_shape, return_flatten_dim=False): + out = _check_is_flatten(old_shape, new_shape, return_flatten_dim=return_flatten_dim) + if return_flatten_dim: + out, (i, j) = out + # if out: + # j = len(new_shape) - j - 1 + return out, (i, j) + return out + + +def _create_segments_from_int(split_size, max_size): + if split_size <= 0: + raise RuntimeError( + f"split_size must be a positive integer, but got {split_size}." + ) + splits = [ + (start, min(start + split_size, max_size)) + for start in range(0, max_size, split_size) + ] + return splits + + +def _create_segments_from_list( + split_size: list[int] | tuple[int], + max_size: int, +): + splits = [ + (start, min(start + size, max_size)) + for start, size in zip( + [0] + list(itertools.accumulate(split_size[:-1])), + split_size, + ) + ] + total_split_size = sum(split_size) + if total_split_size != max_size: + raise RuntimeError( + f"Split method expects split_size to sum exactly to {max_size}, " + f"but got sum({split_size}) = {total_split_size}" + ) + + return splits + + +# Register JSON backends +register_backend(group="json", backends={"json": "json", "orjson": "orjson"}) + + +@implement_for("json") +def _json_dumps(data, **kwargs): + """JSON serialization using standard json module.""" + import json + + return json.dumps(data, **kwargs) + + +@implement_for("orjson") +def _json_dumps(data, **kwargs): # noqa: F811 + """JSON serialization using orjson module.""" + import orjson + + # orjson doesn't support separators parameter, so we need to handle it differently + if "separators" in kwargs: + # Remove separators for orjson and use default compact format + kwargs.pop("separators") + return orjson.dumps(data, **kwargs) + + +def json_dumps(data, **kwargs): + """Unified JSON serialization function that works with both json and orjson backends.""" + return _json_dumps(data, **kwargs) + + +def set_json_backend(backend): + """Set the JSON backend to use (either 'json' or 'orjson').""" + if backend not in ["json", "orjson"]: + raise ValueError("Backend must be either 'json' or 'orjson'") + set_backend("json", backend) + + +def get_json_backend(): + """Get the current JSON backend.""" + return get_backend("json") + + +if importlib.util.find_spec("orjson") is not None: + set_json_backend("orjson") +else: + set_json_backend("json") + + +class LinkedList(list): + """A thin wrapper around a list that automatically updates a tensordict when the list is modified.""" + + def __init__(self, *args, td: TensorDictBase, **kwargs): + super().__init__(*args, **kwargs) + self._td = weakref.ref(td) + + def __setitem__(self, key, value): + super().__setitem__(key, value) + td = self._td() + if td is not None: + td[key] = value + + def append(self, object: Any) -> None: + self._td = lambda: None + return super().append(object) + + def extend(self, object: Any) -> None: + self._td = lambda: None + return super().extend(object) + + def insert(self, index: int, object: Any) -> None: + self._td = lambda: None + return super().insert(index, object) + + def pop(self, index: int = -1) -> Any: + self._td = lambda: None + return super().pop(index) + + def remove(self, value: Any) -> None: + self._td = lambda: None + return super().remove(value) + + def clear(self) -> None: + self._td = lambda: None + return super().clear() + + def __delitem__(self, index: int) -> None: + self._td = lambda: None + return super().__delitem__(index) + + def __repr__(self) -> str: + return f"LinkedList({super().__repr__()})" + + def __str__(self) -> str: + return f"LinkedList({super().__str__()})" + + def __eq__(self, other: Any) -> bool: + return list(self) == other + + def __ne__(self, other: Any) -> bool: + return list(self) != other + + def __getstate__(self) -> object: + return list(self).__getstate__() + + def __iadd__(self, other: Any) -> Any: + self._td = lambda: None + return super().__iadd__(other) + + def __imul__(self, other: Any) -> Any: + self._td = lambda: None + return super().__imul__(other) + + +# register LinkedList in PyTree +@implement_for("torch", "2.3", None) +def _register_pytree_node(): + torch.utils._pytree.register_pytree_node( + LinkedList, + torch.utils._pytree._list_flatten, + torch.utils._pytree._list_unflatten, + serialized_type_name="builtins.list", + flatten_with_keys_fn=torch.utils._pytree._list_flatten_with_keys, + ) + + +@implement_for("torch", None, "2.3") +def _register_pytree_node(): # noqa: F811 # type: ignore + pass + + +_register_pytree_node() diff --git a/lib/python3.12/site-packages/torchvision.libs/libsharpyuv.5c41a003.so.0 b/lib/python3.12/site-packages/torchvision.libs/libsharpyuv.5c41a003.so.0 new file mode 100644 index 0000000000000000000000000000000000000000..f77d830efb43c49a8e614599b30b3a90cbd08bab Binary files /dev/null and b/lib/python3.12/site-packages/torchvision.libs/libsharpyuv.5c41a003.so.0 differ diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/INSTALLER b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/LICENSE b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..4555634d3e51beb788e94ac535d01b6482514a5b --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/LICENSE @@ -0,0 +1,35 @@ +From xFormers: + +Copyright (c) Facebook, Inc. and its affiliates + + +=== + +BSD 3-Clause License + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/METADATA b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..ad5c2b1e9cdcef262bc6bddff52f95deafc9242c --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/METADATA @@ -0,0 +1,30 @@ +Metadata-Version: 2.2 +Name: xformers +Version: 0.0.29.post2 +Summary: XFormers: A collection of composable Transformer building blocks. +Home-page: https://facebookresearch.github.io/xformers/ +Author: Facebook AI Research +Author-email: oncall+xformers@xmail.facebook.com +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: License :: OSI Approved :: BSD License +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Classifier: Operating System :: OS Independent +Requires-Python: >=3.9 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: numpy +Requires-Dist: torch==2.6.0 +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: description +Dynamic: description-content-type +Dynamic: home-page +Dynamic: requires-dist +Dynamic: requires-python +Dynamic: summary + +XFormers: A collection of composable Transformer building blocks.XFormers aims at being able to reproduce most architectures in the Transformer-family SOTA,defined as compatible and combined building blocks as opposed to monolithic models diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/RECORD b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..5dc66d7364687f4257f01ea000eb1bc015ef6919 --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/RECORD @@ -0,0 +1,318 @@ +xformers-0.0.29.post2.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +xformers-0.0.29.post2.dist-info/LICENSE,sha256=iwaVhtHd3wMc1GkkT2Y_GzmnThI1eC5iVcJJQ-ACt-o,1610 +xformers-0.0.29.post2.dist-info/METADATA,sha256=QlCMDkKA_bL69kO2kAKeYOIsNsK6RZOG1aTqyuAkLJA,1211 +xformers-0.0.29.post2.dist-info/RECORD,, +xformers-0.0.29.post2.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers-0.0.29.post2.dist-info/WHEEL,sha256=AiGbv82xhZWUVcPQkB8YCfdiItGrxCLvVhYDzK7pAs8,113 +xformers-0.0.29.post2.dist-info/top_level.txt,sha256=4Px1VcGhKk0j3XhKXjA8HtTm6EQOb0hazeJ5nQsNlKk,9 +xformers/_C.so,sha256=hKtA-JIOzkarlvS6YOks0WFP43v1JmVFOBRECt5ES4o,47840080 +xformers/_C_flashattention3.so,sha256=QEmviV1zmDnu8PWPJRYel1p0lkZ80R1YOIXm6hck_PY,120765064 +xformers/__init__.py,sha256=TKDA28gOAQdAzwVvU9IEpGWz677QRKgVFmnqvngAg2w,1710 +xformers/__pycache__/__init__.cpython-312.pyc,, +xformers/__pycache__/_cpp_lib.cpython-312.pyc,, +xformers/__pycache__/_deprecation_warning.cpython-312.pyc,, +xformers/__pycache__/attn_bias_utils.cpython-312.pyc,, +xformers/__pycache__/checkpoint.cpython-312.pyc,, +xformers/__pycache__/info.cpython-312.pyc,, +xformers/__pycache__/test.cpython-312.pyc,, +xformers/__pycache__/utils.cpython-312.pyc,, +xformers/__pycache__/version.cpython-312.pyc,, +xformers/_cpp_lib.py,sha256=diLPbor1Tp80qZKvJLeoiYKYRfBf8UnsO2UbaMRuQLI,4958 +xformers/_deprecation_warning.py,sha256=a2-JfNsT7EGLG5ZFarz7GTdnMe6EA3cli4BNfr3Wtog,456 +xformers/_flash_attn/__init__.py,sha256=33Vo6R_5k7y4iJR3Bk8Op3Uzxa0dZc4OdchOtdzzJSE,291 +xformers/_flash_attn/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/bert_padding.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/flash_attn_interface.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/flash_attn_triton.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/flash_attn_triton_og.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/flash_blocksparse_attention.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/flash_blocksparse_attn_interface.cpython-312.pyc,, +xformers/_flash_attn/__pycache__/fused_softmax.cpython-312.pyc,, +xformers/_flash_attn/bert_padding.py,sha256=gF1EmsdJ-HpQ86MRQ4VxDw-Sb_RVISdQALdNnoByHlw,9930 +xformers/_flash_attn/flash_attn_interface.py,sha256=dHIPTJx9uYVyWSkLdUcYu6KxgbnHsP6qkgwbXqJW9jo,59398 +xformers/_flash_attn/flash_attn_triton.py,sha256=Du81zbh8Ls70ExEsm00opziGvjGFfcZCoZDUO2zut9Q,41112 +xformers/_flash_attn/flash_attn_triton_amd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/bench.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/bwd_prefill.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/bwd_ref.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/fwd_decode.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/fwd_prefill.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/fwd_ref.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/interface_fa.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/interface_torch.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/test.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/__pycache__/utils.cpython-312.pyc,, +xformers/_flash_attn/flash_attn_triton_amd/bench.py,sha256=wIGZHcYI_Ria2BB6FBJ87GKRxzboRP1LOf2L__aPeA4,9837 +xformers/_flash_attn/flash_attn_triton_amd/bwd_prefill.py,sha256=iCe1jLP5_osxzEGyMBpLUW9koGezp91A5tsI0o6OQQM,20269 +xformers/_flash_attn/flash_attn_triton_amd/bwd_ref.py,sha256=BL2_4jYjRvUodX4GjgpvyQhVDgm2a6OcnVZJnozEz_A,9972 +xformers/_flash_attn/flash_attn_triton_amd/fwd_decode.py,sha256=vrk6GQqo9Tp8-SWSUzREWftqyokI5BQlTA_AAdmbEYA,23435 +xformers/_flash_attn/flash_attn_triton_amd/fwd_prefill.py,sha256=ynZUspV95iXQ0Ccld3jQ0ZxGuZHWrLWwysDL4FS71R0,32986 +xformers/_flash_attn/flash_attn_triton_amd/fwd_ref.py,sha256=ORf0qC92d3FFqPe0l6VA0gMjc2bVlYwUwF-g4Ilq2ko,11362 +xformers/_flash_attn/flash_attn_triton_amd/interface_fa.py,sha256=7NTNBEsr3RQfylD-0IrQXN_ECryEtR7BlvJFVPz-HDM,16292 +xformers/_flash_attn/flash_attn_triton_amd/interface_torch.py,sha256=DU_iepQ4h5FFvYxM4qDR5eh76l55iUiCLXdGBsc6KYo,3308 +xformers/_flash_attn/flash_attn_triton_amd/test.py,sha256=7jN9tQIRNJ_xpfPHrERU9aeG70OewtQKIYqyRFxed7o,30832 +xformers/_flash_attn/flash_attn_triton_amd/utils.py,sha256=XoKPb0Zzrjbbn7kX_j3YIqqhDW0xIhX4EjzqYvwPmng,12247 +xformers/_flash_attn/flash_attn_triton_og.py,sha256=LmvDju7LJG-wOYhoR6Zc2AmdPK2oWyB1VJpMjRhnWnE,11328 +xformers/_flash_attn/flash_blocksparse_attention.py,sha256=gsdH9VtYaVcTcP1rzZYPy1V_wUqgdvVcsB1h4Mk7RGs,7472 +xformers/_flash_attn/flash_blocksparse_attn_interface.py,sha256=2qK2KvVCt851_j8ZzHvjS-aMfdgVDu1yne67-iScWfo,7265 +xformers/_flash_attn/fused_softmax.py,sha256=0-XbXo7R1a5h4-EpUzPy--lwlGytfTDW34WGM5nmBAY,7793 +xformers/_flash_attn/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/layers/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/layers/__pycache__/patch_embed.cpython-312.pyc,, +xformers/_flash_attn/layers/__pycache__/rotary.cpython-312.pyc,, +xformers/_flash_attn/layers/patch_embed.py,sha256=H58CgME_qSOPTZLOG08wFgrQS1j34pvNwMPrkTj3Ek4,2136 +xformers/_flash_attn/layers/rotary.py,sha256=MqsUZ-Gxa0OcYLtL8OsjHIOkqyTacQHkMGpqADa2e6Q,21239 +xformers/_flash_attn/losses/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/losses/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/losses/__pycache__/cross_entropy.cpython-312.pyc,, +xformers/_flash_attn/losses/cross_entropy.py,sha256=tj5IoeUZuSzA1_82UFr7o-1WuoHyKAc1gVS6fWzAbDQ,3197 +xformers/_flash_attn/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/models/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/baichuan.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/bert.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/bigcode.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/btlm.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/falcon.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/gpt.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/gpt_neox.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/gptj.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/llama.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/opt.cpython-312.pyc,, +xformers/_flash_attn/models/__pycache__/vit.cpython-312.pyc,, +xformers/_flash_attn/models/baichuan.py,sha256=eFNWwoRQ02AIeQP0OoK8pNvYw0dqnHOshLigCQPkAEc,5730 +xformers/_flash_attn/models/bert.py,sha256=dMM6-Pj814pgQdsKkgkwg_grNZ7snM2juSgoUB14R7Q,33232 +xformers/_flash_attn/models/bigcode.py,sha256=mkYeItoJtmWVf2wKkUs5oXjwdbTdGSo5eHxi0-1maZ8,9383 +xformers/_flash_attn/models/btlm.py,sha256=d8YDjYTa2G1DutYu-YuVf15S_Dn6oKn8-HzERoersLA,4631 +xformers/_flash_attn/models/falcon.py,sha256=mA3wGv1a4zhbrUSlFNVVmTgVjiXc1sFTOi55eYpgSPo,6033 +xformers/_flash_attn/models/gpt.py,sha256=QGBMCw_osxD4VMMj1uC6TMlXlM5lIInxSUKmq5J5kSU,47669 +xformers/_flash_attn/models/gpt_neox.py,sha256=_704a9KQ2PcnID8uMV7yZ4ggjGlh1zZH5gszue6D1bI,5159 +xformers/_flash_attn/models/gptj.py,sha256=k2eqMNyMbU7CJVM_BHBjlKt0ByFz6ITSETqS1mJa89g,4436 +xformers/_flash_attn/models/llama.py,sha256=bDRI308iRpeJngZLrQlLTGYAmwYotqzUxnjBMirfn-k,16581 +xformers/_flash_attn/models/opt.py,sha256=L0ZIWKpSP44lcEbiVCzVT9un_5gFMAW6cvnS3KHcb-A,5164 +xformers/_flash_attn/models/vit.py,sha256=7i0WUI_jZvQ5TMoSKPPzf77ZcyMDfDJuQaINzXN_iQU,14074 +xformers/_flash_attn/modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/modules/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/modules/__pycache__/block.cpython-312.pyc,, +xformers/_flash_attn/modules/__pycache__/embedding.cpython-312.pyc,, +xformers/_flash_attn/modules/__pycache__/mha.cpython-312.pyc,, +xformers/_flash_attn/modules/__pycache__/mlp.cpython-312.pyc,, +xformers/_flash_attn/modules/block.py,sha256=WLi7JKj9_Zpk89ppzC7WTIoykJJ7TLOJbUSZePNnW1E,17349 +xformers/_flash_attn/modules/embedding.py,sha256=RCVeeiomlGNkLeQD8G6Udvex-NDI_xKD45hXjgZ2lbQ,8693 +xformers/_flash_attn/modules/mha.py,sha256=V6Ynog9pb_G9UVxetRjXlmWGExZlxmJkYVwAExXqUEk,43297 +xformers/_flash_attn/modules/mlp.py,sha256=G6KPQagfKq1DRn7hQRJ3OHznFJLZHj_PiidZE_zcLgg,6033 +xformers/_flash_attn/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/ops/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/ops/__pycache__/activations.cpython-312.pyc,, +xformers/_flash_attn/ops/__pycache__/fused_dense.cpython-312.pyc,, +xformers/_flash_attn/ops/__pycache__/layer_norm.cpython-312.pyc,, +xformers/_flash_attn/ops/__pycache__/rms_norm.cpython-312.pyc,, +xformers/_flash_attn/ops/activations.py,sha256=t5lzNg1In8LP6bKeTnyeMizwqjv27JGbJ6ylPdGvZYg,3939 +xformers/_flash_attn/ops/fused_dense.py,sha256=ACJKqkIfxZibxI3nb5ycb3pXBKaL_CM63rUUyQYNAUE,27907 +xformers/_flash_attn/ops/layer_norm.py,sha256=zr7NXIm-2mtEynTp1CS0fbFGI2Mqdp41dY4AfDWF6EQ,22443 +xformers/_flash_attn/ops/rms_norm.py,sha256=XEnihcj0a4aSz4LO55m5iKGVn4HKTeKN8TIyHjuDgxI,3988 +xformers/_flash_attn/ops/triton/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 +xformers/_flash_attn/ops/triton/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/cross_entropy.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/k_activations.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/layer_norm.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/linear.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/mlp.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/__pycache__/rotary.cpython-312.pyc,, +xformers/_flash_attn/ops/triton/cross_entropy.py,sha256=hjSfLhv4cKt-N8hTfpgkGMFdxhs8B4II6VIkEtck8EM,12845 +xformers/_flash_attn/ops/triton/k_activations.py,sha256=-Z3vIyO4JkqBMipKsPvhzmxljtBdIhJCsl_M-_ESqBo,4034 +xformers/_flash_attn/ops/triton/layer_norm.py,sha256=rNJwuijsZ6sKDtKHlkbT0qDzbi6vetVjjibpy9YRHFQ,35715 +xformers/_flash_attn/ops/triton/linear.py,sha256=OtRvKz8xdpl-7v3q_ZTaS9fdBt9XrzMyapgRr50uBbM,20841 +xformers/_flash_attn/ops/triton/mlp.py,sha256=_5lbZJFZg_pXeXYITGt4V_6LkB_yddClB_jt-diCOdw,6068 +xformers/_flash_attn/ops/triton/rotary.py,sha256=WH7tELBLZ23znuxnYUAzP7YWqwMXJmRgUQ8B64Vjdn4,8583 +xformers/_flash_attn/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xformers/_flash_attn/utils/__pycache__/__init__.cpython-312.pyc,, +xformers/_flash_attn/utils/__pycache__/benchmark.cpython-312.pyc,, +xformers/_flash_attn/utils/__pycache__/distributed.cpython-312.pyc,, +xformers/_flash_attn/utils/__pycache__/generation.cpython-312.pyc,, +xformers/_flash_attn/utils/__pycache__/pretrained.cpython-312.pyc,, +xformers/_flash_attn/utils/benchmark.py,sha256=JDtzdVhFyMIQqs3edbcXdXnmDf-O7RVpmZmn2ZFCvI0,7369 +xformers/_flash_attn/utils/distributed.py,sha256=qhcybRXtslssuV9LYaQy37haPaPtklM4YUMDx9UvnnQ,5825 +xformers/_flash_attn/utils/generation.py,sha256=9IVPvkf_hlbsxCWgECUA03293qHxXzWDFCTAOdqbAVo,30694 +xformers/_flash_attn/utils/pretrained.py,sha256=VZ6qk90sBJA7M86gRzPsNc_CkQXkj5HyrJvwl0I355k,3246 +xformers/attn_bias_utils.py,sha256=2MD_FqAf7Xci5W7jbeu6E3IYxqj0-4q6aWSXA_qzhko,19178 +xformers/benchmarks/LRA/__init__.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/benchmarks/LRA/__pycache__/__init__.cpython-312.pyc,, +xformers/benchmarks/LRA/__pycache__/batch_fetch_results.cpython-312.pyc,, +xformers/benchmarks/LRA/__pycache__/batch_submit.cpython-312.pyc,, +xformers/benchmarks/LRA/__pycache__/run_grid_search.cpython-312.pyc,, +xformers/benchmarks/LRA/__pycache__/run_tasks.cpython-312.pyc,, +xformers/benchmarks/LRA/__pycache__/run_with_submitit.cpython-312.pyc,, +xformers/benchmarks/LRA/batch_fetch_results.py,sha256=QsTgzFU45VP636TPiewnVCKD57bOR_YkjZq2KjMLzqc,3508 +xformers/benchmarks/LRA/batch_submit.py,sha256=7BgzEgSCJNN4vaQ8m-hEY2x5DxSkaPzp_Jorr2fS5p4,1685 +xformers/benchmarks/LRA/code/__init__.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/benchmarks/LRA/code/__pycache__/__init__.cpython-312.pyc,, +xformers/benchmarks/LRA/code/__pycache__/dataset.cpython-312.pyc,, +xformers/benchmarks/LRA/code/__pycache__/model_wrapper.cpython-312.pyc,, +xformers/benchmarks/LRA/code/dataset.py,sha256=6iMOhcM8rY75Ub_t3fopBxGjo1zCWwSIRAC3eHqAsIs,1403 +xformers/benchmarks/LRA/code/model_wrapper.py,sha256=18YKbVAASeoaGZqhM-Ps3scvblLbVoImHWKJqY7CQjw,9777 +xformers/benchmarks/LRA/run_grid_search.py,sha256=yLhNnQv4HDbEdQVApCdJiZzFes9VeGZjPOafHAaZ5Ew,5327 +xformers/benchmarks/LRA/run_tasks.py,sha256=GZP6cUgE1ftpQrp2IOf-MWbYJFcokt9bsLA8Q6PNyIU,9294 +xformers/benchmarks/LRA/run_with_submitit.py,sha256=WUtQJLYWnVVcp6mbmBQuqiA3pdvayEpZzqEpW5FZRz0,4612 +xformers/benchmarks/__init__.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/benchmarks/__pycache__/__init__.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_attn_decoding.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_core.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_indexing.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_mem_eff_attention.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_merge_attentions.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_nystrom_utils.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_revnet.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_sddmm.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_sequence_parallel_fused.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_sp24.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_swiglu.cpython-312.pyc,, +xformers/benchmarks/__pycache__/benchmark_tiled_matmul.cpython-312.pyc,, +xformers/benchmarks/__pycache__/utils.cpython-312.pyc,, +xformers/benchmarks/benchmark_attn_decoding.py,sha256=Jp45OCalf0Wfk8WZ1y8NBqyccup9lSnFs9R1Oq6Gjqg,12368 +xformers/benchmarks/benchmark_core.py,sha256=83x0wl7ouGqRgECQiRE3kzTXQwXHP3uWpQHvqxJWBlw,8626 +xformers/benchmarks/benchmark_indexing.py,sha256=W55qQHBuwlEp98X3XzoCrlHQg7NF00b9gc52_qDzhSE,5288 +xformers/benchmarks/benchmark_mem_eff_attention.py,sha256=HG8mfTCy1oFOsnOOQZ8zOGrOWotRTFU0fVJ4d4_xn0M,10101 +xformers/benchmarks/benchmark_merge_attentions.py,sha256=_6iSEhJGoeljSrVbOdi-tncYHPQNE7lj1wrTQlubSLo,3042 +xformers/benchmarks/benchmark_nystrom_utils.py,sha256=mUML4twJT1LCoCwtvS_3isfcHDS4XTH18leB_kix-OY,3230 +xformers/benchmarks/benchmark_revnet.py,sha256=qG7SxXpoJ-bi5cajWukyft7qVprmVfqqtx9PRQqIdI8,2671 +xformers/benchmarks/benchmark_sddmm.py,sha256=Tvq1rgXpeaj7kGjCKn_6pX2QA7TOt_aNG-nLCYa9_Eo,3765 +xformers/benchmarks/benchmark_sequence_parallel_fused.py,sha256=-Ob3G3c0xTJs--bw5yZQ203UqY1maOtKM3K5ZCDB6tw,14454 +xformers/benchmarks/benchmark_sp24.py,sha256=xcp-cR8t7U5WGfvbDoZeDUoTLXD-w7vVQerfd2kybVw,4861 +xformers/benchmarks/benchmark_swiglu.py,sha256=UD6tmg4JaPYct6KLwqnE_7UsFewKY_GpgV8XZMa1hFM,4295 +xformers/benchmarks/benchmark_tiled_matmul.py,sha256=oJW_-b1AkXEN1J6h35DVDzvZIjvDFMYld1KWq8h7Twk,3446 +xformers/benchmarks/utils.py,sha256=PSLBq4P26wNZqcIoo78afIDoqTDpJsJsS6-Fdus5zgE,24507 +xformers/checkpoint.py,sha256=PUYtKs1QE-3ZKoBfWhIT-aFHMD_uYJq4ZBNePE-Ug6o,20337 +xformers/components/__init__.py,sha256=YYnOMdsb2l8qdNR09HFFTu-wf34-y6g2G_nFb7oxaac,1004 +xformers/components/__pycache__/__init__.cpython-312.pyc,, +xformers/components/__pycache__/input_projection.cpython-312.pyc,, +xformers/components/__pycache__/residual.cpython-312.pyc,, +xformers/components/attention/__init__.py,sha256=ssO2FRcK9wALgkMc_H0QUujFfZjIZtYLImKPmmLpC4s,3311 +xformers/components/attention/__pycache__/__init__.cpython-312.pyc,, +xformers/components/attention/__pycache__/_sputnik_sparse.cpython-312.pyc,, +xformers/components/attention/__pycache__/attention_mask.cpython-312.pyc,, +xformers/components/attention/__pycache__/attention_patterns.cpython-312.pyc,, +xformers/components/attention/__pycache__/base.cpython-312.pyc,, +xformers/components/attention/__pycache__/core.cpython-312.pyc,, +xformers/components/attention/__pycache__/fourier_mix.cpython-312.pyc,, +xformers/components/attention/__pycache__/scaled_dot_product.cpython-312.pyc,, +xformers/components/attention/__pycache__/sparsity_config.cpython-312.pyc,, +xformers/components/attention/__pycache__/utils.cpython-312.pyc,, +xformers/components/attention/_sputnik_sparse.py,sha256=bibk2S-Coatt45gO9fzLTaAgUMKXbRi5U0eR_-QAbnw,3164 +xformers/components/attention/attention_mask.py,sha256=Y0zyM5JDnFzfE0TGPx4BxEeQlqIPX-XnZpO_lskQn4w,4585 +xformers/components/attention/attention_patterns.py,sha256=mZ-s8R-DS6XrOP_-ExZZku7XWVKUOEVfUdHpB9D5kek,9945 +xformers/components/attention/base.py,sha256=X8qmru_ln4eSac68rnYZzfAF3pbQDz5l6UxfRE-pR6Y,3305 +xformers/components/attention/core.py,sha256=KqkGcrwVsZmlovA0yQV3MKPAbGIoKx2AkyBlfxM3xHE,7641 +xformers/components/attention/fourier_mix.py,sha256=AKMjFgmSBomER9A8c9mbuERjnLR9KAcMafawRJ0eCko,1184 +xformers/components/attention/scaled_dot_product.py,sha256=vUcRVzQT3ol6hIqrTcCudlaCru1W1ZcZgjaTIkl7Lr4,4504 +xformers/components/attention/sparsity_config.py,sha256=lZLGeKVcqnKVGKqkjE5Av2yGIHB5zZvU9IwrR3UrvZo,41608 +xformers/components/attention/utils.py,sha256=04aEVvDG9LjzD9NfsQ3Kj-kZSnMFCppCFrJx6NPoV8g,3803 +xformers/components/input_projection.py,sha256=7oH2EfhMsujvjpXxIHzg5A3gT5xA0vJ2trb2-pZoh3g,3121 +xformers/components/residual.py,sha256=jTAnAytB9dZoXv-xEw9NBL2NnSqo3C68q9jlOpfS04I,5693 +xformers/cpp_lib.json,sha256=PwsWDW6OYaOEpQFkBwhXw1kThvv9KyKHREaABXt7rsU,395 +xformers/info.py,sha256=EzzBOTz7vp1OUPDlFT5yBSs3ItqAzMhfK975bv-A2FM,2672 +xformers/ops/__init__.py,sha256=3JYxy2MquvYUcdoZzSez-dTX257Gth7430QjW1xex7g,3502 +xformers/ops/__pycache__/__init__.cpython-312.pyc,, +xformers/ops/__pycache__/common.cpython-312.pyc,, +xformers/ops/__pycache__/differentiable_collectives.cpython-312.pyc,, +xformers/ops/__pycache__/indexing.cpython-312.pyc,, +xformers/ops/__pycache__/ipc.cpython-312.pyc,, +xformers/ops/__pycache__/modpar_layers.cpython-312.pyc,, +xformers/ops/__pycache__/rmsnorm.cpython-312.pyc,, +xformers/ops/__pycache__/rope_padded.cpython-312.pyc,, +xformers/ops/__pycache__/seqpar.cpython-312.pyc,, +xformers/ops/__pycache__/sequence_parallel_fused_ops.cpython-312.pyc,, +xformers/ops/__pycache__/sp24.cpython-312.pyc,, +xformers/ops/__pycache__/swiglu_op.cpython-312.pyc,, +xformers/ops/__pycache__/tiled_matmul.cpython-312.pyc,, +xformers/ops/__pycache__/unbind.cpython-312.pyc,, +xformers/ops/_triton/__init__.py,sha256=G7u-HQ1NEf-mlR3LN5qupDMVCHHA94H1S_aXoStfKOI,741 +xformers/ops/_triton/__pycache__/__init__.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/k_index_select_cat.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/k_scaled_index_add.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/matmul_perf_model.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/rmsnorm_kernels.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/rope_padded_kernels.cpython-312.pyc,, +xformers/ops/_triton/__pycache__/tiled_matmul_kernels.cpython-312.pyc,, +xformers/ops/_triton/k_index_select_cat.py,sha256=_gFAuuUbKd_CH0Awkc_sspxE0lmd4Y-gLndLu3G3brA,6194 +xformers/ops/_triton/k_scaled_index_add.py,sha256=KCEgK-9s_A1jdEVvBn84GlA4py7ZBrQSuxWyYaJG_jc,12802 +xformers/ops/_triton/matmul_perf_model.py,sha256=FpuoFYqUsm0zCn0G6htd1kAri-GhShCuj_MZPvdHXp4,8385 +xformers/ops/_triton/rmsnorm_kernels.py,sha256=2aLv_WHvVebn2HMvLLfIxLcHIScgWsMtz9PtdjUZLDU,5124 +xformers/ops/_triton/rope_padded_kernels.py,sha256=L9QqbmfGYUPMMmVkKvN08OG28BmIDL3aCXsNbhkl4bw,7221 +xformers/ops/_triton/tiled_matmul_kernels.py,sha256=p1G-VMcJnthKc71Bc4cE7Hul66U5Dy_-ur2YnlJ2BC4,13832 +xformers/ops/common.py,sha256=bgeQP7DxTkWhj_xiM-GUZ7QBoxRzg-VY3NfEOns5hhs,1911 +xformers/ops/differentiable_collectives.py,sha256=moSee9BpQDlfXgEq_P3IqxYcF9NjJGi8jbge_cJ9yGI,5356 +xformers/ops/fmha/__init__.py,sha256=EtZAYIbQBkiPWPCWuKiQpe2f3oSG7uGX8JyDoF9jH8I,30491 +xformers/ops/fmha/__pycache__/__init__.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/attn_bias.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/ck.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/ck_decoder.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/ck_splitk.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/common.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/cutlass.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/dispatch.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/flash.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/flash3.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/torch_attention_compat.cpython-312.pyc,, +xformers/ops/fmha/__pycache__/triton_splitk.cpython-312.pyc,, +xformers/ops/fmha/_triton/__init__.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/ops/fmha/_triton/__pycache__/__init__.cpython-312.pyc,, +xformers/ops/fmha/_triton/__pycache__/splitk_kernels.cpython-312.pyc,, +xformers/ops/fmha/_triton/splitk_kernels.py,sha256=2aGCZ6UftNykO-XQBSxUZYce1MzQA0xhSGAE7jqH2AM,41033 +xformers/ops/fmha/attn_bias.py,sha256=gtiyxPr2cVmpcUGRo-8u3KSY8jsVVuQL5ASsyeS_SrI,65449 +xformers/ops/fmha/ck.py,sha256=LrHNuPTsF9DJZqFIKYPMw-4El_ORH3hCmVfmpZMcjvs,17742 +xformers/ops/fmha/ck_decoder.py,sha256=3Lm0DA4TYgBu7KmlKEntnaDSL21iv__d3-z-BH5RAB4,5092 +xformers/ops/fmha/ck_splitk.py,sha256=4_kVxgiKM0YBUsQaWV071I6y4kDNmNore47rAZ_i-oc,6411 +xformers/ops/fmha/common.py,sha256=pw-pcpErmjfHJxL6LgxLm2hWTtn6UbVkvD4JrTia9MI,18579 +xformers/ops/fmha/cutlass.py,sha256=2QRon36JqrXth0pAixPsslYRIKTpVPnkWzXQinCOBbw,17148 +xformers/ops/fmha/dispatch.py,sha256=CLhYpRPS0vQOolffKlm3OlpLM1EJwwSgnaAxS3ZZK7A,6186 +xformers/ops/fmha/flash.py,sha256=mXTBqK7YXBqhoB8eqW_jARbRBgtW2MYqX9Ovas0lguY,28496 +xformers/ops/fmha/flash3.py,sha256=MaQ0S-ZxaIrAS0TATyrEwBOU-SFYC6kWTgU4sOvpZB0,17829 +xformers/ops/fmha/torch_attention_compat.py,sha256=n02uHdad4M7xwO2k3lDTs40vbQ0WfKskZ-DSdy7G914,5745 +xformers/ops/fmha/triton_splitk.py,sha256=8bqMisrKRfncHFr8WkNRigQuSu-oFtGup9ru2fG5Orw,36605 +xformers/ops/indexing.py,sha256=uQDU1rqgh-huDvOdtMekgSQSMXVndrC0n3vn8CFHNrI,6927 +xformers/ops/ipc.py,sha256=Q5YRa5G1LrAbm8syGZiSBzPvMqQEvR-gqML6t_18G-0,6649 +xformers/ops/modpar_layers.py,sha256=gOg2cAn3G-xhp1bQTM33dA1adQc5bmZp6v-jF4bh_sM,5703 +xformers/ops/rmsnorm.py,sha256=ofNCqraOJZBt_z1WoTMuW5HgfCr2wpJ_zkL_b0e-n2s,3358 +xformers/ops/rope_padded.py,sha256=bI6JCNHO0Rggp9oYW23KV4ONsoPqk4tGT3EGejqRZvY,12100 +xformers/ops/seqpar.py,sha256=dTNVh1DPAYJeOP-dmopGEuulQQ9WjLvo7BU0B7Y8zcU,12063 +xformers/ops/sequence_parallel_fused_ops.py,sha256=UVaJf0onPork4jw5LAVqa0-h-jryu0jpLdwevhQuSog,36973 +xformers/ops/sp24.py,sha256=MUPa8Klm1pqijH4Dm3CCRWJSDxOOEJkm2bjwHzG5UM4,27125 +xformers/ops/swiglu_op.py,sha256=mhb1wn7pDKlGs2BXzB7n4pr12djwik-rCjrxwRK8xEw,17327 +xformers/ops/tiled_matmul.py,sha256=mzGdPYDKIBwgwitOOo0R6ld_Q5wtm4mHawElx3bMPqY,12792 +xformers/ops/unbind.py,sha256=KWEif28duVWbzj14bL-O3jJTxMuTZPvuGqCkZQhsgYQ,3925 +xformers/profiler/__init__.py,sha256=4jTj_MGJ-HkCkHwkx6QneBm27dMMPxvNa6kLwro86bk,421 +xformers/profiler/__pycache__/__init__.cpython-312.pyc,, +xformers/profiler/__pycache__/api.cpython-312.pyc,, +xformers/profiler/__pycache__/device_limits.cpython-312.pyc,, +xformers/profiler/__pycache__/find_slowest.cpython-312.pyc,, +xformers/profiler/__pycache__/profile_analyzer.cpython-312.pyc,, +xformers/profiler/__pycache__/profiler.cpython-312.pyc,, +xformers/profiler/__pycache__/profiler_dcgm.cpython-312.pyc,, +xformers/profiler/__pycache__/profiler_dcgm_impl.cpython-312.pyc,, +xformers/profiler/api.py,sha256=t-Cp6VqSEBk7Nm0zChnpJSIAKCMuY7VoYn9mySna7xQ,2685 +xformers/profiler/device_limits.py,sha256=HzPTO6bHMk7HxHelSh3A0Bk8qyuxL3-SsFYGBXdjT6c,3588 +xformers/profiler/find_slowest.py,sha256=o5Kky0Lrdz5ea_5CMJE8NX-EFyaMylwBNVLREdU8GEc,5087 +xformers/profiler/profile_analyzer.py,sha256=v1TKzsE9lm5kkFoCLCXbmc3XMSG_7QD_J3QbeqmyYLc,8717 +xformers/profiler/profiler.py,sha256=sD_9eoeKttMBVfqhODw5LIzp4MYFNx8t7xC4qp5hAYw,12756 +xformers/profiler/profiler_dcgm.py,sha256=1NWMET7oVtUoBYhEFoQ-F17s1B8xmqcrL3Hjddry-H0,1153 +xformers/profiler/profiler_dcgm_impl.py,sha256=umLRJQCYOIyYHwFsyXDk9eNmbs4DUmuvAwe6tUahN6k,8472 +xformers/sparse/__init__.py,sha256=NpDMLawDg0ra3mR7LpNnkua8KOpeY6h2AL4or88rGws,317 +xformers/sparse/__pycache__/__init__.cpython-312.pyc,, +xformers/sparse/__pycache__/_csr_ops.cpython-312.pyc,, +xformers/sparse/__pycache__/blocksparse_tensor.cpython-312.pyc,, +xformers/sparse/__pycache__/csr_tensor.cpython-312.pyc,, +xformers/sparse/__pycache__/utils.cpython-312.pyc,, +xformers/sparse/_csr_ops.py,sha256=p8bGD91Y52LuYieY97ZCtlOpYFNnnPRjXDuBoK5_Bp0,4873 +xformers/sparse/blocksparse_tensor.py,sha256=5f1iIqrK_R0cza1anOIJo4CkaMIUo1WbMlJ0LjPbd8s,8894 +xformers/sparse/csr_tensor.py,sha256=TZbVB4aP0WVG2Gr8gB3orkPbNWdjeGxktxYzHAvkYfc,14191 +xformers/sparse/utils.py,sha256=yxugnW0FLK7J4nc3xTPlRMcyVVQE7xwUcgt47kmjP-o,4274 +xformers/test.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/triton/__init__.py,sha256=BBoPvMfJjZ0Fi25JCRS8sPYUfLbHneLncV3JeGDIGHg,198 +xformers/triton/__pycache__/__init__.cpython-312.pyc,, +xformers/triton/__pycache__/importing.cpython-312.pyc,, +xformers/triton/__pycache__/vararg_kernel.cpython-312.pyc,, +xformers/triton/importing.py,sha256=-gHQG6er7JP4Z88mA5DJL5j6CT2mtdSFd989ZOWiWBk,908 +xformers/triton/vararg_kernel.py,sha256=O7_HmdejlNKIdG_bO2i8aHToqLNwrJuqNHOGqarwzu4,8956 +xformers/utils.py,sha256=eCX-EDLXtb0pOvQAeQVwC8LiMqeBiX2NEyb89jqEVas,5064 +xformers/version.py,sha256=n75bWVQoJa0gsi-54tbbkZFmIIXCoXDrss-RtaHUNos,42 diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/REQUESTED b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/WHEEL b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..9b2f5ecffb31862237ebd35519cb6ee109572d43 --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (75.8.0) +Root-Is-Purelib: false +Tag: cp312-cp312-manylinux_2_28_x86_64 + diff --git a/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/top_level.txt b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..56e3c822be24e7529619f4406fbfb2ad6dfae3fd --- /dev/null +++ b/lib/python3.12/site-packages/xformers-0.0.29.post2.dist-info/top_level.txt @@ -0,0 +1 @@ +xformers